TermRemoteCtl/apps/windows_agent/tests/TermRemoteCtl.Agent.IntegrationTests/Realtime/TerminalWebSocketHandlerTests.cs

519 lines
22 KiB
C#

using System.Net;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.Mvc.Testing;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.DependencyInjection.Extensions;
using Microsoft.Extensions.Options;
using TermRemoteCtl.Agent.Configuration;
using TermRemoteCtl.Agent.Sessions;
using TermRemoteCtl.Agent.Terminal;
namespace TermRemoteCtl.Agent.IntegrationTests.Realtime;
public sealed class TerminalWebSocketHandlerTests
{
[Fact]
[Trait("Track", "Mainline")]
public async Task Attach_Streams_Output_And_Forwards_Input()
{
await using var fixture = new TerminalApiFixture();
var registry = fixture.Services.GetRequiredService<SessionRegistry>();
fixture.TerminalHost.Registry = registry;
var session = registry.Create("Shell", DateTimeOffset.UtcNow);
using WebSocket socket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
var attachedFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var attachedPayload = JsonSerializer.Deserialize<TerminalAttachResponse>(
attachedFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(attachedPayload);
Assert.Equal("attached", attachedPayload!.Type);
Assert.Equal(session.SessionId, attachedPayload.SessionId);
var restoreFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var restorePayload = JsonSerializer.Deserialize<TerminalRestoreResponse>(
restoreFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(restorePayload);
Assert.Equal("restore", restorePayload!.Type);
Assert.Equal(session.SessionId, restorePayload.SessionId);
Assert.Equal(1L, restorePayload.Sequence);
fixture.TerminalHost.EmitOutput(session.SessionId, "abc");
fixture.TerminalHost.EmitOutput(session.SessionId, "def");
var firstOutputFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var firstOutputPayload = JsonSerializer.Deserialize<TerminalOutputResponse>(
firstOutputFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(firstOutputPayload);
Assert.Equal("output", firstOutputPayload!.Type);
Assert.Equal(2L, firstOutputPayload.Sequence);
Assert.Equal("abc", firstOutputPayload.Chunk);
var secondOutputFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var secondOutputPayload = JsonSerializer.Deserialize<TerminalOutputResponse>(
secondOutputFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(secondOutputPayload);
Assert.Equal("output", secondOutputPayload!.Type);
Assert.Equal(3L, secondOutputPayload.Sequence);
Assert.Equal("def", secondOutputPayload.Chunk);
var inputMessage = JsonSerializer.Serialize(new { type = "input", input = "dir" });
await socket.SendAsync(Encoding.UTF8.GetBytes(inputMessage), WebSocketMessageType.Text, true, CancellationToken.None);
await WaitForConditionAsync(() => fixture.TerminalHost.Inputs.Contains(("input", session.SessionId, "dir")), TimeSpan.FromSeconds(2));
await WaitForConditionAsync(
() => fixture.Diagnostics.Events.Contains(("backend.input.received", session.SessionId, "dir")),
TimeSpan.FromSeconds(2));
}
[Fact]
[Trait("Track", "Mainline")]
public async Task Attach_Replays_Recent_Output_For_Existing_Session()
{
await using var fixture = new TerminalApiFixture();
var registry = fixture.Services.GetRequiredService<SessionRegistry>();
fixture.TerminalHost.Registry = registry;
var session = registry.Create("Shell", DateTimeOffset.UtcNow);
await registry.AppendOutputAsync(session.SessionId, "prompt> dir\r\nnext> ", CancellationToken.None);
using WebSocket socket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
var attachedFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var attachedPayload = JsonSerializer.Deserialize<TerminalAttachResponse>(
attachedFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(attachedPayload);
Assert.Equal("attached", attachedPayload!.Type);
var restoreFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var restorePayload = JsonSerializer.Deserialize<TerminalRestoreResponse>(
restoreFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(restorePayload);
Assert.Equal("restore", restorePayload!.Type);
Assert.Equal("prompt> dir\r\nnext> ", restorePayload.ScreenText);
Assert.Equal(string.Empty, restorePayload.PendingInput);
Assert.Equal(2L, restorePayload.Sequence);
}
[Fact]
[Trait("Track", "Mainline")]
public async Task Attach_Does_Not_Duplicate_Output_Produced_During_Replay_Boundary()
{
await using var fixture = new TerminalApiFixture();
var registry = fixture.Services.GetRequiredService<SessionRegistry>();
fixture.TerminalHost.Registry = registry;
var session = registry.Create("Shell", DateTimeOffset.UtcNow);
await registry.AppendOutputAsync(session.SessionId, "prompt> dir\r\n", CancellationToken.None);
fixture.TerminalHost.EmitOutputOnNextSubscription(session.SessionId, "next> ");
using WebSocket socket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
var attachedFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var attachedPayload = JsonSerializer.Deserialize<TerminalAttachResponse>(
attachedFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(attachedPayload);
Assert.Equal("attached", attachedPayload!.Type);
var restoreFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var restorePayload = JsonSerializer.Deserialize<TerminalRestoreResponse>(
restoreFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(restorePayload);
Assert.Equal("prompt> dir\r\n", restorePayload!.ScreenText);
Assert.Equal(string.Empty, restorePayload.PendingInput);
Assert.Equal(2L, restorePayload.Sequence);
var liveFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var livePayload = JsonSerializer.Deserialize<TerminalOutputResponse>(
liveFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(livePayload);
Assert.Equal("output", livePayload!.Type);
Assert.Equal(3L, livePayload.Sequence);
Assert.Equal("next> ", livePayload.Chunk);
await AssertNoAdditionalTextFrameAsync(socket, TimeSpan.FromMilliseconds(200));
}
[Fact]
[Trait("Track", "Mainline")]
public async Task Reattach_Replays_Visible_User_Input_When_No_Output_Echo_Has_Arrived_Yet()
{
await using var fixture = new TerminalApiFixture();
var registry = fixture.Services.GetRequiredService<SessionRegistry>();
fixture.TerminalHost.Registry = registry;
var session = registry.Create("Shell", DateTimeOffset.UtcNow);
using (WebSocket socket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None))
{
_ = await ReceiveTextAsync(socket, CancellationToken.None);
_ = await ReceiveTextAsync(socket, CancellationToken.None);
var inputMessage = JsonSerializer.Serialize(new { type = "input", input = "dir" });
await socket.SendAsync(Encoding.UTF8.GetBytes(inputMessage), WebSocketMessageType.Text, true, CancellationToken.None);
await WaitForConditionAsync(() => fixture.TerminalHost.Inputs.Contains(("input", session.SessionId, "dir")), TimeSpan.FromSeconds(2));
}
using WebSocket replaySocket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
_ = await ReceiveTextAsync(replaySocket, CancellationToken.None);
var restoreFrame = await ReceiveTextAsync(replaySocket, CancellationToken.None);
var restorePayload = JsonSerializer.Deserialize<TerminalRestoreResponse>(
restoreFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(restorePayload);
Assert.Equal("dir", restorePayload!.PendingInput);
Assert.Equal(string.Empty, restorePayload.ScreenText);
Assert.Equal(4L, restorePayload.Sequence);
}
[Fact]
[Trait("Track", "Mainline")]
public async Task Reattach_Returns_Restore_Payload_With_Pending_Input()
{
await using var fixture = new TerminalApiFixture();
var registry = fixture.Services.GetRequiredService<SessionRegistry>();
fixture.TerminalHost.Registry = registry;
var session = registry.Create("Shell", DateTimeOffset.UtcNow);
await registry.RecordInputAsync(session.SessionId, "dir", CancellationToken.None);
using WebSocket socket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
_ = await ReceiveTextAsync(socket, CancellationToken.None);
var restoreFrame = await ReceiveTextAsync(socket, CancellationToken.None);
Assert.Contains("\"type\":\"restore\"", restoreFrame);
Assert.Contains("\"pendingInput\":\"dir\"", restoreFrame);
Assert.Contains("\"sequence\":2", restoreFrame);
}
[Fact]
[Trait("Track", "Mainline")]
public async Task Input_Acknowledgement_Deduplicates_Replayed_Client_Input()
{
await using var fixture = new TerminalApiFixture();
var registry = fixture.Services.GetRequiredService<SessionRegistry>();
fixture.TerminalHost.Registry = registry;
var session = registry.Create("Shell", DateTimeOffset.UtcNow);
using WebSocket firstSocket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
_ = await ReceiveTextAsync(firstSocket, CancellationToken.None);
_ = await ReceiveTextAsync(firstSocket, CancellationToken.None);
var inputMessage = JsonSerializer.Serialize(new { type = "input", input = "dir", inputId = "input-1" });
await firstSocket.SendAsync(Encoding.UTF8.GetBytes(inputMessage), WebSocketMessageType.Text, true, CancellationToken.None);
var firstAckFrame = await ReceiveTextAsync(firstSocket, CancellationToken.None);
Assert.Contains("\"type\":\"inputAck\"", firstAckFrame);
Assert.Contains("\"inputId\":\"input-1\"", firstAckFrame);
using WebSocket secondSocket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
_ = await ReceiveTextAsync(secondSocket, CancellationToken.None);
_ = await ReceiveTextAsync(secondSocket, CancellationToken.None);
await secondSocket.SendAsync(Encoding.UTF8.GetBytes(inputMessage), WebSocketMessageType.Text, true, CancellationToken.None);
var duplicateAckFrame = await ReceiveTextAsync(secondSocket, CancellationToken.None);
Assert.Contains("\"type\":\"inputAck\"", duplicateAckFrame);
Assert.Contains("\"inputId\":\"input-1\"", duplicateAckFrame);
Assert.Single(fixture.TerminalHost.Inputs.Where(item => item == ("input", session.SessionId, "dir")));
}
[Fact]
[Trait("Track", "Mainline")]
public async Task Attach_Includes_Screen_Snapshot_When_Backend_Screen_Protocol_Is_Enabled()
{
await using var fixture = new TerminalApiFixture(enableBackendScreenProtocol: true);
var registry = fixture.Services.GetRequiredService<SessionRegistry>();
fixture.TerminalHost.Registry = registry;
var session = registry.Create("Shell", DateTimeOffset.UtcNow);
await registry.RecordResizeAsync(session.SessionId, 40, 10, CancellationToken.None);
await registry.RecordOutputAsync(session.SessionId, "\u001b[2J\u001b[Hprompt> dir", CancellationToken.None);
using WebSocket socket = await fixture.Server.CreateWebSocketClient().ConnectAsync(
new Uri($"ws://localhost/ws/terminal?sessionId={session.SessionId}"),
CancellationToken.None);
_ = await ReceiveTextAsync(socket, CancellationToken.None);
var restoreFrame = await ReceiveTextAsync(socket, CancellationToken.None);
var restorePayload = JsonSerializer.Deserialize<TerminalRestoreResponse>(
restoreFrame,
new JsonSerializerOptions(JsonSerializerDefaults.Web));
Assert.NotNull(restorePayload);
Assert.NotNull(restorePayload!.ScreenSnapshot);
Assert.Equal("primary", restorePayload.ScreenSnapshot!.ActiveBuffer);
Assert.Equal(2L, restorePayload.ScreenSnapshot.SourceSequence);
}
private static async Task<string> ReceiveTextAsync(WebSocket socket, CancellationToken cancellationToken)
{
var buffer = new byte[4096];
using var stream = new MemoryStream();
while (true)
{
var result = await socket.ReceiveAsync(buffer, cancellationToken);
if (result.MessageType == WebSocketMessageType.Close)
{
throw new InvalidOperationException("Socket closed before a text message was received.");
}
stream.Write(buffer, 0, result.Count);
if (result.EndOfMessage)
{
return Encoding.UTF8.GetString(stream.ToArray());
}
}
}
private static async Task WaitForConditionAsync(Func<bool> condition, TimeSpan timeout)
{
var deadline = DateTimeOffset.UtcNow.Add(timeout);
while (DateTimeOffset.UtcNow < deadline)
{
if (condition())
{
return;
}
await Task.Delay(25);
}
throw new TimeoutException("Condition was not met.");
}
private static async Task AssertNoAdditionalTextFrameAsync(WebSocket socket, TimeSpan timeout)
{
using var cts = new CancellationTokenSource(timeout);
await Assert.ThrowsAnyAsync<OperationCanceledException>(
async () => await ReceiveTextAsync(socket, cts.Token));
}
private sealed class TerminalApiFixture : WebApplicationFactory<Program>
{
private readonly string _dataRoot = Path.Combine(Path.GetTempPath(), "TermRemoteCtl.Tests", Guid.NewGuid().ToString("N"));
private readonly TestTerminalSessionHost _terminalHost;
private readonly RecordingTerminalDiagnosticsSink _diagnostics = new();
private readonly bool _enableBackendScreenProtocol;
public TerminalApiFixture(bool enableBackendScreenProtocol = false)
{
_terminalHost = new TestTerminalSessionHost();
_enableBackendScreenProtocol = enableBackendScreenProtocol;
}
public TestTerminalSessionHost TerminalHost => _terminalHost;
public RecordingTerminalDiagnosticsSink Diagnostics => _diagnostics;
protected override void ConfigureWebHost(IWebHostBuilder builder)
{
builder.UseEnvironment("Development");
builder.ConfigureAppConfiguration((_, configBuilder) =>
{
configBuilder.AddInMemoryCollection(new Dictionary<string, string?>
{
["Agent:DataRoot"] = _dataRoot,
["Agent:BindAddress"] = "127.0.0.1",
["Agent:HttpsPort"] = "9443",
["Agent:WebSocketFrameFlushMilliseconds"] = "33",
["Agent:RingBufferLineLimit"] = "4000",
["Agent:EnableBackendScreenProtocol"] = _enableBackendScreenProtocol.ToString()
});
});
builder.ConfigureServices(services =>
{
services.PostConfigure<AgentOptions>(options =>
{
options.DataRoot = _dataRoot;
options.BindAddress = "127.0.0.1";
options.HttpsPort = 9443;
options.WebSocketFrameFlushMilliseconds = 33;
options.RingBufferLineLimit = 4000;
options.EnableBackendScreenProtocol = _enableBackendScreenProtocol;
});
services.RemoveAll<ISessionHost>();
services.AddSingleton(_terminalHost);
services.AddSingleton<ISessionHost>(sp => sp.GetRequiredService<TestTerminalSessionHost>());
services.AddSingleton<ITerminalDiagnosticsSink>(_diagnostics);
});
}
public new async ValueTask DisposeAsync()
{
await base.DisposeAsync();
if (!Directory.Exists(_dataRoot))
{
return;
}
for (var attempt = 0; attempt < 5; attempt += 1)
{
try
{
Directory.Delete(_dataRoot, true);
return;
}
catch (IOException) when (attempt < 4)
{
await Task.Delay(50);
}
catch (UnauthorizedAccessException) when (attempt < 4)
{
await Task.Delay(50);
}
}
}
}
private sealed class RecordingTerminalDiagnosticsSink : ITerminalDiagnosticsSink
{
private readonly List<(string EventName, string SessionId, string Detail)> _events = [];
public IReadOnlyList<(string EventName, string SessionId, string Detail)> Events => _events;
public void Record(string eventName, string sessionId, string detail)
{
_events.Add((eventName, sessionId, detail));
}
}
private sealed class TestTerminalSessionHost : ISessionHost
{
private readonly List<(string Kind, string SessionId, string Value)> _inputs = new();
private EventHandler<TerminalOutputEventArgs>? _outputReceived;
private readonly object _gate = new();
private (string SessionId, string Chunk)? _pendingSubscriptionEmission;
public event EventHandler<TerminalOutputEventArgs>? OutputReceived
{
add
{
(string SessionId, string Chunk)? emission;
lock (_gate)
{
_outputReceived += value;
emission = _pendingSubscriptionEmission;
_pendingSubscriptionEmission = null;
}
if (emission is { } pending)
{
EmitOutput(pending.SessionId, pending.Chunk);
}
}
remove
{
lock (_gate)
{
_outputReceived -= value;
}
}
}
public SessionRegistry? Registry { get; set; }
public IReadOnlyList<(string Kind, string SessionId, string Value)> Inputs => _inputs;
public Task StartAsync(string sessionId, CancellationToken cancellationToken)
{
return Task.CompletedTask;
}
public Task StopAsync(string sessionId, CancellationToken cancellationToken)
{
return Task.CompletedTask;
}
public Task WriteInputAsync(string sessionId, string input, CancellationToken cancellationToken)
{
_inputs.Add(("input", sessionId, input));
return Task.CompletedTask;
}
public Task ResizeAsync(string sessionId, int columns, int rows, CancellationToken cancellationToken)
{
_inputs.Add(("resize", sessionId, $"{columns}x{rows}"));
return Task.CompletedTask;
}
public void EmitOutputOnNextSubscription(string sessionId, string chunk)
{
lock (_gate)
{
_pendingSubscriptionEmission = (sessionId, chunk);
}
}
public void EmitOutput(string sessionId, string chunk)
{
var ioEvent = Registry?.RecordOutputAsync(sessionId, chunk, CancellationToken.None).GetAwaiter().GetResult();
_outputReceived?.Invoke(
this,
new TerminalOutputEventArgs(sessionId, chunk, ioEvent?.Sequence ?? 0));
}
}
private sealed record TerminalAttachResponse(string SessionId, string Type);
private sealed record TerminalRestoreResponse(
string SessionId,
long Sequence,
string ScreenText,
string PendingInput,
int? CursorRow,
int? CursorColumn,
TerminalScreenSnapshotResponse? ScreenSnapshot,
string Type);
private sealed record TerminalScreenSnapshotResponse(
long ScreenVersion,
long SourceSequence,
string ActiveBuffer);
private sealed record TerminalOutputResponse(
string SessionId,
long Sequence,
string Chunk,
string Type);
}