natsnet/dotnet/src/ZB.MOM.NatsNet.Server/ClientConnection.Gateways.Messages.cs

// Copyright 2018-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0

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

namespace ZB.MOM.NatsNet.Server;

public sealed partial class ClientConnection
{
    internal void SendMsgToGateways(Account? account, byte[] subject, byte[]? reply, byte[] payload)
    {
        if (Server is not NatsServer server || account == null || subject.Length == 0)
            return;

        var outboundGateways = server.GetOutboundGatewayConnections();
        if (outboundGateways.Count == 0)
            return;

        foreach (var gateway in outboundGateways)
        {
            if (!GatewayInterest(account, Encoding.ASCII.GetString(subject)))
                continue;

            var replyToSend = reply;
            if (reply is { Length: > 0 } && gateway.Gateway != null)
            {
                var shouldMap = server.ShouldMapReplyForGatewaySend(reply, gateway.Gateway.UseOldPrefix);
                if (shouldMap)
                    replyToSend = reply;
            }

            var proto = MsgHeader(subject, replyToSend, new Internal.Subscription { Sid = Encoding.ASCII.GetBytes("1") });
            gateway.EnqueueProto(proto);
            gateway.EnqueueProto(payload);
            gateway.EnqueueProto(Encoding.ASCII.GetBytes("\r\n"));
        }
    }

    internal void SendAccountUnsubToGateway(byte[] accountName)
    {
        if (accountName.Length == 0)
            return;

        lock (_mu)
        {
            Gateway ??= new Gateway();
            Gateway.InSim ??= new Dictionary<string, InSide>(StringComparer.Ordinal);

            var key = Encoding.ASCII.GetString(accountName);
            if (!Gateway.InSim.TryGetValue(key, out var entry) || entry != null)
            {
                Gateway.InSim[key] = null!;
                var proto = Encoding.ASCII.GetBytes($"A- {key}\r\n");
                EnqueueProto(proto);
                if (Trace)
                    TraceOutOp(string.Empty, proto.AsSpan(0, proto.Length - 2).ToArray());
            }
        }
    }

    internal bool HandleGatewayReply(byte[] msg)
    {
        if (Server is not NatsServer server || ParseCtx.Pa.Subject is not { Length: > 0 } originalSubject)
            return false;

        var (isRoutedReply, isOldPrefix) = GatewayHandler.IsGWRoutedSubjectAndIsOldPrefix(originalSubject);
        if (!isRoutedReply)
            return false;

        ParseCtx.Pa.Subject = GatewayHandler.GetSubjectFromGWRoutedReply(originalSubject, isOldPrefix);
        ParseCtx.Pa.PaCache = [.. ParseCtx.Pa.Account ?? [], (byte)' ', .. ParseCtx.Pa.Subject];

        var (account, result) = GetAccAndResultFromCache();
        if (account is not Account concreteAccount)
        {
            Debugf("Unknown account {0} for gateway message on subject: {1}",
                ParseCtx.Pa.Account is { Length: > 0 } accountName ? Encoding.ASCII.GetString(accountName) : string.Empty,
                Encoding.ASCII.GetString(ParseCtx.Pa.Subject));

            if (ParseCtx.Pa.Account is { Length: > 0 } gatewayAccount)
                server.GatewayHandleAccountNoInterest(this, gatewayAccount);
            return true;
        }

        if (result != null && (result.PSubs.Count + result.QSubs.Count) > 0)
            ProcessMsgResults(concreteAccount, result, msg, null, ParseCtx.Pa.Subject, ParseCtx.Pa.Reply, PmrFlags.None);

        if (!IsServiceReply(ParseCtx.Pa.Subject))
            SendMsgToGateways(concreteAccount, ParseCtx.Pa.Subject, ParseCtx.Pa.Reply, msg);

        return true;
    }

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

        if (Opts.Verbose)
            SendOK();

        if (Server is not NatsServer server || ParseCtx.Pa.Subject is not { Length: > 0 })
            return;

        if (HandleGatewayReply(msg))
            return;

        var (account, result) = GetAccAndResultFromCache();
        if (account is not Account concreteAccount)
        {
            Debugf("Unknown account {0} for gateway message on subject: {1}",
                ParseCtx.Pa.Account is { Length: > 0 } accountName ? Encoding.ASCII.GetString(accountName) : string.Empty,
                Encoding.ASCII.GetString(ParseCtx.Pa.Subject));

            if (ParseCtx.Pa.Account is { Length: > 0 } gatewayAccount)
                server.GatewayHandleAccountNoInterest(this, gatewayAccount);
            return;
        }

        var noInterest = result == null || result.PSubs.Count == 0;
        if (noInterest)
        {
            server.GatewayHandleSubjectNoInterest(this, concreteAccount, ParseCtx.Pa.Account ?? Encoding.ASCII.GetBytes(concreteAccount.Name), ParseCtx.Pa.Subject);
            if (ParseCtx.Pa.Queues is null || ParseCtx.Pa.Queues.Count == 0)
                return;
        }

        ProcessMsgResults(concreteAccount, result, msg, null, ParseCtx.Pa.Subject, ParseCtx.Pa.Reply, PmrFlags.None);
    }

    internal void GatewayAllSubsReceiveStart(ServerInfo info)
    {
        var account = GatewayHandler.GetAccountFromGatewayCommand(this, info, "start");
        if (string.IsNullOrWhiteSpace(account))
            return;

        Debugf("Gateway {0}: switching account {1} to {2} mode",
            info.Gateway ?? string.Empty, account, GatewayInterestMode.InterestOnly.String());

        Gateway ??= new Gateway();
        Gateway.OutSim ??= new System.Collections.Concurrent.ConcurrentDictionary<string, OutSide>(StringComparer.Ordinal);

        var outSide = Gateway.OutSim.GetOrAdd(account, _ => new OutSide
        {
            Sl = Internal.DataStructures.SubscriptionIndex.NewSublistWithCache(),
        });

        outSide.AcquireWriteLock();
        try { outSide.Mode = GatewayInterestMode.Transitioning; }
        finally { outSide.ReleaseWriteLock(); }
    }

    internal void GatewayAllSubsReceiveComplete(ServerInfo info)
    {
        var account = GatewayHandler.GetAccountFromGatewayCommand(this, info, "complete");
        if (string.IsNullOrWhiteSpace(account))
            return;

        if (Gateway?.OutSim == null || !Gateway.OutSim.TryGetValue(account, out var outSide))
            return;

        outSide.AcquireWriteLock();
        try
        {
            outSide.Ni = null;
            outSide.Mode = GatewayInterestMode.InterestOnly;
        }
        finally
        {
            outSide.ReleaseWriteLock();
        }

        Debugf("Gateway {0}: switching account {1} to {2} mode complete",
            info.Gateway ?? string.Empty, account, GatewayInterestMode.InterestOnly.String());
    }

    internal void GatewaySwitchAccountToSendAllSubs(InSide inSide, string accountName)
    {
        if (Server is not NatsServer server || string.IsNullOrWhiteSpace(accountName))
            return;

        inSide.Ni = null;
        inSide.Mode = GatewayInterestMode.Transitioning;

        var remoteGatewayName = Gateway?.Name ?? string.Empty;
        Debugf("Gateway {0}: switching account {1} to {2} mode",
            remoteGatewayName, accountName, GatewayInterestMode.InterestOnly.String());

        void SendCommand(byte command, bool withLock)
        {
            var info = new ServerInfo
            {
                Gateway = server.GetGatewayName(),
                GatewayCmd = command,
                GatewayCmdPayload = Encoding.ASCII.GetBytes(accountName),
            };

            var infoProto = NatsServer.GenerateInfoJson(info);
            if (withLock)
            {
                lock (_mu)
                {
                    EnqueueProto(infoProto);
                }
                return;
            }

            EnqueueProto(infoProto);
        }

        SendCommand(GatewayHandler.GatewayCmdAllSubsStart, withLock: false);

        _ = server.StartGoRoutine(() =>
        {
            server.SendAccountSubsToGateway(this, accountName);
            SendCommand(GatewayHandler.GatewayCmdAllSubsComplete, withLock: true);

            Debugf("Gateway {0}: switching account {1} to {2} mode complete",
                remoteGatewayName, accountName, GatewayInterestMode.InterestOnly.String());
        });
    }
}