natsnet/dotnet/src/ZB.MOM.NatsNet.Server/ClientConnection.Routes.cs

// Copyright 2012-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");

using System.Text;
using System.Text.Json;
using ZB.MOM.NatsNet.Server.Auth;
using ZB.MOM.NatsNet.Server.Internal;
using ZB.MOM.NatsNet.Server.Protocol;

namespace ZB.MOM.NatsNet.Server;

public sealed partial class ClientConnection
{
    internal void RemoveReplySub(Subscription? sub)
    {
        if (sub?.Sid is not { Length: > 0 } sid || Server is not NatsServer server)
            return;

        var sidText = Encoding.ASCII.GetString(sid);
        var sep = sidText.IndexOf(' ');
        if (sep <= 0)
            return;

        var accountName = sidText[..sep];
        var (account, _) = server.LookupAccount(accountName);
        account?.Sublist?.Remove(sub);

        lock (_mu)
        {
            Subs.Remove(sidText);
        }
    }

    internal Exception? ProcessAccountSub(byte[] arg)
    {
        _ = arg;
        // Gateway account-sub propagation is owned by gateway sessions.
        return null;
    }

    internal void ProcessAccountUnsub(byte[] arg)
    {
        _ = arg;
        // Gateway account-unsub propagation is owned by gateway sessions.
    }

    internal Exception? ProcessRoutedOriginClusterMsgArgs(byte[] arg) =>
        ProtocolParser.ProcessRoutedOriginClusterMsgArgs(ParseCtx, arg);

    internal Exception? ProcessRoutedHeaderMsgArgs(byte[] arg) =>
        ProtocolParser.ProcessRoutedHeaderMsgArgs(ParseCtx, arg);

    internal Exception? ProcessRoutedMsgArgs(byte[] arg) =>
        ProtocolParser.ProcessRoutedMsgArgs(ParseCtx, arg);

    internal void ProcessInboundRoutedMsg(byte[] msg)
    {
        _in.Msgs++;
        _in.Bytes += Math.Max(0, msg.Length - 2);

        if (Opts.Verbose)
            SendOK();

        var pa = ParseCtx.Pa;
        if (pa.Subject is null)
            return;

        var (acc, result) = GetAccAndResultFromCache();
        if (acc is null)
            return;

        if ((result?.PSubs.Count ?? 0) + (result?.QSubs.Count ?? 0) > 0)
            ProcessMsgResults(acc, result, msg, null, pa.Subject, pa.Reply, PmrFlags.None);
    }

    internal Exception? SendRouteConnect(string clusterName, bool tlsRequired)
    {
        var user = string.Empty;
        var pass = string.Empty;
        var routeUrl = Route?.Url;
        if (routeUrl is not null && !string.IsNullOrEmpty(routeUrl.UserInfo))
        {
            var userInfo = routeUrl.UserInfo.Split(':', 2);
            user = userInfo[0];
            if (userInfo.Length > 1)
                pass = userInfo[1];
        }

        if (Server is not NatsServer server)
            return new InvalidOperationException("route server unavailable");

        var connect = new ConnectInfo
        {
            Echo = true,
            Verbose = false,
            Pedantic = false,
            User = user,
            Pass = pass,
            Tls = tlsRequired,
            Name = server.ID(),
            Headers = server.SupportsHeaders(),
            Cluster = clusterName,
            Dynamic = server.IsClusterNameDynamic(),
            Lnoc = true,
        };

        var payload = JsonSerializer.Serialize(connect);
        EnqueueProto(Encoding.ASCII.GetBytes($"CONNECT {payload}\r\n"));
        return null;
    }

    internal void ProcessRouteInfo(ServerInfo info)
    {
        if (Server is not NatsServer server)
            return;

        lock (_mu)
        {
            Route ??= new Route();

            if (Flags.IsSet(ClientFlags.InfoReceived))
            {
                Opts.Import = info.Import;
                Opts.Export = info.Export;
            }

            Route.RemoteId = info.Id;
            Route.RemoteName = info.Name;
            Route.AuthRequired = info.AuthRequired;
            Route.TlsRequired = info.TlsRequired;
            Route.GatewayUrl = info.GatewayUrl ?? string.Empty;
            Route.Lnoc = info.Lnoc;
            Route.Lnocu = info.Lnocu;
            Route.JetStream = info.JetStream;
            Route.ConnectUrls = info.ClientConnectUrls?.ToList() ?? [];
            Route.WsConnUrls = info.WsConnectUrls?.ToList() ?? [];
            Route.LeafnodeUrl = info.LeafNodeUrls is { Length: 1 } leaf ? leaf[0] : string.Empty;
            Route.Hash = NatsServer.GetHash(info.Name);
            Route.IdHash = NatsServer.GetHash(info.Id);

            Opts.Protocol = info.Proto;
            Headers = server.SupportsHeaders() && info.Headers;
            Flags |= ClientFlags.InfoReceived;
        }

        if (NatsServer.NeedsCompression(server.GetOpts().Cluster.Compression.Mode))
            _ = server.NegotiateRouteCompression(this, Route?.DidSolicit == true, Route?.AccName is { Length: > 0 } an ? Encoding.ASCII.GetString(an) : string.Empty, info.Compression ?? string.Empty, server.GetOpts());

        server.UpdateRemoteRoutePerms(this, info);
    }

    internal bool CanImport(string subject) => PubAllowedFullCheck(subject, fullCheck: false, hasLock: true);

    internal bool CanExport(string subject) => CanSubscribe(subject);

    internal void SetRoutePermissions(RoutePermissions? perms)
    {
        if (perms is null)
        {
            Perms = null;
            MPerms = null;
            return;
        }

        SetPermissions(new Permissions
        {
            Publish = perms.Import?.Clone(),
            Subscribe = perms.Export?.Clone(),
        });
    }

    internal (bool IsPinnedAccountRoute, string AccountName, bool KeyHasSubType) GetRoutedSubKeyInfo()
    {
        var accountName = Route?.AccName is { Length: > 0 } an
            ? Encoding.ASCII.GetString(an)
            : string.Empty;
        return (!string.IsNullOrEmpty(accountName), accountName, Route?.Lnocu == true);
    }

    internal void RemoveRemoteSubs()
    {
        if (Server is not NatsServer server)
            return;

        Dictionary<string, Subscription> subs;
        var grouped = new Dictionary<string, List<Subscription>>(StringComparer.Ordinal);
        var (pinned, accountName, keyHasSubType) = GetRoutedSubKeyInfo();

        lock (_mu)
        {
            subs = Subs;
            Subs = new Dictionary<string, Subscription>(StringComparer.Ordinal);
        }

        foreach (var kvp in subs)
        {
            var keyAccount = pinned
                ? accountName
                : RouteHandler.GetAccNameFromRoutedSubKey(kvp.Value, kvp.Key, keyHasSubType);
            if (string.IsNullOrEmpty(keyAccount))
                continue;

            if (!grouped.TryGetValue(keyAccount, out var list))
            {
                list = [];
                grouped[keyAccount] = list;
            }
            list.Add(kvp.Value);
        }

        foreach (var (accName, list) in grouped)
        {
            var (acc, _) = server.LookupAccount(accName);
            acc?.Sublist?.RemoveBatch(list);
        }
    }

    internal List<Subscription> RemoveRemoteSubsForAcc(string name)
    {
        var removed = new List<Subscription>();
        var (_, _, keyHasSubType) = GetRoutedSubKeyInfo();
        lock (_mu)
        {
            foreach (var key in Subs.Keys.ToArray())
            {
                var sub = Subs[key];
                if (RouteHandler.GetAccNameFromRoutedSubKey(sub, key, keyHasSubType) != name)
                    continue;
                removed.Add(sub);
                Subs.Remove(key);
            }
        }
        return removed;
    }

    internal (byte[] Origin, string AccountName, byte[] Subject, byte[] Queue, Exception? Error)
        ParseUnsubProto(byte[] arg, bool accInProto, bool hasOrigin)
    {
        _in.Subs++;

        var args = SplitArg(arg);
        var origin = Array.Empty<byte>();
        var queue = Array.Empty<byte>();
        var subjectIndex = 0;

        if (hasOrigin)
        {
            if (args.Count == 0)
                return (origin, string.Empty, Array.Empty<byte>(), queue, new FormatException($"parse error: '{Encoding.ASCII.GetString(arg)}'"));
            origin = args[0];
            subjectIndex = 1;
        }
        if (accInProto)
            subjectIndex++;

        if (args.Count is not (>= 1) || args.Count < subjectIndex + 1 || args.Count > subjectIndex + 2)
            return (origin, string.Empty, Array.Empty<byte>(), queue, new FormatException($"parse error: '{Encoding.ASCII.GetString(arg)}'"));

        if (args.Count == subjectIndex + 2)
            queue = args[subjectIndex + 1];

        var accountName = accInProto ? Encoding.ASCII.GetString(args[subjectIndex - 1]) : string.Empty;
        return (origin, accountName, args[subjectIndex], queue, null);
    }

    internal Exception? ProcessRemoteUnsub(byte[] arg, bool leafUnsub)
    {
        if (Server is not NatsServer server)
            return null;

        string accountName;
        var accInProto = true;
        bool originSupport;

        lock (_mu)
        {
            originSupport = Route?.Lnocu == true;
            if (Route?.AccName is { Length: > 0 } an)
            {
                accountName = Encoding.ASCII.GetString(an);
                accInProto = false;
            }
            else
            {
                accountName = string.Empty;
            }
        }

        var (_, protoAccName, subject, _, err) = ParseUnsubProto(arg, accInProto, leafUnsub && originSupport);
        if (err is not null)
            return new FormatException($"processRemoteUnsub {err.Message}");

        if (accInProto)
            accountName = protoAccName;

        var (acc, _) = server.LookupAccount(accountName);
        if (acc is null)
        {
            Debugf("Unknown account {0} for subject {1}", accountName, Encoding.ASCII.GetString(subject));
            return null;
        }

        Subscription? sub = null;
        var key = Encoding.ASCII.GetString(arg);
        lock (_mu)
        {
            if (IsClosed())
                return null;
            if (Subs.TryGetValue(key, out sub))
            {
                Subs.Remove(key);
                acc.Sublist?.Remove(sub);
            }
        }

        if (Opts.Verbose)
            SendOK();

        return null;
    }

    internal Exception? ProcessRemoteSub(byte[] protoArg, bool hasOrigin)
    {
        _in.Subs++;
        if (Server is not NatsServer server)
            return null;

        var args = SplitArg(protoArg);
        var (isPinned, accountName, _) = GetRoutedSubKeyInfo();
        var accInProto = !isPinned;
        var subjectIndex = 0;

        if (hasOrigin)
            subjectIndex++;
        if (accInProto)
            subjectIndex++;

        if (args.Count is not (>= 1) || (args.Count != subjectIndex + 1 && args.Count != subjectIndex + 3))
            return new FormatException($"processRemoteSub Parse Error: '{Encoding.ASCII.GetString(protoArg)}'");

        if (accInProto)
            accountName = Encoding.ASCII.GetString(args[subjectIndex - 1]);
        var subject = args[subjectIndex];
        byte[]? queue = null;
        var qw = 1;
        if (args.Count == subjectIndex + 3)
        {
            queue = args[subjectIndex + 1];
            _ = int.TryParse(Encoding.ASCII.GetString(args[subjectIndex + 2]), out qw);
            if (qw <= 0)
                qw = 1;
        }

        var (acc, _) = server.LookupOrRegisterAccount(accountName);
        if (acc is null)
            return null;

        lock (_mu)
        {
            if (IsClosed())
                return null;
            if (Perms is not null && !CanExport(Encoding.ASCII.GetString(subject)))
                return null;
            if (SubsAtLimit())
            {
                MaxSubsExceeded();
                return null;
            }

            var key = Encoding.ASCII.GetString(protoArg);
            if (!Subs.ContainsKey(key))
            {
                var sub = new Subscription
                {
                    Subject = subject,
                    Queue = queue,
                    Sid = Encoding.ASCII.GetBytes(key),
                    Qw = qw,
                };
                Subs[key] = sub;
                acc.Sublist?.Insert(sub);
            }
        }

        if (Opts.Verbose)
            SendOK();
        return null;
    }

    internal byte[] AddRouteSubOrUnsubProtoToBuf(byte[] buf, string accName, Subscription sub, bool isSubProto)
    {
        var list = new List<byte>(buf.Length + 64);
        list.AddRange(buf);

        if (isSubProto)
            list.AddRange(Encoding.ASCII.GetBytes("RS+ "));
        else
            list.AddRange(Encoding.ASCII.GetBytes("RS- "));

        if (Route?.AccName is not { Length: > 0 })
        {
            list.AddRange(Encoding.ASCII.GetBytes(accName));
            list.Add((byte)' ');
        }

        list.AddRange(sub.Subject);
        if (sub.Queue is { Length: > 0 } queue)
        {
            list.Add((byte)' ');
            list.AddRange(queue);
            if (isSubProto)
            {
                list.Add((byte)' ');
                list.AddRange(Encoding.ASCII.GetBytes(Math.Max(sub.Qw, 1).ToString()));
            }
        }

        list.Add((byte)'\r');
        list.Add((byte)'\n');
        return list.ToArray();
    }

    internal void SendRouteSubProtos(IReadOnlyList<Subscription> subs, bool trace, Func<Subscription, bool>? filter = null) =>
        SendRouteSubOrUnSubProtos(subs, isSubProto: true, trace, filter);

    internal void SendRouteUnSubProtos(IReadOnlyList<Subscription> subs, bool trace, Func<Subscription, bool>? filter = null) =>
        SendRouteSubOrUnSubProtos(subs, isSubProto: false, trace, filter);

    internal void SendRouteSubOrUnSubProtos(
        IReadOnlyList<Subscription> subs,
        bool isSubProto,
        bool trace,
        Func<Subscription, bool>? filter = null)
    {
        var buf = Array.Empty<byte>();
        foreach (var sub in subs)
        {
            if (filter is not null && !filter(sub))
                continue;

            var accountName = ServerConstants.DefaultGlobalAccount;

            var startLen = buf.Length;
            buf = AddRouteSubOrUnsubProtoToBuf(buf, accountName, sub, isSubProto);
            if (trace && buf.Length > startLen)
                TraceOutOp(string.Empty, buf.AsSpan(startLen, buf.Length - startLen - 2).ToArray());
        }

        if (buf.Length > 0)
            EnqueueProto(buf);
    }

    internal bool ImportFilter(string subject) => CanImport(subject);

    internal bool IsSolicitedRoute() => Route?.DidSolicit == true;

    internal Exception? ProcessRouteConnect(byte[] arg)
    {
        if (arg is not { Length: > 0 })
            return new FormatException("processRouteConnect parse error");

        ConnectInfo? info;
        try
        {
            info = JsonSerializer.Deserialize<ConnectInfo>(arg);
        }
        catch (Exception ex)
        {
            return ex;
        }

        if (info is null)
            return new FormatException("processRouteConnect missing CONNECT payload");

        lock (_mu)
        {
            Opts.Name = info.Name;
            Opts.Headers = info.Headers;
            Route ??= new Route();
            Route.Lnoc = info.Lnoc;
            Route.Lnocu = info.Lnocu;
        }

        return null;
    }
}