JJServiceManagerRelay/server/RelayHandler.cs

using System.Buffers;
using System.Net.WebSockets;
using System.Security.Cryptography;
using System.Text;

namespace JJServiceManagerRelay;

/// <summary>/relay/{세션ID}?token=…&amp;side=manager|agent — 두 연결을 이어 붙여 바이트를 그대로 넘김</summary>
public sealed class RelayHandler
{
    private readonly Hub _hub;
    private static readonly TimeSpan PairWait = TimeSpan.FromSeconds(30);

    public RelayHandler(Hub hub) => _hub = hub;

    public async Task RunAsync(HttpContext ctx, string sessionId)
    {
        var side = ctx.Request.Query["side"].ToString();
        var token = ctx.Request.Query["token"].ToString();
        if (!Guid.TryParse(sessionId, out var sid) || !_hub.Sessions.TryGetValue(sid, out var s) || (side != "manager" && side != "agent"))
        {
            ctx.Response.StatusCode = 404;
            return;
        }
        var expect = side == "manager" ? s.ManagerToken : s.AgentToken;
        bool tokenOk = CryptographicOperations.FixedTimeEquals(Encoding.UTF8.GetBytes(token), Encoding.UTF8.GetBytes(expect));
        lock (s.Lock)
        {
            bool used = side == "manager" ? s.ManagerUsed : s.AgentUsed;
            if (!tokenOk || used || s.Finished || s.Direct)
            {
                ctx.Response.StatusCode = 403;
                return;
            }
            if (side == "manager") s.ManagerUsed = true; else s.AgentUsed = true;
        }

        using var ws = await ctx.WebSockets.AcceptWebSocketAsync();
        WebSocket other = null;
        lock (s.Lock)
        {
            if (side == "manager") s.ManagerWs = ws; else s.AgentWs = ws;
            if (s.ManagerWs != null && s.AgentWs != null && !s.Paired)
            {
                s.Paired = true;
                other = side == "manager" ? s.AgentWs : s.ManagerWs;
            }
        }

        if (other == null)
        {
            // 상대를 기다림 (상대 쪽 요청이 이어 붙이기를 맡음)
            await Task.WhenAny(s.Done.Task, s.Wake.Task, Task.Delay(PairWait, ctx.RequestAborted).ContinueWith(_ => { }));
            bool paired;
            lock (s.Lock) paired = s.Paired;
            if (!paired)
            {
                await CloseAsync(ws, s.Direct ? "직접 연결됨" : "상대가 연결되지 않았습니다.");
                if (!s.Direct) await _hub.FinishSessionAsync(s, "failed", $"중계: {(side == "manager" ? "대상 PC" : "관리 PC")}가 30초 안에 붙지 않음");
                return;
            }
            await s.Done.Task; // 이어 붙인 쪽이 끝낼 때까지 이 요청을 살려 둠
            return;
        }

        // 두 번째로 붙은 쪽: 양방향으로 넘김
        await _hub.Db.ConnectionModeAsync(s.Id, "relay");
        var managerWs = side == "manager" ? ws : other;
        var agentWs = side == "manager" ? other : ws;
        _hub.Log.Info($"중계 시작: {s.Username} → {s.DeviceName} ({s.Id})");
        string reason = "정상 종료";
        try
        {
            var t1 = PumpAsync(managerWs, agentWs, n => Interlocked.Add(ref s.BytesToDevice, n));
            var t2 = PumpAsync(agentWs, managerWs, n => Interlocked.Add(ref s.BytesFromDevice, n));
            var first = await Task.WhenAny(t1, t2);
            reason = await first;
            // 한쪽이 끝나면 다른 쪽도 닫음
            await CloseAsync(managerWs, reason);
            await CloseAsync(agentWs, reason);
            // 상대가 닫기 응답을 안 하면 2초 뒤 강제로 끊음
            await Task.WhenAny(Task.WhenAll(t1, t2), Task.Delay(2000));
            try { managerWs.Abort(); } catch { }
            try { agentWs.Abort(); } catch { }
        }
        catch (Exception ex)
        {
            reason = "중계 오류: " + ex.Message;
        }
        finally
        {
            _hub.Log.Info($"중계 끝: {s.Username} → {s.DeviceName} ({reason}, 보냄 {s.BytesToDevice}B / 받음 {s.BytesFromDevice}B)");
            await _hub.FinishSessionAsync(s, "relay", reason);
        }
    }

    /// <summary>from에서 받은 프레임을 to로 그대로 보냄. 끝난 사유를 돌려줌</summary>
    private static async Task<string> PumpAsync(WebSocket from, WebSocket to, Action<long> count)
    {
        var buf = ArrayPool<byte>.Shared.Rent(64 * 1024);
        try
        {
            while (true)
            {
                var r = await from.ReceiveAsync(buf, CancellationToken.None);
                if (r.MessageType == WebSocketMessageType.Close) return "상대가 연결을 닫음";
                count(r.Count);
                await to.SendAsync(new ArraySegment<byte>(buf, 0, r.Count), r.MessageType, r.EndOfMessage, CancellationToken.None);
            }
        }
        catch (WebSocketException)
        {
            return "연결 끊김";
        }
        catch (ObjectDisposedException)
        {
            return "연결 끊김";
        }
        catch (OperationCanceledException)
        {
            return "연결 끊김";
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buf);
        }
    }

    private static async Task CloseAsync(WebSocket ws, string reason)
    {
        try
        {
            using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(3));
            var b = reason ?? "";
            while (Encoding.UTF8.GetByteCount(b) > 120) b = b[..^1];
            if (ws.State == WebSocketState.Open || ws.State == WebSocketState.CloseReceived)
                await ws.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, b, cts.Token);
        }
        catch { }
    }
}