ModernUO/Projects/Server.Tests/Tests/Network/NetStateSendBufferTests.cs

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Net.Sockets;
using System.Network;
using Server.Network;
using Xunit;

namespace Server.Tests.Network;

[Collection("Sequential Server Tests")]
public class NetStateSendBufferTests
{
    // Authentication is the 0x91 credentials check on the game server; that is where buffers promote
    private static NetState CreateAuthenticatedNetState(out Socket client)
    {
        var ns = PacketTestUtilities.CreateTestNetState(out client);
        ns._protocolState = NetState.ProtocolState.GameServer_AwaitingGameServerLogin;
        ns.Account = new MockAccount();
        return ns;
    }

    [SkippableFact]
    public void Account_InTheGameServerState_PromotesBothBuffers()
    {
        NetTestHelpers.RequirePreAuthBuffers();
        using var ping = NetTestHelpers.RegisterPing();
        var ns = PacketTestUtilities.CreateTestNetState(out var client);
        try
        {
            Assert.Equal(NetState.InitialSendBufferSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.Equal(NetState.InitialRecvBufferSize, ns._socket.RecvBuffer.PhysicalSize);

            ns._protocolState = NetState.ProtocolState.GameServer_AwaitingGameServerLogin;
            ns.Account = new MockAccount();

            // Send promotes at once; nothing was in flight so no buffer retires
            Assert.Equal(NetState.SendBufferSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.False(ns._sendBufferGrown);

            // Recv promotes at the next completion
            ns._protocolState = NetState.ProtocolState.GameServer_LoggedIn;
            client.Send(new byte[] { NetTestHelpers.PingPacketId, 0x01 });
            NetTestHelpers.SliceUntil(() => ns._socket.RecvBuffer.PhysicalSize == NetState.RecvBufferSize);
            Assert.True(ns.Running);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [SkippableFact]
    public void Account_OnTheLoginServer_KeepsTheInitialBuffers()
    {
        NetTestHelpers.RequirePreAuthBuffers();
        var ns = PacketTestUtilities.CreateTestNetState(out var client);
        try
        {
            ns._protocolState = NetState.ProtocolState.LoginServer_AwaitingLogin;
            ns.Account = new MockAccount();

            Assert.Equal(NetState.InitialSendBufferSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.Equal(NetState.InitialRecvBufferSize, ns._socket.RecvBuffer.PhysicalSize);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [SkippableFact]
    public void Send_BeyondTheInitialBuffer_BeforeCredentials_IsExhausted()
    {
        // Nothing promotes before credentials verify: the 4 KiB ring is the whole pre-auth budget
        NetTestHelpers.RequirePreAuthBuffers();
        var ns = PacketTestUtilities.CreateTestNetState(out var client);
        try
        {
            var data = NetTestHelpers.Pattern(NetState.InitialSendBufferSize + 512, 7);
            ns.Send(data);

            Assert.Contains("exhausted", ns._disconnectReason, StringComparison.OrdinalIgnoreCase);
            Assert.Equal(NetState.InitialSendBufferSize, ns._socket.SendBuffer.PhysicalSize);

            // A graceful server-side close half-closes and waits for the peer's own FIN; closing the
            // peer here stands in for the client eventually going away, so Running can be observed.
            // Once it completes the NetState is disposed and _socket goes null, so nothing below
            // this point may touch it.
            client.Close();
            NetTestHelpers.SliceUntil(() => !ns.Running);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [SkippableFact]
    public void Send_BeyondTheInitialBuffer_WithCredentials_PromotesOnDemand()
    {
        // The late-verdict path from Account_AssignedAfterTheStateFlipped_StillPromotes: an account
        // attached while the setter's own gate does not fire must still let the send path promote
        NetTestHelpers.RequirePreAuthBuffers();
        var ns = PacketTestUtilities.CreateTestNetState(out var client);
        try
        {
            ns.Account = new MockAccount(); // _protocolState is still AwaitingSeed: the setter does not promote
            var data = NetTestHelpers.Pattern(NetState.InitialSendBufferSize + 512, 7);
            ns.Send(data);

            Assert.True(ns.Running);
            Assert.Equal(string.Empty, ns._disconnectReason);
            Assert.Equal(NetState.SendBufferSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.False(ns._sendBufferGrown); // promotion is not growth: nothing to shrink later
            Assert.Equal(data, NetTestHelpers.ReadAll(client, data.Length));
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [SkippableFact]
    public void Send_BacklogBeforeCredentials_IsExhaustedInsteadOfPromoted()
    {
        NetTestHelpers.RequirePreAuthBuffers();
        var ns = PacketTestUtilities.CreateTestNetState(out var client);
        try
        {
            // No Slice() between sends: nothing drains, so a second write finds a non-empty initial buffer
            var half = NetTestHelpers.Pattern(NetState.InitialSendBufferSize / 2, 11);
            ns.Send(half);
            ns.Send(half);
            ns.Send(half); // cannot fit, buffer not empty, no account: refused

            Assert.Equal(NetState.InitialSendBufferSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.Contains("exhausted", ns._disconnectReason, StringComparison.OrdinalIgnoreCase);

            // A graceful server-side close waits for the queued bytes to drain and then for the
            // peer's own FIN; closing the peer here stands in for the client eventually going away.
            // Once it completes the NetState is disposed and _socket goes null, so nothing below
            // this point may touch it.
            client.Close();
            NetTestHelpers.SliceUntil(() => !ns.Running);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [SkippableFact]
    public void Account_AssignedAfterTheStateFlipped_StillPromotes()
    {
        NetTestHelpers.RequirePreAuthBuffers();
        var ns = PacketTestUtilities.CreateTestNetState(out var client);
        try
        {
            ns._protocolState = NetState.ProtocolState.GameServer_LoggedIn;
            ns.Account = new MockAccount();
            Assert.Equal(NetState.SendBufferSize, ns._socket.SendBuffer.PhysicalSize);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    // A random prefix lands under the target; zeros (2 bits each) walk it up a byte at a time
    private static byte[] CompressesToExactly(int compressedLength, int seed)
    {
        var randomLength = compressedLength * 5 / 8;
        var data = new byte[compressedLength * 8];
        NetTestHelpers.Pattern(randomLength, seed).CopyTo(data, 0);

        var scratch = new byte[compressedLength * 2 + 16];
        for (var length = randomLength; length < data.Length; length++)
        {
            var compressed = NetworkCompression.Compress(data.AsSpan(0, length), scratch);
            Assert.InRange(compressed, 1, compressedLength);

            if (compressed == compressedLength)
            {
                return data[..length];
            }
        }

        throw new InvalidOperationException($"No chunk compresses to exactly {compressedLength} bytes");
    }

    // Peer receipt does not free the buffer; the completion must land first
    private static void WaitForDrain(NetState ns) =>
        NetTestHelpers.SliceUntil(
            () => ns._socket.SendBuffer.ReadableBytes == 0 && ns._socket.SendBuffer.InFlightBytes == 0,
            10000
        );

    [Fact]
    public void Send_GrowsInsteadOfDisconnecting_AndStreamStaysIntact()
    {
        var ns = CreateAuthenticatedNetState(out var client);
        var baseSize = ns._socket.SendBuffer.PhysicalSize;

        try
        {
            // No Slice(): nothing drains, so the base buffer overflows
            var chunk = baseSize / 4;
            var expected = new byte[chunk * 6];
            for (var i = 0; i < 6; i++)
            {
                var data = NetTestHelpers.Pattern(chunk, i);
                data.CopyTo(expected, i * chunk);
                ns.Send(data);
            }

            Assert.True(ns.Running);
            Assert.Equal(string.Empty, ns._disconnectReason);
            Assert.True(ns._socket.SendBuffer.PhysicalSize > baseSize);
            Assert.True(ns._sendBufferGrown);

            Assert.Equal(expected, NetTestHelpers.ReadAll(client, expected.Length));
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [Fact]
    public void Send_CompressedGrowth_RetriesAndStreamStaysIntact()
    {
        var ns = CreateAuthenticatedNetState(out var client);
        ns.CompressionEnabled = true;
        var baseSize = ns._socket.SendBuffer.PhysicalSize;

        try
        {
            // No Slice(): compressed output accumulates until the base overflows
            var chunkLength = baseSize / 4;
            var expected = new byte[(chunkLength * 2 + 4) * 6];
            var scratch = new byte[chunkLength * 2 + 4];
            var expectedLength = 0;
            for (var i = 0; i < 6; i++)
            {
                var data = NetTestHelpers.Pattern(chunkLength, i + 100);
                var compressedLength = NetworkCompression.Compress(data, scratch);
                Assert.True(compressedLength > 0);
                scratch.AsSpan(0, compressedLength).CopyTo(expected.AsSpan(expectedLength));
                expectedLength += compressedLength;
                ns.Send(data);
            }

            Assert.True(ns.Running);
            Assert.Equal(string.Empty, ns._disconnectReason);
            Assert.True(ns._socket.SendBuffer.PhysicalSize > baseSize);
            Assert.True(ns._sendBufferGrown);

            Array.Resize(ref expected, expectedLength);
            Assert.Equal(expected, NetTestHelpers.ReadAll(client, expected.Length));
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [Fact]
    public void Send_CompressedGrowth_FullBuffer_GrowsOneTierAndFits()
    {
        var ns = CreateAuthenticatedNetState(out var client);
        ns.CompressionEnabled = true;
        var baseSize = ns._socket.SendBuffer.PhysicalSize;
        var expected = new List<byte>(baseSize * 2);

        void SendAndRecord(byte[] data)
        {
            var scratch = new byte[data.Length * 2 + 4];
            var compressedLength = NetworkCompression.Compress(data, scratch);
            Assert.True(compressedLength > 0);
            expected.AddRange(scratch.AsSpan(0, compressedLength));
            ns.Send(data);
        }

        try
        {
            // Chunks sized so the fill cannot overflow (random bytes never compress; ratio at most 11/8)
            var seed = 600;
            while (ns._socket.SendBuffer.WritableBytes > 4096)
            {
                var writable = ns._socket.SendBuffer.WritableBytes;
                SendAndRecord(NetTestHelpers.Pattern(Math.Min(baseSize / 8, (writable - 2048) * 8 / 11), seed++));
            }

            // Full; no span to compress into
            SendAndRecord(CompressesToExactly(ns._socket.SendBuffer.WritableBytes, seed));
            Assert.Equal(0, ns._socket.SendBuffer.WritableBytes);
            Assert.Equal(baseSize, ns._socket.SendBuffer.PhysicalSize);

            // Raw exceeds the remainder; compressed fits in one tier
            SendAndRecord(new byte[Math.Min(baseSize * 3 / 4, NetworkCompression.DefiniteOverflow)]);

            Assert.True(ns.Running);
            Assert.Equal(string.Empty, ns._disconnectReason);
            Assert.Equal(baseSize * 2, ns._socket.SendBuffer.PhysicalSize);
            Assert.True(ns._sendBufferGrown);

            Assert.Equal(expected.ToArray(), NetTestHelpers.ReadAll(client, expected.Count));
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [Fact]
    public void Send_WithSendInFlight_GrowsAndStreamStaysIntact()
    {
        var ns = CreateAuthenticatedNetState(out var client);
        var baseSize = ns._socket.SendBuffer.PhysicalSize;

        try
        {
            // A small peer buffer keeps the first send in flight through the growth
            client.ReceiveBufferSize = 4096;

            var first = NetTestHelpers.Pattern(baseSize / 2, 300);
            var second = NetTestHelpers.Pattern(baseSize / 2 + 1, 301);

            ns.Send(first);
            NetState.Slice(); // posts the first send
            Assert.True(ns._socket.SendBuffer.InFlightBytes > 0);

            ns.Send(second);

            Assert.True(ns.Running);
            Assert.Equal(string.Empty, ns._disconnectReason);
            Assert.True(ns._socket.SendBuffer.PhysicalSize > baseSize);
            Assert.True(ns._sendBufferGrown);

            var expected = new byte[first.Length + second.Length];
            first.CopyTo(expected, 0);
            second.CopyTo(expected, first.Length);

            Assert.Equal(expected, NetTestHelpers.ReadAll(client, expected.Length));
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [Fact]
    public void Shrink_AfterDrainAndHold_ReturnsToBase()
    {
        var ns = CreateAuthenticatedNetState(out var client);
        var baseSize = ns._socket.SendBuffer.PhysicalSize;

        try
        {
            var chunk = baseSize / 4;
            for (var i = 0; i < 6; i++)
            {
                ns.Send(NetTestHelpers.Pattern(chunk, i));
            }

            Assert.True(ns._sendBufferGrown);
            NetTestHelpers.ReadAll(client, chunk * 6);
            WaitForDrain(ns);

            var grewAt = ns._sendBufferGrewAt;
            Assert.False(ns.TryShrinkSendBuffer(grewAt + NetState.SendBufferHoldMs - 1));
            Assert.True(ns.TryShrinkSendBuffer(grewAt + NetState.SendBufferHoldMs));
            Assert.Equal(baseSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.False(ns._sendBufferGrown);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [Fact]
    public void Shrink_ThroughAliveSweep_ReturnsToBase()
    {
        var ns = CreateAuthenticatedNetState(out var client);
        var baseSize = ns._socket.SendBuffer.PhysicalSize;
        var previousTicks = Core._tickCount;

        try
        {
            var chunk = baseSize / 4;
            for (var i = 0; i < 6; i++)
            {
                ns.Send(NetTestHelpers.Pattern(chunk, i + 400));
            }

            Assert.True(ns._sendBufferGrown);
            NetTestHelpers.ReadAll(client, chunk * 6);
            WaitForDrain(ns);

            // The sweep shrinks; nothing here calls TryShrinkSendBuffer
            Core._tickCount = ns._sendBufferGrewAt + NetState.SendBufferHoldMs;
            NetState.CheckAllAlive();

            Assert.Equal(baseSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.False(ns._sendBufferGrown);
        }
        finally
        {
            Core._tickCount = previousTicks;
            ns.Dispose();
            client.Close();
        }
    }

    [Fact]
    public void Send_PastTheMaximum_FallsBackToExhaustion()
    {
        var ns = CreateAuthenticatedNetState(out var client);

        try
        {
            var chunk = 64 * 1024;
            var writes = NetState.MaxSendBufferSize / chunk + 1;
            for (var i = 0; i < writes; i++)
            {
                ns.Send(NetTestHelpers.Pattern(chunk, i));
            }

            Assert.Contains("Send buffer exhausted", ns._disconnectReason);
            Assert.Equal(NetState.MaxSendBufferSize, ns._socket.SendBuffer.PhysicalSize);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [Fact]
    public void Send_AboveMemoryCeiling_DoesNotGrow()
    {
        var previous = NetState._availableMemoryBytes;
        NetState ns = null;
        Socket client = null;

        try
        {
            ns = CreateAuthenticatedNetState(out client);
            var baseSize = ns._socket.SendBuffer.PhysicalSize;
            NetState._availableMemoryBytes = 1; // any working set is above 80% of one byte

            for (var i = 0; i < 6; i++)
            {
                ns.Send(NetTestHelpers.Pattern(baseSize / 4, i));
            }

            Assert.Equal(baseSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.Contains("Send buffer exhausted", ns._disconnectReason);
        }
        finally
        {
            NetState._availableMemoryBytes = previous;
            ns?.Dispose();
            client?.Close();
        }
    }

    [Fact]
    public void Send_OnClosingSocket_DoesNotGrow()
    {
        var ns = CreateAuthenticatedNetState(out var client);
        var baseSize = ns._socket.SendBuffer.PhysicalSize;

        try
        {
            ns.Disconnect("test");
            NetState.Slice(); // handoff

            // CannotSendPackets short-circuits before any growth
            ns.Send(NetTestHelpers.Pattern(baseSize * 2, 500));

            Assert.Equal(baseSize, ns._socket.SendBuffer.PhysicalSize);
            Assert.False(ns._sendBufferGrown);
        }
        finally
        {
            ns.Dispose();
            client.Close();
        }
    }

    [Theory]
    [InlineData(3 * 1024 * 1024, 4 * 1024 * 1024)]   // not a power of two; rounded up
    [InlineData(512 * 1024 * 1024, 256 * 1024 * 1024)] // above the transport ceiling; capped
    [InlineData(1024, 64 * 1024)]                     // below the minimum; raised
    [InlineData(2 * 1024 * 1024, 2 * 1024 * 1024)]   // already valid; untouched
    public void CoercePowerOfTwoSetting_CoercesToATierTheTransportAccepts(int configured, int expected) =>
        Assert.Equal(expected, NetState.CoercePowerOfTwoSetting("network.sendBufferMaxSize", configured, 64 * 1024));

    [Fact]
    public void CoerceSendBufferGrowthBudget_ClampsNegativeAndRaisesBelowOneSlab()
    {
        const int sendBufferSize = 256 * 1024;
        const int maxSendBufferSize = 2 * 1024 * 1024;
        var minimum = RingSocketManager.MinimumSendBufferGrowthBudget(sendBufferSize);

        Assert.Equal(0L, NetState.CoerceSendBufferGrowthBudget(-1, sendBufferSize, maxSendBufferSize));

        // zero means never grow
        Assert.Equal(0L, NetState.CoerceSendBufferGrowthBudget(0, sendBufferSize, maxSendBufferSize));

        Assert.Equal(minimum, NetState.CoerceSendBufferGrowthBudget(1, sendBufferSize, maxSendBufferSize));
        Assert.Equal(minimum * 4, NetState.CoerceSendBufferGrowthBudget(minimum * 4, sendBufferSize, maxSendBufferSize));

        // growth off; budget untouched
        Assert.Equal(1L, NetState.CoerceSendBufferGrowthBudget(1, sendBufferSize, sendBufferSize));
    }
}