138 lines
4.8 KiB
C#
138 lines
4.8 KiB
C#
using System.Collections.Concurrent;
|
|
using System.Net.WebSockets;
|
|
|
|
namespace AmtScanner.Api.Services;
|
|
|
|
/// <summary>
|
|
/// 屏幕流代理服务 - 管理到 Agent 的 WebSocket 连接
|
|
/// </summary>
|
|
public class ScreenStreamProxyService
|
|
{
|
|
private readonly ILogger<ScreenStreamProxyService> _logger;
|
|
private readonly ConcurrentDictionary<string, AgentConnection> _agentConnections = new();
|
|
|
|
public ScreenStreamProxyService(ILogger<ScreenStreamProxyService> logger)
|
|
{
|
|
_logger = logger;
|
|
}
|
|
|
|
/// <summary>
|
|
/// 代理屏幕流 - 从 Agent 获取屏幕流并转发给客户端
|
|
/// </summary>
|
|
public async Task ProxyScreenStreamAsync(
|
|
WebSocket clientWebSocket,
|
|
string agentIp,
|
|
int agentPort,
|
|
CancellationToken cancellationToken)
|
|
{
|
|
ClientWebSocket? agentWebSocket = null;
|
|
|
|
try
|
|
{
|
|
// 连接到 Agent 的 WebSocket 服务
|
|
agentWebSocket = new ClientWebSocket();
|
|
var agentUri = new Uri($"ws://{agentIp}:{agentPort}/");
|
|
|
|
_logger.LogInformation("正在连接到 Agent: {Uri}", agentUri);
|
|
|
|
using var connectCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
|
|
using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, connectCts.Token);
|
|
|
|
await agentWebSocket.ConnectAsync(agentUri, linkedCts.Token);
|
|
|
|
_logger.LogInformation("已连接到 Agent: {Ip}:{Port}", agentIp, agentPort);
|
|
|
|
// 转发数据
|
|
var buffer = new byte[1024 * 1024]; // 1MB buffer for images
|
|
|
|
while (agentWebSocket.State == WebSocketState.Open &&
|
|
clientWebSocket.State == WebSocketState.Open &&
|
|
!cancellationToken.IsCancellationRequested)
|
|
{
|
|
try
|
|
{
|
|
var result = await agentWebSocket.ReceiveAsync(
|
|
new ArraySegment<byte>(buffer),
|
|
cancellationToken);
|
|
|
|
if (result.MessageType == WebSocketMessageType.Close)
|
|
{
|
|
_logger.LogInformation("Agent 关闭了连接");
|
|
break;
|
|
}
|
|
|
|
if (result.MessageType == WebSocketMessageType.Binary && result.Count > 0)
|
|
{
|
|
// 转发给客户端
|
|
await clientWebSocket.SendAsync(
|
|
new ArraySegment<byte>(buffer, 0, result.Count),
|
|
WebSocketMessageType.Binary,
|
|
result.EndOfMessage,
|
|
cancellationToken);
|
|
}
|
|
}
|
|
catch (WebSocketException ex)
|
|
{
|
|
_logger.LogWarning("WebSocket 错误: {Message}", ex.Message);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
catch (WebSocketException ex)
|
|
{
|
|
_logger.LogError("连接 Agent 失败: {Message}", ex.Message);
|
|
|
|
// 发送错误消息给客户端
|
|
if (clientWebSocket.State == WebSocketState.Open)
|
|
{
|
|
var errorMsg = System.Text.Encoding.UTF8.GetBytes($"ERROR:无法连接到设备 {agentIp}:{agentPort}");
|
|
await clientWebSocket.SendAsync(
|
|
new ArraySegment<byte>(errorMsg),
|
|
WebSocketMessageType.Text,
|
|
true,
|
|
cancellationToken);
|
|
}
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
_logger.LogInformation("代理连接已取消");
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
_logger.LogError(ex, "代理屏幕流时发生错误");
|
|
}
|
|
finally
|
|
{
|
|
if (agentWebSocket != null)
|
|
{
|
|
try
|
|
{
|
|
if (agentWebSocket.State == WebSocketState.Open)
|
|
{
|
|
await agentWebSocket.CloseAsync(
|
|
WebSocketCloseStatus.NormalClosure,
|
|
"Proxy closing",
|
|
CancellationToken.None);
|
|
}
|
|
agentWebSocket.Dispose();
|
|
}
|
|
catch { }
|
|
}
|
|
|
|
_logger.LogInformation("代理连接已关闭: {Ip}:{Port}", agentIp, agentPort);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Agent 连接信息
|
|
/// </summary>
|
|
public class AgentConnection
|
|
{
|
|
public string Uuid { get; set; } = "";
|
|
public string IpAddress { get; set; } = "";
|
|
public int Port { get; set; }
|
|
public ClientWebSocket? WebSocket { get; set; }
|
|
public DateTime ConnectedAt { get; set; }
|
|
}
|