diff --git a/gatuna-win/SessionManager.cs b/gatuna-win/SessionManager.cs index 560a646..2e3617e 100644 --- a/gatuna-win/SessionManager.cs +++ b/gatuna-win/SessionManager.cs @@ -100,7 +100,7 @@ sealed class SessionManager : IDisposable case Frame.Close close: if (_sessions.TryRemove(close.SessionId, out var s)) - s.Dispose(); + s.OnRemoteClose(); break; case Frame.Pong pong: @@ -394,6 +394,7 @@ sealed class Session( readonly RecvState _recv = new(); readonly object _sendLock = new(); readonly object _recvLock = new(); + volatile bool _closeSent; public void Start() { @@ -403,6 +404,17 @@ sealed class Session( _ = RetransmitTimer(); } + /// + /// Called when a CLOSE frame arrives from the server. The server has + /// already torn down its side — we just need to flush queued data to + /// the local socket and close gracefully. Do NOT echo CLOSE back. + /// + public void OnRemoteClose() + { + _cts.Cancel(); + _deliverChannel.Writer.TryComplete(); + } + /// Handle a DATA frame from the tunnel. public void HandleData(Frame.Data data) { @@ -440,10 +452,10 @@ sealed class Session( async Task PumpSocketToTunnel() { + var buf = new byte[Proto.MaxPayload]; try { var stream = client.GetStream(); - var buf = new byte[Proto.MaxPayload]; using var reg = _cts.Token.Register(() => client.Dispose()); while (!_cts.IsCancellationRequested) { @@ -476,6 +488,8 @@ sealed class Session( } } catch { } + _cts.Cancel(); + _deliverChannel.Writer.TryComplete(); SendClose(); onClosed(); } @@ -485,10 +499,11 @@ sealed class Session( try { var stream = client.GetStream(); - await foreach (var payload in _deliverChannel.Reader.ReadAllAsync(_cts.Token)) - await stream.WriteAsync(payload, _cts.Token); + await foreach (var payload in _deliverChannel.Reader.ReadAllAsync()) + await stream.WriteAsync(payload); } catch { } + finally { client.Close(); } } async Task RetransmitTimer() @@ -539,7 +554,12 @@ sealed class Session( link.SendTo(serverMac, new Frame.Data(sessionId, seq, ackSeq, [])); } - void SendClose() => link.SendTo(serverMac, new Frame.Close(sessionId, null)); + void SendClose() + { + if (_closeSent) return; + _closeSent = true; + link.SendTo(serverMac, new Frame.Close(sessionId, null)); + } public void Dispose() { diff --git a/gatunad/src/main.rs b/gatunad/src/main.rs index 56974fa..62f24da 100644 --- a/gatunad/src/main.rs +++ b/gatunad/src/main.rs @@ -224,8 +224,9 @@ async fn handle_frame( } } Frame::Close { session_id, reason: _ } => { - info!("CLOSE session {session_id} from {src}"); - store.lock().expect("store poisoned").remove(&session_id); + if store.lock().expect("store poisoned").remove(&session_id).is_some() { + info!("CLOSE session {session_id} from {src}"); + } } Frame::Ping { nonce } => { let pong = Frame::Pong { nonce }; diff --git a/gatunad/src/session.rs b/gatunad/src/session.rs index 3886703..0a1551d 100644 --- a/gatunad/src/session.rs +++ b/gatunad/src/session.rs @@ -275,12 +275,12 @@ pub fn spawn_pump( } if should_close { + store.lock().expect("store poisoned").remove(&session_id); let close = Frame::Close { session_id, reason: Some(REASON_MAX_RETRIES), }; let _ = tx.try_send((peer_mac, close.encode())); - store.lock().expect("store poisoned").remove(&session_id); break; } } @@ -330,8 +330,8 @@ pub fn spawn_pump( break; } } + store.lock().expect("store poisoned").remove(&session_id); let close = Frame::Close { session_id, reason: None }; let _ = tx.send((peer_mac, close.encode())).await; - store.lock().expect("store poisoned").remove(&session_id); }); }