foundry_wallets/wallet_browser/
server.rs1use 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#[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 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 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 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 pub fn port(&self) -> u16 {
90 self.port
91 }
92
93 pub fn open_browser(&self) -> bool {
95 self.open_browser
96 }
97
98 pub fn timeout(&self) -> Duration {
100 self.timeout
101 }
102
103 pub fn session_token(&self) -> &str {
105 self.state.session_token()
106 }
107
108 pub async fn is_connected(&self) -> bool {
110 self.state.is_connected().await
111 }
112
113 pub async fn get_connection(&self) -> Option<Connection> {
115 self.state.get_connection().await
116 }
117
118 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 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 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}