1  use tokio::net::TcpListener;
 2  use tokio::task;
 3  use abzu_transport::transports::websocket::WsStream;
 4  use abzu_transport::AbzuInterface;
 5  
 6  #[tokio::test]
 7  async fn test_pfs_handshake_roundtrip() {
 8      // 1. Setup a standard TCP listener
 9      let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10      let addr_socket = listener.local_addr().unwrap();
11      let addr_str = addr_socket.to_string();
12  
13      // 2. Spawn the listener task (Server)
14      let server_task = task::spawn(async move {
15          // Accept the raw TCP connection
16          let (stream, _) = listener.accept().await.unwrap();
17          
18          // Accept and perform the PFS handshake
19          match WsStream::accept_pfs(stream).await {
20              Ok(handshake) => {
21                  let stream = handshake.stream;
22                  let peer_pubkey = handshake.peer_pubkey;
23                  
24                  // Echo the peer's public key back to prove we have it
25                  let peer_key_vec = peer_pubkey.to_vec();
26                  stream.send(&peer_key_vec).await.unwrap();
27                  
28                  // Also receive a message (echo test)
29                  let msg = stream.recv().await.unwrap();
30                  assert_eq!(msg, b"client_secret_payload");
31                  
32                  // Return server's success status
33                  true
34              },
35              Err(e) => {
36                  println!("Server handshake failed: {}", e);
37                  false
38              }
39          }
40      });
41  
42      // 3. Spawn the client task
43      let addr_for_client = addr_str.clone();
44      let client_task = task::spawn(async move {
45          // Connect and perform PFS handshake
46          match WsStream::connect_pfs(&addr_for_client).await {
47              Ok(handshake) => {
48                  let stream = handshake.stream;
49                  
50                  // Send a secret payload
51                  stream.send(b"client_secret_payload").await.unwrap();
52                  
53                  // Receive the server's echo of our public key (not really standard, just for this test)
54                  let msg = stream.recv().await.unwrap();
55                  
56                  // We don't easily know our own ephemeral key here to verify, but we can check format
57                  assert_eq!(msg.len(), 32);
58                  
59                  true
60              },
61              Err(e) => {
62                  println!("Client handshake failed: {}", e);
63                  false
64              }
65          }
66      });
67  
68      // 4. Wait for both to complete
69      let server_result = server_task.await.unwrap();
70      let client_result = client_task.await.unwrap();
71  
72      assert!(server_result, "Server failed");
73      assert!(client_result, "Client failed");
74  }