gatuna/gatuna-win/SessionManager.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using System.Threading.Channels;

namespace gatuna;

sealed class SessionManager : IDisposable
{
    TunnelLink? _link;
    readonly ConcurrentDictionary<uint, Session> _sessions = new();
    readonly ConcurrentDictionary<int, ListenerState> _listeners = new();

    byte[]? _serverMac;
    string _serverHostname = "";
    UpstreamEntry[] _upstreams = [];
    PingTest? _pingTest;

    // Serialized OPEN: only one outstanding at a time.
    readonly object _openLock = new();
    PendingOpen? _pending;
    readonly Queue<PendingOpen> _openQueue = new();

    public event Action<string>? Log;
    public event Action<string, byte[], UpstreamEntry[]>? ManifestReceived;

    public UpstreamEntry[] Upstreams => _upstreams;
    public byte[]? ServerMac => _serverMac;
    public string ServerHostname => _serverHostname;
    public TunnelLink? Link => _link;

    public void AttachLink(TunnelLink link)
    {
        _link = link;
        link.FrameReceived += HandleFrame;
    }

    public void DetachLink()
    {
        if (_link != null)
            _link.FrameReceived -= HandleFrame;
        _link = null;
    }

    public void Discover()
    {
        if (_link == null) return;
        _link.SendBroadcast(new Frame.Discover());
    }

    public PingTest? StartPing()
    {
        if (_link == null || _serverMac == null)
            return null;
        _pingTest?.Stop();
        _pingTest = new PingTest(_link, _serverMac);
        _pingTest.Start();
        return _pingTest;
    }

    public void StopPing()
    {
        _pingTest?.Stop();
        _pingTest = null;
    }

    public void HandleFrame(Frame frame, byte[] srcMac)
    {
        switch (frame)
        {
            case Frame.Manifest manifest:
                _serverMac = srcMac;
                _serverHostname = manifest.Hostname;
                _upstreams = manifest.Entries;
                Log?.Invoke($"manifest: {manifest.Entries.Length} upstreams from {manifest.Hostname} ({BitConverter.ToString(srcMac)})");
                ManifestReceived?.Invoke(manifest.Hostname, srcMac, manifest.Entries);
                break;

            case Frame.OpenAck ack:
                HandleOpenAck(ack, srcMac);
                break;

            case Frame.OpenNak nak:
                Log?.Invoke($"OPEN_NAK upstream {nak.UpstreamId} reason {nak.Reason}");
                lock (_openLock)
                {
                    if (_pending != null)
                        NetUtil.RstClose(_pending.Client);
                }
                ProcessQueue();
                break;

            case Frame.Data data:
                if (_sessions.TryGetValue(data.SessionId, out var session))
                    session.HandleData(data);
                else if (_link != null && _serverMac != null)
                    _link.SendTo(_serverMac,
                        new Frame.Close(data.SessionId, Proto.ReasonUnknownSession));
                break;

            case Frame.Close close:
                if (_sessions.TryRemove(close.SessionId, out var s))
                    s.Dispose();
                break;

            case Frame.Pong pong:
                _pingTest?.HandlePong(pong.Nonce);
                break;

            case Frame.Ping _:
                break;
        }
    }

    /// <summary>
    /// Compute a deterministic mirror port from the server MAC and the
    /// upstream port. XOR the upstream port with (mac[0]<<8 | mac[5]),
    /// then ensure the result is outside the privileged range.
    /// </summary>
    static ushort ComputeMirrorPort(byte[] serverMac, ushort upstreamPort)
    {
        var k = (ushort)((serverMac[0] << 8) | serverMac[5]);
        var port = (ushort)(upstreamPort ^ k);
        if (port < 1024)
            port += 1024;
        return port;
    }

    /// <summary>
    /// Start a local TCP listener for the given upstream. Returns the mirror
    /// port, or 0 on failure.
    /// </summary>
    public int StartListener(UpstreamEntry upstream)
    {
        if (_serverMac == null)
            return 0;

        var preferred = ComputeMirrorPort(_serverMac, upstream.Port);

        // Try the deterministic port first; fall back to OS assignment.
        TcpListener listener;
        int port;
        try
        {
            listener = new TcpListener(IPAddress.Loopback, preferred);
            listener.Start();
            port = ((IPEndPoint)listener.LocalEndpoint).Port;
        }
        catch
        {
            listener = new TcpListener(IPAddress.Loopback, 0);
            listener.Start();
            port = ((IPEndPoint)listener.LocalEndpoint).Port;
            Log?.Invoke($"port {preferred} in use, fell back to {port}");
        }

        var state = new ListenerState(listener, upstream);
        _listeners[port] = state;
        _ = AcceptLoop(state);
        Log?.Invoke($"listening 127.0.0.1:{port} -> upstream {upstream.Id} ({upstream.ProtoName}:{upstream.Port})");
        return port;
    }

    public void StopListener(int port)
    {
        if (_listeners.TryRemove(port, out var state))
        {
            state.Listener.Stop();
            Log?.Invoke($"stopped listener port {port}");
        }
    }

    async Task AcceptLoop(ListenerState state)
    {
        while (true)
        {
            TcpClient client;
            try
            {
                client = await state.Listener.AcceptTcpClientAsync();
            }
            catch { break; }
            EnqueueOpen(client, state.Upstream.Id, state.Upstream.Protocol);
        }
    }

    void EnqueueOpen(TcpClient client, byte upstreamId, byte proto)
    {
        lock (_openLock)
        {
            if (_pending == null)
            {
                _pending = new PendingOpen(client, upstreamId, proto);
                SendOpen(_pending);
            }
            else
            {
                _openQueue.Enqueue(new PendingOpen(client, upstreamId, proto));
            }
        }
    }

    void SendOpen(PendingOpen po)
    {
        if (_link == null || _serverMac == null)
        {
            Log?.Invoke("no server; cannot OPEN");
            po.Client.Dispose();
            return;
        }
        _link.SendTo(_serverMac, new Frame.Open(po.UpstreamId, po.Proto));
    }

    void ProcessQueue()
    {
        lock (_openLock)
        {
            if (_openQueue.Count > 0)
            {
                _pending = _openQueue.Dequeue();
                SendOpen(_pending);
            }
            else
            {
                _pending = null;
            }
        }
    }

    void HandleOpenAck(Frame.OpenAck ack, byte[] srcMac)
    {
        PendingOpen? po;
        lock (_openLock)
            po = _pending;

        if (po == null || po.UpstreamId != ack.UpstreamId)
        {
            Log?.Invoke($"OPEN_ACK upstream {ack.UpstreamId} session {ack.SessionId} — no matching pending");
            return;
        }

        var session = new Session(
            ack.SessionId, ack.Proto, po.Client, srcMac, _link!,
            () => _sessions.TryRemove(ack.SessionId, out _),
            msg => Log?.Invoke(msg));
        _sessions[ack.SessionId] = session;
        session.Start();
        Log?.Invoke($"session {ack.SessionId} upstream {ack.UpstreamId} established");
        ProcessQueue();
    }

    public void StopAll()
    {
        StopPing();
        foreach (var kv in _listeners)
            kv.Value.Listener.Stop();
        _listeners.Clear();
        foreach (var s in _sessions.Values)
            s.Dispose();
        _sessions.Clear();
    }

    public void Dispose()
    {
        DetachLink();
        StopAll();
    }
}

sealed class ListenerState(TcpListener listener, UpstreamEntry upstream)
{
    public TcpListener Listener { get; } = listener;
    public UpstreamEntry Upstream { get; } = upstream;
}

sealed class PendingOpen(TcpClient client, byte upstreamId, byte proto)
{
    public TcpClient Client { get; } = client;
    public byte UpstreamId { get; } = upstreamId;
    public byte Proto { get; } = proto;
}

/// L2 reliability: sender-side state.
class SendState
{
    public uint SendSeq;
    public uint AckedSeq;
    // seq -> (frame_bytes, send_time_ticks, retry_count)
    public readonly SortedList<uint, (byte[], long, uint)> RetransmitBuffer = new();

    public uint NextSeq()
    {
        var s = SendSeq;
        SendSeq++;
        return s;
    }

    public void RecordSent(uint seq, byte[] frameBytes)
    {
        RetransmitBuffer[seq] = (frameBytes, Environment.TickCount64, 0);
    }

    public void ProcessAck(uint ackSeq)
    {
        var toRemove = RetransmitBuffer.Keys
            .Where(k => k <= ackSeq)
            .ToList();
        foreach (var k in toRemove)
            RetransmitBuffer.Remove(k);
        if (ackSeq > AckedSeq)
            AckedSeq = ackSeq;
    }

    /// Returns frames to retransmit and whether the session should close.
    public (List<byte[]> resend, bool shouldClose) CheckRetransmit()
    {
        var resend = new List<byte[]>();
        var shouldClose = false;
        var now = Environment.TickCount64;
        const int timeoutMs = 5;
        const uint maxRetries = 10;

        foreach (var kv in RetransmitBuffer.ToList())
        {
            if (now - kv.Value.Item2 > timeoutMs)
            {
                if (kv.Value.Item3 >= maxRetries)
                {
                    shouldClose = true;
                    break;
                }
                RetransmitBuffer[kv.Key] = (kv.Value.Item1, now, kv.Value.Item3 + 1);
                resend.Add(kv.Value.Item1);
            }
        }
        return (resend, shouldClose);
    }
}

/// L2 reliability: receiver-side state.
class RecvState
{
    public uint ExpectedSeq;
    public uint DeliverSeq;
    // seq -> payload (out-of-order buffer)
    public readonly SortedList<uint, byte[]> ReceiveBuffer = new();

    /// Process an incoming DATA frame. Returns (payloads to deliver, needAck).
    public (List<byte[]> deliver, bool needAck) ProcessData(uint seq, byte[] payload)
    {
        if (seq < ExpectedSeq)
        {
            // Duplicate.
            return ([], true);
        }
        if (seq == ExpectedSeq)
        {
            // In-order: deliver and drain buffer.
            var deliver = new List<byte[]> { payload };
            ExpectedSeq++;
            DeliverSeq = ExpectedSeq - 1;
            while (ReceiveBuffer.Remove(ExpectedSeq, out var buffered))
            {
                deliver.Add(buffered);
                ExpectedSeq++;
                DeliverSeq = ExpectedSeq - 1;
            }
            return (deliver, false);
        }
        else
        {
            // Out-of-order: buffer.
            ReceiveBuffer[seq] = payload;
            return ([], true);
        }
    }

    /// The ack_seq to report in outgoing DATA frames.
    public uint CurrentAckSeq => ExpectedSeq == 0 ? uint.MaxValue : ExpectedSeq - 1;
}

sealed class Session(
    uint sessionId,
    byte proto,
    TcpClient client,
    byte[] serverMac,
    TunnelLink link,
    Action onClosed,
    Action<string>? log) : IDisposable
{
    readonly CancellationTokenSource _cts = new();
    readonly bool _isTcp = proto == Proto.ProtoTcp;
    readonly SendState _send = new();
    readonly RecvState _recv = new();
    readonly object _sendLock = new();
    readonly object _recvLock = new();

    public void Start()
    {
        _ = PumpSocketToTunnel();
        _ = PumpTunnelToSocket();
        if (_isTcp)
            _ = RetransmitTimer();
    }

    /// Handle a DATA frame from the tunnel.
    public void HandleData(Frame.Data data)
    {
        if (!_isTcp)
        {
            // Best-effort: deliver directly, ignore seq/ack.
            if (data.Payload.Length > 0)
                _deliverChannel.Writer.TryWrite(data.Payload);
            return;
        }

        // TCP: process ack_seq to advance send window.
        lock (_sendLock)
            _send.ProcessAck(data.AckSeq);

        // Process seq for in-order delivery.
        List<byte[]> deliver;
        bool needAck;
        lock (_recvLock)
            (deliver, needAck) = _recv.ProcessData(data.Seq, data.Payload);

        // Deliver to the TCP socket via the channel.
        foreach (var chunk in deliver)
        {
            if (!_deliverChannel.Writer.TryWrite(chunk))
                log?.Invoke($"session {sessionId}: deliver channel full");
        }

        // Send pure ACK if duplicate or out-of-order.
        if (needAck)
            SendPureAck();
    }

    readonly Channel<byte[]> _deliverChannel = Channel.CreateBounded<byte[]>(256);

    async Task PumpSocketToTunnel()
    {
        try
        {
            var stream = client.GetStream();
            var buf = new byte[Proto.MaxPayload];
            using var reg = _cts.Token.Register(() => client.Dispose());
            while (!_cts.IsCancellationRequested)
            {
                var n = await stream.ReadAsync(buf, _cts.Token);
                if (n == 0) break;

                uint seq = 0, ackSeq = 0;
                byte[]? frameBytes = null;

                if (_isTcp)
                {
                    lock (_sendLock)
                    {
                        seq = _send.NextSeq();
                        lock (_recvLock)
                            ackSeq = _recv.CurrentAckSeq;
                    }
                }

                var frame = new Frame.Data(sessionId, seq, ackSeq, buf[..n]);

                if (_isTcp)
                {
                    frameBytes = FrameCodec.Encode(frame);
                    lock (_sendLock)
                        _send.RecordSent(seq, frameBytes);
                }

                link.SendTo(serverMac, frame);
            }
        }
        catch { }
        SendClose();
        onClosed();
    }

    async Task PumpTunnelToSocket()
    {
        try
        {
            var stream = client.GetStream();
            await foreach (var payload in _deliverChannel.Reader.ReadAllAsync(_cts.Token))
                await stream.WriteAsync(payload, _cts.Token);
        }
        catch { }
    }

    async Task RetransmitTimer()
    {
        using var timer = new PeriodicTimer(TimeSpan.FromMilliseconds(1));
        try
        {
            while (!_cts.IsCancellationRequested)
            {
                await timer.WaitForNextTickAsync(_cts.Token);

                List<byte[]>? resend = null;
                bool shouldClose = false;
                lock (_sendLock)
                    (resend, shouldClose) = _send.CheckRetransmit();

                if (resend != null)
                {
                    foreach (var frameBytes in resend)
                    {
                        // Send raw bytes directly (already encoded).
                        link.SendRaw(serverMac, frameBytes);
                    }
                }

                if (shouldClose)
                {
                    link.SendTo(serverMac, new Frame.Close(sessionId, Proto.ReasonMaxRetries));
                    onClosed();
                    break;
                }
            }
        }
        catch { }
    }

    void SendPureAck()
    {
        if (!_isTcp) return;
        uint seq, ackSeq;
        lock (_sendLock)
        {
            seq = _send.NextSeq();
            lock (_recvLock)
                ackSeq = _recv.CurrentAckSeq;
        }
        // Pure ACK: empty payload. Not stored in retransmit buffer.
        link.SendTo(serverMac, new Frame.Data(sessionId, seq, ackSeq, []));
    }

    void SendClose() => link.SendTo(serverMac, new Frame.Close(sessionId, null));

    public void Dispose()
    {
        _cts.Cancel();
        _deliverChannel.Writer.TryComplete();
        SendClose();
        NetUtil.RstClose(client);
        _cts.Dispose();
    }
}

/// Close a TcpClient with a TCP RST instead of a FIN.
static partial class NetUtil
{
    public static void RstClose(TcpClient c)
    {
        try
        {
            c.LingerState = new LingerOption(true, 0);
            c.Close();
        }
        catch { }
    }
}