natsnet/dotnet/src/ZB.MOM.NatsNet.Server/Internal/SignalHandler.cs

// Copyright 2012-2025 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Adapted from server/signal.go and server/service.go in the NATS server Go source.

using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Text;

namespace ZB.MOM.NatsNet.Server.Internal;

/// <summary>
/// Maps <see cref="ServerCommand"/> to OS signal-like behavior.
/// Mirrors <c>CommandToSignal</c> and <c>ProcessSignal</c> from signal.go.
/// In .NET, signal sending is replaced by process-level signaling on Unix.
/// </summary>
public static class SignalHandler
{
    private const string ResolvePidError = "unable to resolve pid, try providing one";
    private static string _processName = "nats-server";
    internal static Func<List<int>> ResolvePidsHandler { get; set; } = ResolvePids;
    internal static Func<int, UnixSignal, Exception?> SendSignalHandler { get; set; } = SendSignal;

    internal static void ResetTestHooks()
    {
        ResolvePidsHandler = ResolvePids;
        SendSignalHandler = SendSignal;
    }

    /// <summary>
    /// Sets the process name used for resolving PIDs.
    /// Mirrors <c>SetProcessName</c> in signal.go.
    /// </summary>
    public static void SetProcessName(string name) => _processName = name;

    /// <summary>
    /// Sends a signal command to a running NATS server process.
    /// On Unix, maps commands to kill signals.
    /// On Windows, this is a no-op (service manager handles signals).
    /// Mirrors <c>ProcessSignal</c> in signal.go.
    /// </summary>
    public static Exception? ProcessSignal(ServerCommand command, string pidExpr = "")
    {
        if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
            return new PlatformNotSupportedException("Signal processing not supported on Windows; use service manager.");

        try
        {
            var pids = new List<int>(1);
            var pidStr = pidExpr.TrimEnd('*');
            var isGlob = pidExpr.EndsWith('*');

            if (!string.IsNullOrEmpty(pidStr))
            {
                if (!int.TryParse(pidStr, out var pid))
                    return new InvalidOperationException($"invalid pid: {pidStr}");
                pids.Add(pid);
            }

            if (string.IsNullOrEmpty(pidStr) || isGlob)
                pids = ResolvePidsHandler();

            if (pids.Count > 1 && !isGlob)
            {
                var sb = new StringBuilder($"multiple {_processName} processes running:");
                foreach (var p in pids)
                    sb.Append('\n').Append(p);
                return new InvalidOperationException(sb.ToString());
            }

            if (pids.Count == 0)
                return new InvalidOperationException($"no {_processName} processes running");

            UnixSignal signal;
            try
            {
                signal = CommandToUnixSignal(command);
            }
            catch (Exception ex)
            {
                return ex;
            }

            var errBuilder = new StringBuilder();
            foreach (var pid in pids)
            {
                var pidText = pid.ToString();
                if (pidStr.Length > 0 && pidText != pidStr)
                {
                    if (!isGlob || !pidText.StartsWith(pidStr, StringComparison.Ordinal))
                        continue;
                }

                var err = SendSignalHandler(pid, signal);
                if (err != null)
                {
                    errBuilder
                        .Append('\n')
                        .Append("signal \"")
                        .Append(CommandToString(command))
                        .Append("\" ")
                        .Append(pid)
                        .Append(": ")
                        .Append(err.Message);
                }
            }

            if (errBuilder.Length > 0)
                return new InvalidOperationException(errBuilder.ToString());

            return null;
        }
        catch (Exception ex)
        {
            return ex;
        }
    }

    /// <summary>
    /// Resolves PIDs of running nats-server processes via pgrep.
    /// Mirrors <c>resolvePids</c> in signal.go.
    /// </summary>
    public static List<int> ResolvePids()
    {
        var pids = new List<int>(8);
        try
        {
            var psi = new ProcessStartInfo("pgrep", _processName)
            {
                RedirectStandardOutput = true,
                UseShellExecute = false,
                CreateNoWindow = true,
            };
            using var proc = Process.Start(psi);
            if (proc == null)
                throw new InvalidOperationException(ResolvePidError);

            var output = proc.StandardOutput.ReadToEnd();
            proc.WaitForExit();
            if (proc.ExitCode != 0)
                return pids;

            var currentPid = Environment.ProcessId;
            foreach (var line in output.Split('\n', StringSplitOptions.RemoveEmptyEntries))
            {
                if (!int.TryParse(line.Trim(), out var pid))
                    throw new InvalidOperationException(ResolvePidError);

                if (pid != currentPid)
                    pids.Add(pid);
            }
        }
        catch (InvalidOperationException ex) when (ex.Message == ResolvePidError)
        {
            throw;
        }
        catch
        {
            throw new InvalidOperationException(ResolvePidError);
        }

        return pids;
    }

    /// <summary>
    /// Maps a server command to Unix signal.
    /// Mirrors <c>CommandToSignal</c> in signal.go.
    /// </summary>
    public static UnixSignal CommandToUnixSignal(ServerCommand command) => command switch
    {
        ServerCommand.Stop => UnixSignal.SigKill,
        ServerCommand.Quit => UnixSignal.SigInt,
        ServerCommand.Reopen => UnixSignal.SigUsr1,
        ServerCommand.Reload => UnixSignal.SigHup,
        ServerCommand.LameDuckMode => UnixSignal.SigUsr2,
        ServerCommand.Term => UnixSignal.SigTerm,
        _ => throw new ArgumentOutOfRangeException(nameof(command), $"unknown signal \"{CommandToString(command)}\""),
    };

    /// <summary>
    /// Go parity alias for <see cref="CommandToUnixSignal"/>.
    /// Mirrors <c>CommandToSignal</c> in signal.go.
    /// </summary>
    public static UnixSignal CommandToSignal(ServerCommand command) => CommandToUnixSignal(command);

    private static Exception? SendSignal(int pid, UnixSignal signal)
    {
        try
        {
            Process.GetProcessById(pid).Kill(signal == UnixSignal.SigKill);
            return null;
        }
        catch (Exception ex)
        {
            return ex;
        }
    }

    private static string CommandToString(ServerCommand command) => command switch
    {
        ServerCommand.Stop => "stop",
        ServerCommand.Quit => "quit",
        ServerCommand.Reopen => "reopen",
        ServerCommand.Reload => "reload",
        ServerCommand.LameDuckMode => "ldm",
        ServerCommand.Term => "term",
        _ => command.ToString().ToLowerInvariant(),
    };

    /// <summary>
    /// Runs the server (non-Windows). Mirrors <c>Run</c> in service.go.
    /// </summary>
    public static void Run(Action startServer)
    {
        var error = ServiceManager.Run(startServer);
        if (error is not null)
            throw error;
    }

    /// <summary>
    /// Returns false on non-Windows. Mirrors <c>isWindowsService</c>.
    /// </summary>
    public static bool IsWindowsService() => ServiceManager.IsWindowsService();
}

/// <summary>Unix signal codes for NATS command mapping.</summary>
public enum UnixSignal
{
    SigInt = 2,
    SigKill = 9,
    SigUsr1 = 10,
    SigHup = 1,
    SigUsr2 = 12,
    SigTerm = 15,
}