Optimize handshake

src/commit.rs

@@ -9,7 +9,7 @@ use anyhow::{Error, Result};
 use bitcoincore_rpc::bitcoin::hex::DisplayHex;
 use hex::FromHex;
 use sdk_common::pcd::{Member, Pcd, RoleDefinition};
-use sdk_common::serialization::{MemberOutPointMap, OutPointProcessMap};
+use sdk_common::serialization::{OutPointMemberMap, OutPointProcessMap};
 use sdk_common::silentpayments::create_transaction;
 use sdk_common::sp_client::spclient::Recipient;
 use sdk_common::network::CommitMessage;
@@ -110,9 +110,9 @@ fn handle_initial_transaction(tx: Transaction, commit_msg: &CommitMessage) -> Re
     Ok(root_commitment)
 }
 
-pub static MEMBERLIST: OnceLock<Mutex<HashMap<Member, OutPoint>>> = OnceLock::new();
+pub static MEMBERLIST: OnceLock<Mutex<HashMap<OutPoint, Member>>> = OnceLock::new();
 
-pub fn lock_members() -> Result<MutexGuard<'static, HashMap<Member, OutPoint>>, anyhow::Error> {
+pub fn lock_members() -> Result<MutexGuard<'static, HashMap<OutPoint, Member>>, anyhow::Error> {
     MEMBERLIST
         .get_or_init(|| Mutex::new(HashMap::new()))
         .lock_anyhow()
@@ -141,8 +141,8 @@ fn handle_member_list(pcd_commitment: &Value, roles: &HashMap<String, RoleDefini
             let member = owner_role.members.get(0).unwrap();
             let mut memberlist = lock_members()?;
             memberlist.insert(
-                member.clone(),
                 root_commitment,
+                member.to_owned(),
             );
 
             return Ok(());
@@ -174,7 +174,7 @@ pub fn dump_cached_members() -> Result<(), anyhow::Error> {
 
     let members_file = &storage.members_file;
 
-    let members_map = MemberOutPointMap(members);
+    let members_map = OutPointMemberMap(members);
     let json = serde_json::to_value(&members_map)?;
     members_file.save(&json)?;
 
@@ -199,6 +199,21 @@ pub fn dump_cached_processes() -> Result<(), anyhow::Error> {
     Ok(())
 }
 
+// pub fn load_cached_processes() -> Result<(), anyhow::Error> {
+//     let storage = STORAGE.get().ok_or(Error::msg("STORAGE is not initialized"))?.lock_anyhow()?;
+
+//     let processes_file = &storage.processes_file;
+
+//     let raw_data = processes_file.load()?;
+//     let processes: HashMap<OutPoint, Process> = serde_json::from_slice(&raw_data)?;
+
+//     let mut cached_processes = lock_processes()?;
+
+//     *cached_processes = processes;
+    
+//     Ok(())
+// }
+
 // Register a new state when validation tokens are empty
 fn register_new_state(commitment: &mut Process, commit_msg: CommitMessage) -> Result<OutPoint> {
     let concurrent_states = commitment.get_latest_concurrent_states()?;

src/main.rs

@@ -16,7 +16,7 @@ use futures_util::{future, pin_mut, stream::TryStreamExt, FutureExt, StreamExt};
 use log::{debug, error, warn};
 use message::{broadcast_message, process_message, BroadcastType, MessageCache, MESSAGECACHE};
 use scan::{check_transaction_alone, compute_partial_tweak_to_transaction};
-use sdk_common::{network::HandshakeMessage, pcd::Member, process::{lock_processes, Process, CACHEDPROCESSES}, serialization::{MemberOutPointMap, OutPointProcessMap}, sp_client::{bitcoin::{
+use sdk_common::{network::HandshakeMessage, pcd::Member, process::{lock_processes, Process, CACHEDPROCESSES}, serialization::{OutPointMemberMap, OutPointProcessMap}, sp_client::{bitcoin::{
     consensus::deserialize,
     hex::{DisplayHex, FromHex},
     Amount, Network, Transaction,
@@ -134,7 +134,7 @@ pub struct DiskStorage {
 
 pub static STORAGE: OnceLock<Mutex<DiskStorage>> = OnceLock::new();
 
-const FAUCET_AMT: Amount = Amount::from_sat(100_000);
+const FAUCET_AMT: Amount = Amount::from_sat(10_000);
 
 #[derive(Debug)]
 struct SilentPaymentWallet {
@@ -169,10 +169,16 @@ async fn handle_connection(raw_stream: TcpStream, addr: SocketAddr, our_sp_addre
 
     let peers = PEERMAP.get().expect("Peer Map not initialized");
 
-    let ws_stream = tokio_tungstenite::accept_async(raw_stream)
-        .await
-        .expect("Error during the websocket handshake occurred");
-    debug!("WebSocket connection established");
+    let ws_stream = match tokio_tungstenite::accept_async(raw_stream).await {
+        Ok(stream) => {
+            debug!("WebSocket connection established");
+            stream
+        }
+        Err(e) => {
+            log::error!("WebSocket handshake failed for {}: {}", addr, e);
+            return;
+        }
+    };
 
     // Insert the write part of this peer to the peer map.
     let (tx, rx) = unbounded_channel();
@@ -189,8 +195,8 @@ async fn handle_connection(raw_stream: TcpStream, addr: SocketAddr, our_sp_addre
 
     let init_msg = HandshakeMessage::new(
         our_sp_address.to_string(), 
-        members,
-        processes,
+        OutPointMemberMap(members),
+        OutPointProcessMap(processes),
     );
 
     if let Err(e) = broadcast_message(
@@ -407,9 +413,9 @@ async fn main() -> Result<()> {
         }
     };
 
-    let members: HashMap<Member, OutPoint> = match members_file.load() {
+    let members: HashMap<OutPoint, Member> = match members_file.load() {
         Ok(members) => {
-            let deserialized: MemberOutPointMap = serde_json::from_value(members)?;
+            let deserialized: OutPointMemberMap = serde_json::from_value(members)?;
             deserialized.0
         },
         Err(_) => {