foundry_wallets/wallet_browser/
server.rs

1use std::{
2    net::SocketAddr,
3    sync::Arc,
4    time::{Duration, Instant},
5};
6
7use alloy_dyn_abi::TypedData;
8use alloy_primitives::{Address, Bytes, TxHash};
9use tokio::{
10    net::TcpListener,
11    sync::{Mutex, oneshot},
12};
13use uuid::Uuid;
14
15use crate::wallet_browser::{
16    error::BrowserWalletError,
17    router::build_router,
18    state::BrowserWalletState,
19    types::{
20        BrowserSignRequest, BrowserSignTypedDataRequest, BrowserTransactionRequest, Connection,
21        SignRequest, SignType,
22    },
23};
24
25/// Browser wallet server.
26#[derive(Debug, Clone)]
27pub struct BrowserWalletServer {
28    port: u16,
29    state: Arc<BrowserWalletState>,
30    shutdown_tx: Option<Arc<Mutex<Option<oneshot::Sender<()>>>>>,
31    open_browser: bool,
32    timeout: Duration,
33}
34
35impl BrowserWalletServer {
36    /// Create a new browser wallet server.
37    pub fn new(port: u16, open_browser: bool, timeout: Duration, development: bool) -> Self {
38        Self {
39            port,
40            state: Arc::new(BrowserWalletState::new(Uuid::new_v4().to_string(), development)),
41            shutdown_tx: None,
42            open_browser,
43            timeout,
44        }
45    }
46
47    /// Start the server and open browser.
48    pub async fn start(&mut self) -> Result<(), BrowserWalletError> {
49        let router = build_router(self.state.clone(), self.port).await;
50
51        let addr = SocketAddr::from(([127, 0, 0, 1], self.port));
52        let listener = TcpListener::bind(addr)
53            .await
54            .map_err(|e| BrowserWalletError::ServerError(e.to_string()))?;
55        self.port = listener.local_addr().unwrap().port();
56
57        let (shutdown_tx, shutdown_rx) = oneshot::channel();
58        self.shutdown_tx = Some(Arc::new(Mutex::new(Some(shutdown_tx))));
59
60        tokio::spawn(async move {
61            let server = axum::serve(listener, router);
62            let _ = server
63                .with_graceful_shutdown(async {
64                    let _ = shutdown_rx.await;
65                })
66                .await;
67        });
68
69        if self.open_browser {
70            webbrowser::open(&format!("http://127.0.0.1:{}", self.port)).map_err(|e| {
71                BrowserWalletError::ServerError(format!("Failed to open browser: {e}"))
72            })?;
73        }
74
75        Ok(())
76    }
77
78    /// Stop the server.
79    pub async fn stop(&mut self) -> Result<(), BrowserWalletError> {
80        if let Some(shutdown_arc) = self.shutdown_tx.take()
81            && let Some(tx) = shutdown_arc.lock().await.take()
82        {
83            let _ = tx.send(());
84        }
85        Ok(())
86    }
87
88    /// Get the server port.
89    pub fn port(&self) -> u16 {
90        self.port
91    }
92
93    /// Check if the browser should be opened.
94    pub fn open_browser(&self) -> bool {
95        self.open_browser
96    }
97
98    /// Get the timeout duration.
99    pub fn timeout(&self) -> Duration {
100        self.timeout
101    }
102
103    /// Get the session token.
104    pub fn session_token(&self) -> &str {
105        self.state.session_token()
106    }
107
108    /// Check if a wallet is connected.
109    pub async fn is_connected(&self) -> bool {
110        self.state.is_connected().await
111    }
112
113    /// Get current wallet connection.
114    pub async fn get_connection(&self) -> Option<Connection> {
115        self.state.get_connection().await
116    }
117
118    /// Request a transaction to be signed and sent via the browser wallet.
119    pub async fn request_transaction(
120        &self,
121        request: BrowserTransactionRequest,
122    ) -> Result<TxHash, BrowserWalletError> {
123        if !self.is_connected().await {
124            return Err(BrowserWalletError::NotConnected);
125        }
126
127        let tx_id = request.id;
128
129        self.state.add_transaction_request(request).await;
130
131        let start = Instant::now();
132
133        loop {
134            if let Some(response) = self.state.get_transaction_response(&tx_id).await {
135                if let Some(hash) = response.hash {
136                    return Ok(hash);
137                } else if let Some(error) = response.error {
138                    return Err(BrowserWalletError::Rejected {
139                        operation: "Transaction",
140                        reason: error,
141                    });
142                } else {
143                    return Err(BrowserWalletError::ServerError(
144                        "Transaction response missing both hash and error".to_string(),
145                    ));
146                }
147            }
148
149            if start.elapsed() > self.timeout {
150                self.state.remove_transaction_request(&tx_id).await;
151                return Err(BrowserWalletError::Timeout { operation: "Transaction" });
152            }
153
154            tokio::time::sleep(Duration::from_millis(100)).await;
155        }
156    }
157
158    /// Request a message to be signed via the browser wallet.
159    pub async fn request_signing(
160        &self,
161        request: BrowserSignRequest,
162    ) -> Result<Bytes, BrowserWalletError> {
163        if !self.is_connected().await {
164            return Err(BrowserWalletError::NotConnected);
165        }
166
167        let tx_id = request.id;
168
169        self.state.add_signing_request(request).await;
170
171        let start = Instant::now();
172
173        loop {
174            if let Some(response) = self.state.get_signing_response(&tx_id).await {
175                if let Some(signature) = response.signature {
176                    return Ok(signature);
177                } else if let Some(error) = response.error {
178                    return Err(BrowserWalletError::Rejected {
179                        operation: "Signing",
180                        reason: error,
181                    });
182                } else {
183                    return Err(BrowserWalletError::ServerError(
184                        "Signing response missing both signature and error".to_string(),
185                    ));
186                }
187            }
188
189            if start.elapsed() > self.timeout {
190                self.state.remove_signing_request(&tx_id).await;
191                return Err(BrowserWalletError::Timeout { operation: "Signing" });
192            }
193
194            tokio::time::sleep(Duration::from_millis(100)).await;
195        }
196    }
197
198    /// Request EIP-712 typed data signing via the browser wallet.
199    pub async fn request_typed_data_signing(
200        &self,
201        address: Address,
202        typed_data: TypedData,
203    ) -> Result<Bytes, BrowserWalletError> {
204        let request = BrowserSignTypedDataRequest { id: Uuid::new_v4(), address, typed_data };
205
206        let sign_request = BrowserSignRequest {
207            id: request.id,
208            sign_type: SignType::SignTypedDataV4,
209            request: SignRequest {
210                message: serde_json::to_string(&request.typed_data).map_err(|e| {
211                    BrowserWalletError::ServerError(format!("Failed to serialize typed data: {e}"))
212                })?,
213                address: request.address,
214            },
215        };
216
217        self.request_signing(sign_request).await
218    }
219}