src: renames projects
This commit is contained in:
@@ -0,0 +1,333 @@
|
||||
using System.Collections.Concurrent;
|
||||
using System.Net.WebSockets;
|
||||
using ReverseLlama.Protocol;
|
||||
|
||||
namespace ReverseLlama.Server;
|
||||
|
||||
internal sealed class TunnelConnection
|
||||
{
|
||||
private readonly ConcurrentDictionary<string, PendingCommand> _commands = new();
|
||||
private readonly object _modelsLock = new();
|
||||
private readonly ConcurrentDictionary<string, PendingProxyRequest> _pending = new();
|
||||
private readonly SemaphoreSlim _sendLock = new(1, 1);
|
||||
private readonly WebSocket _socket;
|
||||
private readonly ILogger<TunnelConnection> _logger;
|
||||
private string[] _activeModels = [];
|
||||
private string[] _models = [];
|
||||
private DateTimeOffset? _modelsUpdatedAt;
|
||||
|
||||
public TunnelConnection(string clientId, WebSocket socket, ILogger<TunnelConnection> logger)
|
||||
{
|
||||
ClientId = clientId;
|
||||
_socket = socket;
|
||||
_logger = logger;
|
||||
}
|
||||
|
||||
public string ClientId { get; }
|
||||
|
||||
public string Id { get; } = Guid.NewGuid().ToString("n");
|
||||
|
||||
public bool IsOpen => _socket.State == WebSocketState.Open;
|
||||
|
||||
public int PendingRequestCount => _pending.Count;
|
||||
|
||||
public IReadOnlyList<string> Models
|
||||
{
|
||||
get
|
||||
{
|
||||
lock (_modelsLock)
|
||||
{
|
||||
return _models;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public IReadOnlyList<string> ActiveModels
|
||||
{
|
||||
get
|
||||
{
|
||||
lock (_modelsLock)
|
||||
{
|
||||
return _activeModels;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public DateTimeOffset? ModelsUpdatedAt
|
||||
{
|
||||
get
|
||||
{
|
||||
lock (_modelsLock)
|
||||
{
|
||||
return _modelsUpdatedAt;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public PendingProxyRequest RegisterPending(string requestId)
|
||||
{
|
||||
var pending = new PendingProxyRequest();
|
||||
|
||||
if (!_pending.TryAdd(requestId, pending))
|
||||
{
|
||||
throw new InvalidOperationException($"Request id {requestId} is already registered.");
|
||||
}
|
||||
|
||||
return pending;
|
||||
}
|
||||
|
||||
public void RemovePending(string requestId)
|
||||
{
|
||||
_pending.TryRemove(requestId, out _);
|
||||
}
|
||||
|
||||
public bool HasModel(string model)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(model))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
var requested = model.Trim();
|
||||
|
||||
foreach (var available in Models.Concat(ActiveModels))
|
||||
{
|
||||
if (ModelNamesMatch(requested, available))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
public void UpdateModels(IEnumerable<string> models, IEnumerable<string> activeModels)
|
||||
{
|
||||
var snapshot = models
|
||||
.Where(model => !string.IsNullOrWhiteSpace(model))
|
||||
.Select(model => model.Trim())
|
||||
.Distinct(StringComparer.OrdinalIgnoreCase)
|
||||
.OrderBy(model => model, StringComparer.OrdinalIgnoreCase)
|
||||
.ToArray();
|
||||
var activeSnapshot = activeModels
|
||||
.Where(model => !string.IsNullOrWhiteSpace(model))
|
||||
.Select(model => model.Trim())
|
||||
.Distinct(StringComparer.OrdinalIgnoreCase)
|
||||
.OrderBy(model => model, StringComparer.OrdinalIgnoreCase)
|
||||
.ToArray();
|
||||
|
||||
lock (_modelsLock)
|
||||
{
|
||||
_models = snapshot;
|
||||
_activeModels = activeSnapshot;
|
||||
_modelsUpdatedAt = DateTimeOffset.UtcNow;
|
||||
}
|
||||
}
|
||||
|
||||
public async Task<TunnelMessage> SendModelCommandAsync(
|
||||
string command,
|
||||
string model,
|
||||
string? payloadJson,
|
||||
TimeSpan timeout,
|
||||
CancellationToken cancellationToken)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(command))
|
||||
{
|
||||
throw new ArgumentException("Command is required.", nameof(command));
|
||||
}
|
||||
|
||||
if (string.IsNullOrWhiteSpace(model))
|
||||
{
|
||||
throw new ArgumentException("Model is required.", nameof(model));
|
||||
}
|
||||
|
||||
var requestId = Guid.NewGuid().ToString("n");
|
||||
var pending = new PendingCommand();
|
||||
|
||||
if (!_commands.TryAdd(requestId, pending))
|
||||
{
|
||||
throw new InvalidOperationException($"Command id {requestId} is already registered.");
|
||||
}
|
||||
|
||||
using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
timeoutCts.CancelAfter(timeout);
|
||||
|
||||
try
|
||||
{
|
||||
await SendAsync(
|
||||
new TunnelMessage
|
||||
{
|
||||
Type = TunnelMessageTypes.ModelCommand,
|
||||
RequestId = requestId,
|
||||
Command = command,
|
||||
Model = model,
|
||||
PayloadJson = payloadJson
|
||||
},
|
||||
timeoutCts.Token);
|
||||
|
||||
return await pending.WaitAsync(timeoutCts.Token);
|
||||
}
|
||||
finally
|
||||
{
|
||||
_commands.TryRemove(requestId, out _);
|
||||
}
|
||||
}
|
||||
|
||||
public async Task SendAsync(TunnelMessage message, CancellationToken cancellationToken)
|
||||
{
|
||||
await _sendLock.WaitAsync(cancellationToken);
|
||||
try
|
||||
{
|
||||
if (_socket.State != WebSocketState.Open)
|
||||
{
|
||||
throw new InvalidOperationException("The tunnel client is not connected.");
|
||||
}
|
||||
|
||||
await WebSocketMessageTransport.SendAsync(_socket, message, cancellationToken);
|
||||
}
|
||||
finally
|
||||
{
|
||||
_sendLock.Release();
|
||||
}
|
||||
}
|
||||
|
||||
public async Task RunReceiveLoopAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
try
|
||||
{
|
||||
while (_socket.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
var message = await WebSocketMessageTransport.ReceiveAsync(_socket, cancellationToken);
|
||||
if (message is null)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
Dispatch(message);
|
||||
}
|
||||
}
|
||||
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
_logger.LogWarning(exception, "Tunnel receive loop failed.");
|
||||
}
|
||||
finally
|
||||
{
|
||||
FailAll("Tunnel client disconnected.");
|
||||
}
|
||||
}
|
||||
|
||||
public async Task CloseAsync(string reason, string? closeDescription = null)
|
||||
{
|
||||
try
|
||||
{
|
||||
if (_socket.State is WebSocketState.Open or WebSocketState.CloseReceived)
|
||||
{
|
||||
await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, closeDescription ?? reason, CancellationToken.None);
|
||||
}
|
||||
}
|
||||
catch
|
||||
{
|
||||
}
|
||||
finally
|
||||
{
|
||||
FailAll(reason);
|
||||
}
|
||||
}
|
||||
|
||||
private void Dispatch(TunnelMessage message)
|
||||
{
|
||||
if (message.Type == TunnelMessageTypes.ModelSnapshot)
|
||||
{
|
||||
UpdateModels(message.Models, message.ActiveModels);
|
||||
_logger.LogInformation(
|
||||
"Tunnel client {ClientId} reported {ModelCount} listed and {ActiveModelCount} active model(s).",
|
||||
ClientId,
|
||||
Models.Count,
|
||||
ActiveModels.Count);
|
||||
return;
|
||||
}
|
||||
|
||||
if (message.Type == TunnelMessageTypes.ModelCommandResult)
|
||||
{
|
||||
if (_commands.TryRemove(message.RequestId, out var pendingCommand))
|
||||
{
|
||||
if (!string.IsNullOrWhiteSpace(message.Error))
|
||||
{
|
||||
pendingCommand.Fail(message.Error);
|
||||
}
|
||||
else
|
||||
{
|
||||
pendingCommand.Complete(message);
|
||||
}
|
||||
}
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (string.IsNullOrWhiteSpace(message.RequestId))
|
||||
{
|
||||
_logger.LogDebug("Ignoring tunnel message without a request id: {MessageType}", message.Type);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!_pending.TryGetValue(message.RequestId, out var pending))
|
||||
{
|
||||
_logger.LogDebug("Ignoring tunnel message for unknown request {RequestId}: {MessageType}", message.RequestId, message.Type);
|
||||
return;
|
||||
}
|
||||
|
||||
switch (message.Type)
|
||||
{
|
||||
case TunnelMessageTypes.HttpResponseHeaders:
|
||||
pending.SetResponseHeaders(message);
|
||||
break;
|
||||
|
||||
case TunnelMessageTypes.HttpResponseBody:
|
||||
pending.AddBody(message.Body ?? []);
|
||||
break;
|
||||
|
||||
case TunnelMessageTypes.HttpResponseComplete:
|
||||
pending.Complete();
|
||||
break;
|
||||
|
||||
case TunnelMessageTypes.Error:
|
||||
pending.Fail(message.Error ?? "The tunnel client reported an error.");
|
||||
break;
|
||||
|
||||
default:
|
||||
_logger.LogDebug("Ignoring unsupported tunnel message type from client: {MessageType}", message.Type);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
private static bool ModelNamesMatch(string requested, string available) =>
|
||||
string.Equals(requested, available, StringComparison.OrdinalIgnoreCase)
|
||||
|| string.Equals(StripLatestTag(requested), StripLatestTag(available), StringComparison.OrdinalIgnoreCase);
|
||||
|
||||
private static string StripLatestTag(string model) =>
|
||||
model.EndsWith(":latest", StringComparison.OrdinalIgnoreCase)
|
||||
? model[..^":latest".Length]
|
||||
: model;
|
||||
|
||||
private void FailAll(string reason)
|
||||
{
|
||||
foreach (var pair in _pending.ToArray())
|
||||
{
|
||||
if (_pending.TryRemove(pair.Key, out var pending))
|
||||
{
|
||||
pending.Fail(reason);
|
||||
}
|
||||
}
|
||||
|
||||
foreach (var pair in _commands.ToArray())
|
||||
{
|
||||
if (_commands.TryRemove(pair.Key, out var pending))
|
||||
{
|
||||
pending.Fail(reason);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user