From 9ea2978c8ddca7edf49341076249a3bd3750fd14 Mon Sep 17 00:00:00 2001 From: Klesti Selimaj Date: Wed, 18 Feb 2026 22:01:45 +0100 Subject: [PATCH] Working client --- examples/client.rs | 4 +- examples/server.rs | 2 +- src/handshake.rs | 92 ++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/session.rs | 20 ++-------- 5 files changed, 98 insertions(+), 21 deletions(-) diff --git a/examples/client.rs b/examples/client.rs index 5d09db5..853832b 100644 --- a/examples/client.rs +++ b/examples/client.rs @@ -1,12 +1,10 @@ use std::sync::Arc; -use tokio::net::TcpStream; use session_rs::session::Session; #[tokio::main(flavor = "current_thread")] async fn main() -> session_rs::Result<()> { - let stream = TcpStream::connect("127.0.0.1:8080").await?; - let session = Arc::new(Session::new(stream).await?); + let session = Arc::new(Session::new_server("127.0.0.1:8080", "/").await?); // Spawn read loop let read_session = Arc::clone(&session); diff --git a/examples/server.rs b/examples/server.rs index 35b294c..5ab6336 100644 --- a/examples/server.rs +++ b/examples/server.rs @@ -14,7 +14,7 @@ async fn main() -> session_rs::Result<()> { tokio::spawn(async move { // Wrap session in Arc so tasks can share it - let session = match Session::new(stream).await { + let session = match Session::new_client(stream).await { Ok(s) => Arc::new(s), Err(e) => { eprintln!("Handshake failed: {:?}", e); diff --git a/src/handshake.rs b/src/handshake.rs index a3c98e8..56edc63 100644 --- a/src/handshake.rs +++ b/src/handshake.rs @@ -1,8 +1,13 @@ +use sha1::{Digest, Sha1}; +use std::sync::Arc; use tokio::{ io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, net::TcpStream, + sync::Mutex, }; +use crate::session::Session; + pub async fn handle_websocket_handshake(stream: &mut TcpStream) -> std::io::Result<()> { let (read_half, mut write_half) = stream.split(); let mut reader = BufReader::new(read_half); @@ -75,3 +80,90 @@ pub async fn handle_websocket_handshake(stream: &mut TcpStream) -> std::io::Resu write_half.write_all(response.as_bytes()).await?; Ok(()) } + +impl Session { + pub async fn new_client(mut stream: TcpStream) -> crate::Result { + crate::handshake::handle_websocket_handshake(&mut stream).await?; + + let (read, write) = stream.into_split(); + + Ok(Self { + reader: Arc::new(Mutex::new(read)), + writer: Arc::new(Mutex::new(write)), + id: rand::random(), + }) + } + + /// Connect to a WebSocket server and perform the handshake + pub async fn new_server(addr: &str, path: &str) -> crate::Result { + // 1. TCP connect + let mut stream = TcpStream::connect(addr).await?; + + // 2. Generate Sec-WebSocket-Key + let key_bytes: [u8; 16] = rand::random(); + let key = base64::encode(&key_bytes); + + // 3. Send HTTP Upgrade request + let request = format!( + "GET {} HTTP/1.1\r\n\ + Host: {}\r\n\ + Upgrade: websocket\r\n\ + Connection: Upgrade\r\n\ + Sec-WebSocket-Key: {}\r\n\ + Sec-WebSocket-Version: 13\r\n\ + \r\n", + path, addr, key + ); + stream.write_all(request.as_bytes()).await?; + stream.flush().await?; + + // 4. Read HTTP response + let mut reader = BufReader::new(&mut stream); + let mut status_line = String::new(); + reader.read_line(&mut status_line).await?; + if !status_line.starts_with("HTTP/1.1 101") { + return Err(crate::Error::HandshakeFailed(format!( + "Expected 101 Switching Protocols, got: {}", + status_line.trim_end() + ))); + } + + // Read headers + let mut sec_accept = None; + loop { + let mut line = String::new(); + reader.read_line(&mut line).await?; + let line = line.trim_end(); + if line.is_empty() { + break; // end of headers + } + if let Some((k, v)) = line.split_once(':') { + if k.eq_ignore_ascii_case("sec-websocket-accept") { + sec_accept = Some(v.trim().to_string()); + } + } + } + + // 5. Verify Sec-WebSocket-Accept + let expected = { + let mut sha1 = Sha1::new(); + sha1.update(key.as_bytes()); + sha1.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); + base64::encode(sha1.finalize()) + }; + if sec_accept.as_deref() != Some(expected.as_str()) { + return Err(crate::Error::HandshakeFailed( + "Sec-WebSocket-Accept mismatch".into(), + )); + } + + // 6. Upgrade succeeded, split stream + let (read, write) = stream.into_split(); + + Ok(Self { + reader: Arc::new(Mutex::new(read)), + writer: Arc::new(Mutex::new(write)), + id: rand::random(), + }) + } +} diff --git a/src/lib.rs b/src/lib.rs index 64c4fe1..0af8d9a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,7 @@ pub enum Error { Io(std::io::Error), Json(serde_json::Error), InvalidFrame(String), + HandshakeFailed(String), ConnectionClosed, } diff --git a/src/session.rs b/src/session.rs index 333af59..732f745 100644 --- a/src/session.rs +++ b/src/session.rs @@ -9,23 +9,9 @@ use tokio::{ }; pub struct Session { - reader: Arc>, - writer: Arc>, - id: u64, -} - -impl Session { - pub async fn new(mut stream: TcpStream) -> crate::Result { - crate::handshake::handle_websocket_handshake(&mut stream).await?; - - let (read, write) = stream.into_split(); - - Ok(Self { - reader: Arc::new(Mutex::new(read)), - writer: Arc::new(Mutex::new(write)), - id: rand::random(), - }) - } + pub(crate) reader: Arc>, + pub(crate) writer: Arc>, + pub(crate) id: u64, } impl Clone for Session {