using System; using System.Collections.Generic; using System.IO; using System.Runtime.InteropServices; using Prometheus; using Robust.Shared.Asynchronous; using Robust.Shared.Log; namespace Robust.Shared.Network.Transfer; internal abstract partial class BaseTransferManager { internal static readonly Counter SentDataMetrics = Metrics.CreateCounter( "robust_transfer_sent_bytes", "Number of bytes sent via the transfer system"); internal static readonly Counter ReceivedDataMetrics = Metrics.CreateCounter( "robust_transfer_received_bytes", "Number of bytes received via the transfer system"); private readonly NetMessageAccept _side; private readonly ITaskManager _taskManager; protected readonly Dictionary RegisteredKeys = []; protected readonly ISawmill Sawmill; private protected BaseTransferManager( ILogManager logManager, NetMessageAccept side, ITaskManager taskManager) { _side = side; _taskManager = taskManager; Sawmill = logManager.GetSawmill("net.transfer"); } public void RegisterTransferMessage( string key, Action? rxCallback = null, NetMessageAccept accept = NetMessageAccept.Both) { if ((accept & ~NetMessageAccept.Both) != 0) throw new ArgumentException("Invalid accept given: must be client, server, or both"); ref var slot = ref CollectionsMarshal.GetValueRefOrAddDefault(RegisteredKeys, key, out var exists); if (exists) throw new InvalidOperationException($"Key '{key}' was already registered!"); slot = new RegisteredKey(); if ((accept & _side) > 0) slot.Callback = rxCallback; } internal void TransferReceived(string key, INetChannel channel, Stream stream) { if (!RegisteredKeys.TryGetValue(key, out var registered)) throw new Exception($"Unknown key: {key}"); if (registered.Callback == null) throw new Exception($"Key is send-only: {key}"); _taskManager.RunOnMainThread(() => { registered.Callback(new TransferReceivedEvent(key, channel, stream)); }); } protected void CheckRegistered(TransferStartInfo info) { if (!RegisteredKeys.ContainsKey(info.MessageKey)) throw new ArgumentException($"Key is not registered: {info.MessageKey}"); } protected sealed class RegisteredKey { public Action? Callback; } }