gatuna/win/gatuna-client/SessionManager.cs

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

namespace gatuna_client;

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

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

    // 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 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)
                        _pending.Client.Dispose();
                }
                ProcessQueue();
                break;

            case Frame.Data data:
                if (_sessions.TryGetValue(data.SessionId, out var session))
                    session.Deliver(data.Payload);
                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;
        }
    }

    /// <summary>
    /// Start a local TCP listener for the given upstream. Returns the mirror
    /// port, or 0 on failure.
    /// </summary>
    public int StartListener(UpstreamEntry upstream)
    {
        var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var port = ((IPEndPoint)listener.LocalEndpoint).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);
        }
    }

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

    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));
    }

    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, 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()
    {
        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)
{
    public TcpClient Client { get; } = client;
    public byte UpstreamId { get; } = upstreamId;
}

sealed class Session(
    uint sessionId,
    TcpClient client,
    byte[] serverMac,
    TunnelLink link,
    Action onClosed,
    Action<string>? log) : IDisposable
{
    readonly CancellationTokenSource _cts = new();
    readonly Channel<byte[]> _incoming = Channel.CreateBounded<byte[]>(256);

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

    public void Deliver(byte[] payload)
    {
        if (!_incoming.Writer.TryWrite(payload))
            log?.Invoke($"session {sessionId}: incoming channel full");
    }

    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;
                link.SendTo(serverMac, new Frame.Data(sessionId, buf[..n]));
            }
        }
        catch { }
        SendClose();
        onClosed();
    }

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

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

    public void Dispose()
    {
        _cts.Cancel();
        _incoming.Writer.TryComplete();
        SendClose();
        try { client.Dispose(); } catch { }
        _cts.Dispose();
    }
}