From 872f1f95c89f3d5364dffea25cc931cdd5e1c0d1 Mon Sep 17 00:00:00 2001 From: lejianwen <84855512@qq.com> Date: Sun, 31 Aug 2025 16:10:52 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E4=BC=98=E5=8C=96=E5=B9=B6=E5=8A=A0?= =?UTF-8?q?=E5=85=A5websocket=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/rendezvous_server.rs | 409 ++++++++++++++++++++++++++------------- 1 file changed, 272 insertions(+), 137 deletions(-) diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs index 6f8d7e5..84569d1 100644 --- a/src/rendezvous_server.rs +++ b/src/rendezvous_server.rs @@ -58,13 +58,42 @@ enum Data { const REG_TIMEOUT: i32 = 30_000; type TcpStreamSink = SplitSink, Bytes>; type WsSink = SplitSink, tungstenite::Message>; -enum SinkType { - TcpStream(TcpStreamSink), - Ws(WsSink), +struct SafeWsSink { + sink: WsSink, + encrypt: Option, } -struct Sink { - tx: SinkType, - key: Arc>>, +struct SafeTcpStreamSink { + sink: TcpStreamSink, + encrypt: Option, +} +enum Sink { + // TcpStream(TcpStreamSink), + // Ws(WsSink), + Wss(SafeWsSink), + Tss(SafeTcpStreamSink), +} + +impl Sink { + async fn send(&mut self, msg: &RendezvousMessage) { + if let Ok(mut bytes) = msg.write_to_bytes() { + match self { + // Sink::TcpStream(mut s) => allow_err!(s.send(Bytes::from(bytes)).await), + // Sink::Ws(mut s) => allow_err!(s.send(tungstenite::Message::Binary(bytes)).await), + Sink::Wss(s) => { + if let Some(key) = s.encrypt.as_mut() { + bytes = key.enc(&bytes); + } + allow_err!(s.sink.send(tungstenite::Message::Binary(bytes)).await) + } + Sink::Tss(s) => { + if let Some(key) = s.encrypt.as_mut() { + bytes = key.enc(&bytes); + } + allow_err!(s.sink.send(Bytes::from(bytes)).await) + } + } + } + } } type Sender = mpsc::UnboundedSender; type Receiver = mpsc::UnboundedReceiver; @@ -95,6 +124,7 @@ pub struct RendezvousServer { relay_servers0: Arc, rendezvous_servers: Arc>, inner: Arc, + ws_map: Arc>> } enum LoopFailure { @@ -153,6 +183,7 @@ impl RendezvousServer { secure_tcp_pk_b, secure_tcp_sk_b, }), + ws_map: Arc::new(Mutex::new(HashMap::new())) }; log::info!("mask: {:?}", rs.inner.mask); log::info!("local-ip: {:?}", rs.inner.local_ip); @@ -357,7 +388,13 @@ impl RendezvousServer { // B registered if !rp.id.is_empty() { log::trace!("New peer registered: {:?} {:?}", &rp.id, &addr); - self.update_addr(rp.id, addr, socket).await?; + let request_pk = self.update_addr(rp.id, addr).await; + let mut msg_out = RendezvousMessage::new(); + msg_out.set_register_peer_response(RegisterPeerResponse { + request_pk, + ..Default::default() + }); + socket.send(&msg_out, addr).await?; if self.inner.serial > rp.serial { let mut msg_out = RendezvousMessage::new(); msg_out.set_configure_update(ConfigUpdate { @@ -370,89 +407,25 @@ impl RendezvousServer { } } Some(rendezvous_message::Union::RegisterPk(rk)) => { - if rk.uuid.is_empty() || rk.pk.is_empty() { - return Ok(()); - } - let id = rk.id; - let ip = addr.ip().to_string(); - if id.len() < 6 { - return send_rk_res(socket, addr, UUID_MISMATCH).await; - } else if !self.check_ip_blocker(&ip, &id).await { - return send_rk_res(socket, addr, TOO_FREQUENT).await; - } - let peer = self.pm.get_or(&id).await; - let (changed, ip_changed) = { - let peer = peer.read().await; - if peer.uuid.is_empty() { - (true, false) - } else { - if peer.uuid == rk.uuid { - if peer.info.ip != ip && peer.pk != rk.pk { - log::warn!( - "Peer {} ip/pk mismatch: {}/{:?} vs {}/{:?}", - id, - ip, - rk.pk, - peer.info.ip, - peer.pk, - ); - drop(peer); - return send_rk_res(socket, addr, UUID_MISMATCH).await; - } - } else { - log::warn!( - "Peer {} uuid mismatch: {:?} vs {:?}", - id, - rk.uuid, - peer.uuid - ); - drop(peer); - return send_rk_res(socket, addr, UUID_MISMATCH).await; - } - let ip_changed = peer.info.ip != ip; - ( - peer.uuid != rk.uuid || peer.pk != rk.pk || ip_changed, - ip_changed, - ) + let response = self.handle_register_pk(rk, addr, false).await; + match response { + Err(err) => { + let mut msg_out = RendezvousMessage::new(); + msg_out.set_register_pk_response(RegisterPkResponse { + result: err.into(), + ..Default::default() + }); + socket.send(&msg_out, addr).await?; } - }; - let mut req_pk = peer.read().await.reg_pk; - if req_pk.1.elapsed().as_secs() > 6 { - req_pk.0 = 0; - } else if req_pk.0 > 2 { - return send_rk_res(socket, addr, TOO_FREQUENT).await; - } - req_pk.0 += 1; - req_pk.1 = Instant::now(); - peer.write().await.reg_pk = req_pk; - if ip_changed { - let mut lock = IP_CHANGES.lock().await; - if let Some((tm, ips)) = lock.get_mut(&id) { - if tm.elapsed().as_secs() > IP_CHANGE_DUR { - *tm = Instant::now(); - ips.clear(); - ips.insert(ip.clone(), 1); - } else if let Some(v) = ips.get_mut(&ip) { - *v += 1; - } else { - ips.insert(ip.clone(), 1); - } - } else { - lock.insert( - id.clone(), - (Instant::now(), HashMap::from([(ip.clone(), 1)])), - ); + Ok(res) => { + let mut msg_out = RendezvousMessage::new(); + msg_out.set_register_pk_response(RegisterPkResponse { + result: res.into(), + ..Default::default() + }); + socket.send(&msg_out, addr).await?; } } - if changed { - self.pm.update_pk(id, peer, addr, rk.uuid, rk.pk, ip).await; - } - let mut msg_out = RendezvousMessage::new(); - msg_out.set_register_pk_response(RegisterPkResponse { - result: register_pk_response::Result::OK.into(), - ..Default::default() - }); - socket.send(&msg_out, addr).await? } Some(rendezvous_message::Union::PunchHoleRequest(ph)) => { if self.pm.is_in_memory(&ph.id).await { @@ -519,7 +492,30 @@ impl RendezvousServer { ws: bool, ) -> bool { if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(bytes) { + // log::debug!("Received TCP message from {}: {:?}", addr, msg_in); match msg_in.union { + Some(rendezvous_message::Union::RegisterPeer(rp)) => { + // B registered + if !rp.id.is_empty() { + log::trace!("New peer registered: {:?} {:?}", &rp.id, &addr); + let request_pk = self.update_addr(rp.id, addr).await; + let mut msg_out = RendezvousMessage::new(); + msg_out.set_register_peer_response(RegisterPeerResponse { + request_pk, + ..Default::default() + }); + Self::send_to_sink(sink, msg_out).await; + if self.inner.serial > rp.serial { + let mut msg_out = RendezvousMessage::new(); + msg_out.set_configure_update(ConfigUpdate { + serial: self.inner.serial, + rendezvous_servers: (*self.rendezvous_servers).clone(), + ..Default::default() + }); + Self::send_to_sink(sink, msg_out).await; + } + } + } Some(rendezvous_message::Union::PunchHoleRequest(ph)) => { // there maybe several attempt, so sink can be none if let Some(sink) = sink.take() { @@ -583,14 +579,34 @@ impl RendezvousServer { msg_out.set_test_nat_response(res); Self::send_to_sink(sink, msg_out).await; } - Some(rendezvous_message::Union::RegisterPk(_)) => { - let res = register_pk_response::Result::NOT_SUPPORT; - let mut msg_out = RendezvousMessage::new(); - msg_out.set_register_pk_response(RegisterPkResponse { - result: res.into(), - ..Default::default() - }); - Self::send_to_sink(sink, msg_out).await; + Some(rendezvous_message::Union::RegisterPk(rk)) => { + let response = self.handle_register_pk(rk, addr, ws).await; + match response { + Err(err) => { + let mut msg_out = RendezvousMessage::new(); + msg_out.set_register_pk_response(RegisterPkResponse { + result: err.into(), + ..Default::default() + }); + Self::send_to_sink(sink, msg_out).await; + return false; + } + Ok(res) => { + let mut msg_out = RendezvousMessage::new(); + msg_out.set_register_pk_response(RegisterPkResponse { + result: res.into(), + ..Default::default() + }); + Self::send_to_sink(sink, msg_out).await; + if ws { + // for ws, we can only get addr when register_pk + if let Some(sink) = sink.take() { + self.ws_map.lock().await.insert(try_into_v4(addr), sink); + } + } + return true; + } + } } Some(rendezvous_message::Union::KeyExchange(ex)) => { log::trace!("KeyExchange {:?} <- bytes: {:?}", addr, hex::encode(&bytes)); @@ -612,7 +628,10 @@ impl RendezvousServer { match key { Some(key) => { if let Some(sink) = sink.as_mut() { - sink.key.lock().await.replace(Encrypt::new(key)); + match sink { + Sink::Wss(s) => s.encrypt = Some(Encrypt::new(key)), + Sink::Tss(s) => s.encrypt = Some(Encrypt::new(key)), + } } log::debug!("KeyExchange symetric key set"); return true; @@ -654,13 +673,110 @@ impl RendezvousServer { states } + async fn handle_register_pk( + &mut self, + rk: RegisterPk, + addr: SocketAddr, + ws: bool, + ) -> Result { + if rk.uuid.is_empty() || rk.pk.is_empty() { + return Err(INVALID_ID_FORMAT); + } + let id = rk.id; + let ip = addr.ip().to_string(); + if id.len() < 6 { + return Err(UUID_MISMATCH); + //return Err(send_rk_res(socket, addr, UUID_MISMATCH).await); + } else if !self.check_ip_blocker(&ip, &id).await { + return Err(TOO_FREQUENT); + //return Err(send_rk_res(socket, addr, TOO_FREQUENT).await); + } + let peer = self.pm.get_or(&id).await; + let (changed, ip_changed) = { + let peer = peer.read().await; + if peer.uuid.is_empty() { + (true, false) + } else { + if peer.uuid == rk.uuid { + if peer.info.ip != ip && peer.pk != rk.pk { + log::warn!( + "Peer {} ip/pk mismatch: {}/{:?} vs {}/{:?}", + id, + ip, + rk.pk, + peer.info.ip, + peer.pk, + ); + drop(peer); + return Err(UUID_MISMATCH); + //return Err(send_rk_res(socket, addr, UUID_MISMATCH).await); + } + } else { + log::warn!( + "Peer {} uuid mismatch: {:?} vs {:?}", + id, + rk.uuid, + peer.uuid + ); + drop(peer); + return Err(UUID_MISMATCH); + //return Err(send_rk_res(socket, addr, UUID_MISMATCH).await); + } + let ip_changed = peer.info.ip != ip; + ( + peer.uuid != rk.uuid || peer.pk != rk.pk || ip_changed, + ip_changed, + ) + } + }; + let mut req_pk = peer.read().await.reg_pk; + if req_pk.1.elapsed().as_secs() > 6 { + req_pk.0 = 0; + } else if req_pk.0 > 2 { + return Err(TOO_FREQUENT); + //return Err(send_rk_res(socket, addr, TOO_FREQUENT).await); + } + req_pk.0 += 1; + req_pk.1 = Instant::now(); + peer.write().await.reg_pk = req_pk; + if ip_changed { + let mut lock = IP_CHANGES.lock().await; + if let Some((tm, ips)) = lock.get_mut(&id) { + if tm.elapsed().as_secs() > IP_CHANGE_DUR { + *tm = Instant::now(); + ips.clear(); + ips.insert(ip.clone(), 1); + } else if let Some(v) = ips.get_mut(&ip) { + *v += 1; + } else { + ips.insert(ip.clone(), 1); + } + } else { + lock.insert( + id.clone(), + (Instant::now(), HashMap::from([(ip.clone(), 1)])), + ); + } + } + if changed || ws { + // update peer info，解决tcp过程中不更新在线时间的问题 + self.pm.update_pk(id, peer, addr, rk.uuid, rk.pk, ip).await; + } + Ok(register_pk_response::Result::OK) + // let mut msg_out = RendezvousMessage::new(); + // msg_out.set_register_pk_response(RegisterPkResponse { + // result: register_pk_response::Result::OK.into(), + // ..Default::default() + // }); + // Ok(msg_out) + } + #[inline] async fn update_addr( &mut self, id: String, socket_addr: SocketAddr, - socket: &mut FramedSocket, - ) -> ResultType<()> { + ) -> bool { let (request_pk, ip_change) = if let Some(old) = self.pm.get_in_memory(&id).await { let mut old = old.write().await; let ip = socket_addr.ip(); @@ -690,12 +806,13 @@ impl RendezvousServer { if let Some(old) = ip_change { log::info!("IP change of {} from {} to {}", id, old, socket_addr); } - let mut msg_out = RendezvousMessage::new(); - msg_out.set_register_peer_response(RegisterPeerResponse { - request_pk, - ..Default::default() - }); - socket.send(&msg_out, socket_addr).await + request_pk + // let mut msg_out = RendezvousMessage::new(); + // msg_out.set_register_peer_response(RegisterPeerResponse { + // request_pk, + // ..Default::default() + // }); + // socket.send(&msg_out, socket_addr).await } #[inline] @@ -771,7 +888,7 @@ impl RendezvousServer { ph: PunchHoleRequest, key: &str, ws: bool, - ) -> ResultType<(RendezvousMessage, Option)> { + ) -> ResultType<(RendezvousMessage, Option<(SocketAddr)>)> { let mut ph = ph; if !key.is_empty() && ph.licence_key != key { let mut msg_out = RendezvousMessage::new(); @@ -869,7 +986,8 @@ impl RendezvousServer { ..Default::default() }); } - Ok((msg_out, Some(peer_addr))) + // + Ok((msg_out, Some((peer_addr)))) } else { let mut msg_out = RendezvousMessage::new(); msg_out.set_punch_hole_response(PunchHoleResponse { @@ -909,19 +1027,7 @@ impl RendezvousServer { #[inline] async fn send_to_sink(sink: &mut Option, msg: RendezvousMessage) { if let Some(sink) = sink.as_mut() { - if let Ok(mut bytes) = msg.write_to_bytes() { - if let Some(enc) = &mut sink.key.lock().await.as_mut() { - bytes = enc.enc(&bytes); - } - match &mut sink.tx { - SinkType::TcpStream(s) => { - allow_err!(s.send(Bytes::from(bytes)).await); - } - SinkType::Ws(ws) => { - allow_err!(ws.send(tungstenite::Message::Binary(bytes)).await); - } - } - } + sink.send(&msg).await; } } @@ -946,13 +1052,50 @@ impl RendezvousServer { ) -> ResultType<()> { let (msg, to_addr) = self.handle_punch_hole_request(addr, ph, key, ws).await?; if let Some(addr) = to_addr { - self.tx.send(Data::Msg(msg.into(), addr))?; + let mut sink = self.ws_map.lock().await.remove(&try_into_v4(addr)); + if let Some(s) = sink.as_mut() { + s.send(&msg).await; + } else { + self.tx.send(Data::Msg(msg.into(), addr))?; + } } else { self.send_to_tcp_sync(msg, addr).await?; } Ok(()) } + /** + async fn handle_tcp_punch_hole_request( + &mut self, + addr: SocketAddr, + ph: PunchHoleRequest, + key: &str, + ws: bool, + ) -> ResultType<()> { + let (msg, to_addr) = self.handle_punch_hole_request(addr, ph, key, ws).await?; + log::debug!( + "handle_tcp_punch_hole_request: {:?} to {:?}", + &msg, &to_addr + ); + // + + if let Some(addr) = to_addr { + // self.send_to_tcp_sync(msg, addr).await?; + //这里无法判断目标addr是否是ws，请求无法发送到目标peer + if self.tcp_punch.lock().await.contains_key(&try_into_v4(addr)) { + self.send_to_tcp_sync(msg, addr).await?; //这里如果client使用ws,可以，但如果不用ws就必须用下面的，怎么判断呢 + //这样判断还是不行，如果client从ws切换回udp，tcp_punch中还是有残留，会导致发送失败，应该做个新的hashmap存？？ + } else { + self.tx.send(Data::Msg(msg.into(), addr))?; + } + // self.tx.send(Data::Msg(msg.into(), addr))?; + + } else { + self.send_to_tcp_sync(msg, addr).await?; + } + Ok(()) + }**/ + #[inline] async fn handle_udp_punch_hole_request( &mut self, @@ -1256,10 +1399,10 @@ impl RendezvousServer { }; let ws_stream = tokio_tungstenite::accept_hdr_async(stream, callback).await?; let (a, mut b) = ws_stream.split(); - sink = Some(Sink { - tx: SinkType::Ws(a), - key: Arc::new(Mutex::new(None)), - }); + sink = Some(Sink::Wss(SafeWsSink { + sink: a, + encrypt: None, + })); while let Ok(Some(Ok(msg))) = timeout(30_000, b.next()).await { if let tungstenite::Message::Binary(bytes) = msg { if !self.handle_tcp(&bytes, &mut sink, addr, key, ws).await { @@ -1269,23 +1412,15 @@ impl RendezvousServer { } } else { let (a, mut b) = Framed::new(stream, BytesCodec::new()).split(); - let enc = Arc::new(Mutex::new(None)); - sink = Some(Sink { - tx: SinkType::TcpStream(a), - key: enc.clone(), - }); + sink = Some(Sink::Tss(SafeTcpStreamSink { + sink: a, + encrypt: None, + })); // Avoid key exchange if answering on nat helper port if !key.is_empty() { self.key_exchange_phase1(addr, &mut sink).await; } - while let Ok(Some(Ok(mut bytes))) = timeout(30_000, b.next()).await { - let mut enc_lock = enc.lock().await; - if enc_lock.is_some() { - if let Err(err) = enc_lock.as_mut().unwrap().dec(&mut bytes) { - return Err(err.into()); - } - } - drop(enc_lock); + while let Ok(Some(Ok(bytes))) = timeout(30_000, b.next()).await { if !self.handle_tcp(&bytes, &mut sink, addr, key, ws).await { break; }