Skip to main content

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