v2
This commit is contained in:
		| @@ -5,6 +5,7 @@ use wasm_bindgen::prelude::*; | ||||
| use serde::{Deserialize, Serialize}; | ||||
| use sigsocket_client::{SigSocketClient, SignRequest, SignRequestHandler, Result as SigSocketResult, SigSocketError}; | ||||
| use web_sys::console; | ||||
| use base64::prelude::*; | ||||
|  | ||||
| use crate::vault_bindings::{get_workspace_default_public_key, get_current_keyspace_name, is_unlocked, sign_with_default_keypair}; | ||||
|  | ||||
| @@ -31,7 +32,32 @@ impl ExtensionNotificationHandler { | ||||
|  | ||||
| impl SignRequestHandler for ExtensionNotificationHandler { | ||||
|     fn handle_sign_request(&self, request: &SignRequest) -> SigSocketResult<Vec<u8>> { | ||||
|         // Create event object for JavaScript | ||||
|         console_log!("📨 WASM: Handling sign request: {}", request.id); | ||||
|  | ||||
|         // First, store the request in the WASM client | ||||
|         let store_result = SIGSOCKET_CLIENT.with(|c| { | ||||
|             let mut client_opt = c.borrow_mut(); | ||||
|             if let Some(client) = client_opt.as_mut() { | ||||
|                 // Get the connected public key as the target | ||||
|                 if let Some(target_public_key) = client.connected_public_key() { | ||||
|                     client.add_pending_request(request.clone(), target_public_key.to_string()); | ||||
|                     console_log!("✅ WASM: Stored sign request: {}", request.id); | ||||
|                     Ok(()) | ||||
|                 } else { | ||||
|                     Err(SigSocketError::Other("No connected public key".to_string())) | ||||
|                 } | ||||
|             } else { | ||||
|                 Err(SigSocketError::Other("No SigSocket client available".to_string())) | ||||
|             } | ||||
|         }); | ||||
|  | ||||
|         // If storage failed, return error | ||||
|         if let Err(e) = store_result { | ||||
|             console_log!("❌ WASM: Failed to store request: {:?}", e); | ||||
|             return Err(e); | ||||
|         } | ||||
|  | ||||
|         // Create event object for JavaScript notification | ||||
|         let event = js_sys::Object::new(); | ||||
|         js_sys::Reflect::set(&event, &"type".into(), &"sign_request".into()) | ||||
|             .map_err(|_| SigSocketError::Other("Failed to set event type".to_string()))?; | ||||
| @@ -40,17 +66,16 @@ impl SignRequestHandler for ExtensionNotificationHandler { | ||||
|         js_sys::Reflect::set(&event, &"message".into(), &request.message.clone().into()) | ||||
|             .map_err(|_| SigSocketError::Other("Failed to set message".to_string()))?; | ||||
|  | ||||
|         // Store the request in our pending requests (this will be done by the client) | ||||
|         // and notify the extension | ||||
|         // Notify the extension | ||||
|         match self.callback.call1(&wasm_bindgen::JsValue::NULL, &event) { | ||||
|             Ok(_) => { | ||||
|                 console_log!("Notified extension about sign request: {}", request.id); | ||||
|                 console_log!("✅ WASM: Notified extension about sign request: {}", request.id); | ||||
|                 // Return an error to indicate this request should not be auto-signed | ||||
|                 // The extension will handle the approval flow | ||||
|                 Err(SigSocketError::Other("Request forwarded to extension for approval".to_string())) | ||||
|             } | ||||
|             Err(e) => { | ||||
|                 console_log!("Failed to notify extension: {:?}", e); | ||||
|                 console_log!("❌ WASM: Failed to notify extension: {:?}", e); | ||||
|                 Err(SigSocketError::Other("Extension notification failed".to_string())) | ||||
|             } | ||||
|         } | ||||
| @@ -72,10 +97,12 @@ pub struct SigSocketManager; | ||||
|  | ||||
| #[wasm_bindgen] | ||||
| impl SigSocketManager { | ||||
|     /// Connect to SigSocket server with a specific workspace and event callback | ||||
|     /// Connect to SigSocket server with smart connection management | ||||
|     /// | ||||
|     /// This establishes a real WebSocket connection using the workspace's default public key | ||||
|     /// and integrates with the vault system for security validation. | ||||
|     /// This handles all connection logic: | ||||
|     /// - Reuses existing connection if same workspace | ||||
|     /// - Switches connection if different workspace | ||||
|     /// - Creates new connection if none exists | ||||
|     /// | ||||
|     /// # Arguments | ||||
|     /// * `workspace` - The workspace name to connect with | ||||
| @@ -98,25 +125,56 @@ impl SigSocketManager { | ||||
|         let public_key_bytes = hex::decode(&public_key_hex) | ||||
|             .map_err(|e| JsValue::from_str(&format!("Invalid public key format: {}", e)))?; | ||||
|          | ||||
|         // 3. Create SigSocket client with extension notification handler | ||||
|         let mut client = SigSocketClient::new(server_url, public_key_bytes) | ||||
|             .map_err(|e| JsValue::from_str(&format!("Failed to create client: {:?}", e)))?; | ||||
|         // 3. Check if already connected to same workspace and handle disconnection | ||||
|         let should_connect = SIGSOCKET_CLIENT.with(|c| { | ||||
|             let mut client_opt = c.borrow_mut(); | ||||
|  | ||||
|         // Set up extension notification handler using existing API | ||||
|         let handler = ExtensionNotificationHandler::new(event_callback.clone()); | ||||
|         client.set_sign_handler(handler); | ||||
|             // Check if we already have a client for this workspace | ||||
|             if let Some(existing_client) = client_opt.as_ref() { | ||||
|                 if let Some(existing_key) = existing_client.connected_public_key() { | ||||
|                     if existing_key == hex::encode(&public_key_bytes) && existing_client.is_connected() { | ||||
|                         console_log!("🔄 WASM: Already connected to workspace: {}", workspace); | ||||
|                         return false; // Reuse existing connection | ||||
|                     } else { | ||||
|                         console_log!("🔄 WASM: Switching workspace from {} to {}", | ||||
|                                    existing_key, hex::encode(&public_key_bytes)); | ||||
|  | ||||
|         // 4. Connect to the WebSocket server | ||||
|         client.connect().await | ||||
|             .map_err(|e| JsValue::from_str(&format!("Connection failed: {:?}", e)))?; | ||||
|                         // Disconnect the old client | ||||
|                         *client_opt = None; // This will drop the old client and close WebSocket | ||||
|                         console_log!("🔌 WASM: Disconnected from old workspace"); | ||||
|  | ||||
|         console_log!("SigSocket connected successfully to {}", server_url); | ||||
|                         return true; // Need new connection | ||||
|                     } | ||||
|                 } | ||||
|             } | ||||
|  | ||||
|         // 5. Store the connected client | ||||
|         SIGSOCKET_CLIENT.with(|c| { | ||||
|             *c.borrow_mut() = Some(client); | ||||
|             true // Need new connection, no old one to disconnect | ||||
|         }); | ||||
|  | ||||
|         // 4. Create and connect if needed | ||||
|         if should_connect { | ||||
|             console_log!("🔗 WASM: Creating new connection for workspace: {}", workspace); | ||||
|  | ||||
|             // Create new client | ||||
|             let mut client = SigSocketClient::new(server_url, public_key_bytes.clone()) | ||||
|                 .map_err(|e| JsValue::from_str(&format!("Failed to create client: {:?}", e)))?; | ||||
|  | ||||
|             // Set up extension notification handler | ||||
|             let handler = ExtensionNotificationHandler::new(event_callback.clone()); | ||||
|             client.set_sign_handler(handler); | ||||
|  | ||||
|             // Connect to the WebSocket server | ||||
|             client.connect().await | ||||
|                 .map_err(|e| JsValue::from_str(&format!("Connection failed: {:?}", e)))?; | ||||
|  | ||||
|             console_log!("✅ WASM: Connected to SigSocket server for workspace: {}", workspace); | ||||
|  | ||||
|             // Store the connected client | ||||
|             SIGSOCKET_CLIENT.with(|c| { | ||||
|                 *c.borrow_mut() = Some(client); | ||||
|             }); | ||||
|         } | ||||
|  | ||||
|         // 6. Return connection info | ||||
|         let connection_info = SigSocketConnectionInfo { | ||||
|             workspace: workspace.to_string(), | ||||
| @@ -150,7 +208,7 @@ impl SigSocketManager { | ||||
|     } | ||||
|      | ||||
|     /// Disconnect from SigSocket server | ||||
|     ///  | ||||
|     /// | ||||
|     /// # Returns | ||||
|     /// * `Ok(())` - Successfully disconnected | ||||
|     /// * `Err(error)` - If disconnect failed | ||||
| @@ -158,9 +216,16 @@ impl SigSocketManager { | ||||
|     pub async fn disconnect() -> Result<(), JsValue> { | ||||
|         SIGSOCKET_CLIENT.with(|c| { | ||||
|             let mut client_opt = c.borrow_mut(); | ||||
|             if let Some(_client) = client_opt.take() { | ||||
|                 // client.disconnect().await?; // Will be async in real implementation | ||||
|                 console_log!("SigSocket client disconnected"); | ||||
|             if let Some(client) = client_opt.take() { | ||||
|                 let workspace_info = client.connected_public_key() | ||||
|                     .map(|key| key[..16].to_string()) | ||||
|                     .unwrap_or_else(|| "unknown".to_string()); | ||||
|  | ||||
|                 // Dropping the client will close the WebSocket connection | ||||
|                 drop(client); | ||||
|                 console_log!("🔌 WASM: Disconnected SigSocket client (was: {}...)", workspace_info); | ||||
|             } else { | ||||
|                 console_log!("🔌 WASM: No SigSocket client to disconnect"); | ||||
|             } | ||||
|             Ok(()) | ||||
|         }) | ||||
| @@ -254,27 +319,51 @@ impl SigSocketManager { | ||||
|          | ||||
|         // 3. Sign with vault | ||||
|         let signature_result = sign_with_default_keypair(&message_bytes).await?; | ||||
|         let signature_obj: serde_json::Value = serde_json::from_str(&signature_result.as_string().unwrap()) | ||||
|             .map_err(|e| JsValue::from_str(&format!("Failed to parse signature: {}", e)))?; | ||||
|         let signature_hex = signature_result.as_string() | ||||
|             .ok_or_else(|| JsValue::from_str("Signature result is not a string"))?; | ||||
|  | ||||
|         // Convert hex signature to base64 for SigSocket protocol | ||||
|         let signature_bytes = hex::decode(&signature_hex) | ||||
|             .map_err(|e| JsValue::from_str(&format!("Invalid hex signature: {}", e)))?; | ||||
|         let signature_base64 = base64::prelude::BASE64_STANDARD.encode(&signature_bytes); | ||||
|          | ||||
|         let signature_base64 = signature_obj["signature"].as_str() | ||||
|             .ok_or_else(|| JsValue::from_str("Invalid signature format"))?; | ||||
|          | ||||
|         // 4. Send response to server and remove request | ||||
|         // 4. Get original message for response | ||||
|         let original_message = SIGSOCKET_CLIENT.with(|c| { | ||||
|             let client = c.borrow(); | ||||
|             let client = client.as_ref().ok_or_else(|| JsValue::from_str("Not connected"))?; | ||||
|  | ||||
|             let request = client.get_pending_request(request_id) | ||||
|                 .ok_or_else(|| JsValue::from_str("Request not found"))?; | ||||
|  | ||||
|             Ok::<String, JsValue>(request.request.message.clone()) | ||||
|         })?; | ||||
|  | ||||
|         // 5. Send response to server (create a new scope to avoid borrowing issues) | ||||
|         { | ||||
|             let client_ref = SIGSOCKET_CLIENT.with(|c| { | ||||
|                 c.borrow().as_ref().map(|client| client as *const SigSocketClient) | ||||
|             }).ok_or_else(|| JsValue::from_str("Not connected"))?; | ||||
|  | ||||
|             // SAFETY: We know the client exists and we're using it synchronously | ||||
|             let client = unsafe { &*client_ref }; | ||||
|  | ||||
|             client.send_response(request_id, &original_message, &signature_base64).await | ||||
|                 .map_err(|e| JsValue::from_str(&format!("Failed to send response: {:?}", e)))?; | ||||
|  | ||||
|             console_log!("✅ WASM: Sent signature response to server for request: {}", request_id); | ||||
|         } | ||||
|  | ||||
|         // 6. Remove the request after successful send | ||||
|         SIGSOCKET_CLIENT.with(|c| { | ||||
|             let mut client = c.borrow_mut(); | ||||
|             let client = client.as_mut().ok_or_else(|| JsValue::from_str("Not connected"))?; | ||||
|              | ||||
|             // Send response (will be async in real implementation) | ||||
|             // client.send_response(request_id, &original_request.message, signature_base64).await?; | ||||
|              | ||||
|             // Remove the request | ||||
|             client.remove_pending_request(request_id); | ||||
|              | ||||
|             console_log!("Approved and sent signature for request: {}", request_id); | ||||
|              | ||||
|             Ok(signature_base64.to_string()) | ||||
|         }) | ||||
|             if let Some(client) = client.as_mut() { | ||||
|                 client.remove_pending_request(request_id); | ||||
|                 console_log!("✅ WASM: Removed request from pending list: {}", request_id); | ||||
|             } | ||||
|         }); | ||||
|  | ||||
|         console_log!("🎉 WASM: Successfully approved and sent signature for request: {}", request_id); | ||||
|         Ok(signature_base64) | ||||
|     } | ||||
|      | ||||
|     /// Reject a sign request | ||||
| @@ -288,20 +377,32 @@ impl SigSocketManager { | ||||
|     /// * `Err(error)` - If rejection failed | ||||
|     #[wasm_bindgen] | ||||
|     pub async fn reject_request(request_id: &str, reason: &str) -> Result<(), JsValue> { | ||||
|         // Send rejection to server first | ||||
|         { | ||||
|             let client_ref = SIGSOCKET_CLIENT.with(|c| { | ||||
|                 c.borrow().as_ref().map(|client| client as *const SigSocketClient) | ||||
|             }).ok_or_else(|| JsValue::from_str("Not connected"))?; | ||||
|  | ||||
|             // SAFETY: We know the client exists and we're using it synchronously | ||||
|             let client = unsafe { &*client_ref }; | ||||
|  | ||||
|             client.send_rejection(request_id, reason).await | ||||
|                 .map_err(|e| JsValue::from_str(&format!("Failed to send rejection: {:?}", e)))?; | ||||
|  | ||||
|             console_log!("✅ WASM: Sent rejection to server for request: {}", request_id); | ||||
|         } | ||||
|  | ||||
|         // Remove the request after successful send | ||||
|         SIGSOCKET_CLIENT.with(|c| { | ||||
|             let mut client = c.borrow_mut(); | ||||
|             let client = client.as_mut().ok_or_else(|| JsValue::from_str("Not connected"))?; | ||||
|             if let Some(client) = client.as_mut() { | ||||
|                 client.remove_pending_request(request_id); | ||||
|                 console_log!("✅ WASM: Removed rejected request from pending list: {}", request_id); | ||||
|             } | ||||
|         }); | ||||
|  | ||||
|             // Send rejection (will be async in real implementation) | ||||
|             // client.send_rejection(request_id, reason).await?; | ||||
|  | ||||
|             // Remove the request | ||||
|             client.remove_pending_request(request_id); | ||||
|  | ||||
|             console_log!("Rejected request {}: {}", request_id, reason); | ||||
|  | ||||
|             Ok(()) | ||||
|         }) | ||||
|         console_log!("🚫 WASM: Successfully rejected request: {} (reason: {})", request_id, reason); | ||||
|         Ok(()) | ||||
|     } | ||||
|  | ||||
|     /// Get pending requests filtered by current workspace | ||||
|   | ||||
		Reference in New Issue
	
	Block a user