Rejig Device networks (#7380)

Content.Server/DeviceNetwork/Systems/DeviceNetworkSystem.cs

@@ -1,10 +1,10 @@
 using Content.Server.DeviceNetwork.Components;
+using Content.Shared.DeviceNetwork;
 using JetBrains.Annotations;
-using Robust.Shared.GameObjects;
-using Robust.Shared.IoC;
+using Robust.Shared.Prototypes;
 using Robust.Shared.Random;
-using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
+using static Content.Server.DeviceNetwork.Components.DeviceNetworkComponent;
 
 namespace Content.Server.DeviceNetwork.Systems
 {
@@ -16,15 +16,17 @@ namespace Content.Server.DeviceNetwork.Systems
     public sealed class DeviceNetworkSystem : EntitySystem
     {
         [Dependency] private readonly IRobustRandom _random = default!;
+        [Dependency] private readonly IPrototypeManager _protoMan = default!;
+        [Dependency] private readonly SharedTransformSystem _transformSystem = default!;
 
-        private readonly Dictionary<DeviceNetworkComponent.ConnectionType, List<DeviceNetworkComponent>> _connections = new();
-        private readonly Queue<NetworkPacket> _packets = new();
+        private readonly Dictionary<ConnectionType, DeviceNet> _networks = new();
+        private readonly Queue<DeviceNetworkPacketEvent> _packets = new();
 
         public override void Initialize()
         {
             base.Initialize();
 
-            SubscribeLocalEvent<DeviceNetworkComponent, ComponentStartup>(OnNetworkStarted);
+            SubscribeLocalEvent<DeviceNetworkComponent, MapInitEvent>(OnMapInit);
             SubscribeLocalEvent<DeviceNetworkComponent, ComponentShutdown>(OnNetworkShutdown);
         }
 
@@ -36,72 +38,53 @@ namespace Content.Server.DeviceNetwork.Systems
             {
                 var packet = _packets.Dequeue();
 
-                if(packet.Broadcast)
-                {
-                    BroadcastPacket(packet);
-                    continue;
-                }
                 SendPacket(packet);
             }
         }
 
-        /// <summary>
-        /// Manually connect an entity with a DeviceNetworkComponent.
-        /// </summary>
-        /// <param name="uid">The Entity containing a DeviceNetworkComponent</param>
-        public void Connect(EntityUid uid)
-        {
-            if (EntityManager.TryGetComponent<DeviceNetworkComponent>(uid, out var component))
-            {
-                AddConnection(component);
-            }
-        }
-
         /// <summary>
         /// Sends the given payload as a device network packet to the entity with the given address and frequency.
         /// Addresses are given to the DeviceNetworkComponent of an entity when connecting.
         /// </summary>
         /// <param name="uid">The EntityUid of the sending entity</param>
-        /// <param name="address">The address of the entity that the packet gets sent to when not broadcasting</param>
+        /// <param name="address">The address of the entity that the packet gets sent to. If null, the message is broadcast to all devices on that frequency (except the sender)</param>
         /// <param name="frequency">The frequency to send on</param>
         /// <param name="data">The data to be sent</param>
-        /// <param name="broadcast">Send to all devices on the same device network on the given frequency</param>
-        public void QueuePacket(EntityUid uid, string address, int frequency, NetworkPayload data, bool broadcast = false)
+        public void QueuePacket(EntityUid uid, string? address, NetworkPayload data, uint? frequency = null, DeviceNetworkComponent? device = null)
         {
-            if (EntityManager.TryGetComponent<DeviceNetworkComponent>(uid, out var component))
-            {
-                var packet = new NetworkPacket
-                {
-                    NetId = component.DeviceNetId,
-                    Address = address,
-                    Frequency = frequency,
-                    Broadcast = broadcast,
-                    Data = data,
-                    Sender = component
-                };
+            if (!Resolve(uid, ref device, false))
+                return;
 
-                _packets.Enqueue(packet);
-            }
+            if (device.Address == null)
+                return;
+
+            frequency ??= device.TransmitFrequency;
+
+            if (frequency != null)
+                _packets.Enqueue(new DeviceNetworkPacketEvent(device.DeviceNetId, address, frequency.Value, device.Address, uid, data));
         }
 
         /// <summary>
-        /// Manually disconnect an entity with a DeviceNetworkComponent.
+        /// Automatically attempt to connect some devices when a map starts.
         /// </summary>
-        /// <param name="uid">The Entity containing a DeviceNetworkComponent</param>
-        public void Disconnect(EntityUid uid)
+        private void OnMapInit(EntityUid uid, DeviceNetworkComponent device, MapInitEvent args)
         {
-            if (EntityManager.TryGetComponent<DeviceNetworkComponent>(uid, out var component))
+            if (device.ReceiveFrequency == null
+                && device.ReceiveFrequencyId != null
+                && _protoMan.TryIndex<DeviceFrequencyPrototype>(device.ReceiveFrequencyId, out var receive))
             {
-                RemoveConnection(component);
+                device.ReceiveFrequency = receive.Frequency;
             }
-        }
 
-        /// <summary>
-        /// Automatically connect when an entity with a DeviceNetworkComponent starts up.
-        /// </summary>
-        private void OnNetworkStarted(EntityUid uid, DeviceNetworkComponent component, ComponentStartup args)
-        {
-            AddConnection(component);
+            if (device.TransmitFrequency == null
+                && device.TransmitFrequencyId != null
+                && _protoMan.TryIndex<DeviceFrequencyPrototype>(device.TransmitFrequencyId, out var xmit))
+            {
+                device.TransmitFrequency = xmit.Frequency;
+            }
+
+            if (device.AutoConnect)
+                ConnectDevice(uid, device);
         }
 
         /// <summary>
@@ -109,127 +92,197 @@ namespace Content.Server.DeviceNetwork.Systems
         /// </summary>
         private void OnNetworkShutdown(EntityUid uid, DeviceNetworkComponent component, ComponentShutdown args)
         {
-            RemoveConnection(component);
-        }
-
-        private bool AddConnection(DeviceNetworkComponent connection)
-        {
-            var netId = connection.DeviceNetId;
-            if (!_connections.ContainsKey(netId))
-                _connections[netId] = new List<DeviceNetworkComponent>();
-
-            if (!_connections[netId].Contains(connection))
-            {
-                connection.Address = GenerateValidAddress(netId);
-                _connections[netId].Add(connection);
-                connection.Open = true;
-                return true;
-            }
-
-            return false;
-        }
-
-        private bool RemoveConnection(DeviceNetworkComponent connection)
-        {
-            connection.Address = "";
-            connection.Open = false;
-            return _connections[connection.DeviceNetId].Remove(connection);
+            DisconnectDevice(uid, component, false);
         }
 
         /// <summary>
-        /// Generates a valid address by randomly generating one and checking if it already exists on the device network with the given device netId.
+        /// Connect an entity with a DeviceNetworkComponent. Note that this will re-use an existing address if the
+        /// device already had one configured. If there is a clash, the device cannot join the network.
         /// </summary>
-        private string GenerateValidAddress(DeviceNetworkComponent.ConnectionType netId)
+        public bool ConnectDevice(EntityUid uid, DeviceNetworkComponent? device = null)
         {
-            var unique = false;
-            var connections = _connections[netId];
-            var address = "";
+            if (!Resolve(uid, ref device, false))
+                return false;
 
-            while (!unique)
+            if (!_networks.TryGetValue(device.DeviceNetId, out var network))
             {
-                address = _random.Next().ToString("x");
-                unique = !connections.Exists(connection => connection.Address == address);
+                network = new(device.DeviceNetId, _random);
+                _networks[device.DeviceNetId] = network;
             }
 
-            return address;
+            return network.Add(device);
         }
 
-        private List<DeviceNetworkComponent> ConnectionsForFrequency(DeviceNetworkComponent.ConnectionType netId, int frequency)
+        /// <summary>
+        /// Disconnect an entity with a DeviceNetworkComponent.
+        /// </summary>
+        public bool DisconnectDevice(EntityUid uid, DeviceNetworkComponent? device, bool preventAutoConnect = true)
         {
-            if (!_connections.ContainsKey(netId))
-                return new List<DeviceNetworkComponent>();
+            if (!Resolve(uid, ref device, false))
+                return false;
 
-            var result = _connections[netId].FindAll(connection => connection.Frequency == frequency);
+            // If manually disconnected, don't auto reconnect when a game state is loaded.
+            if (preventAutoConnect)
+                device.AutoConnect = false;
 
-            return result;
-        }
-
-        private bool TryGetConnectionWithAddress(DeviceNetworkComponent.ConnectionType netId, int frequency, string address, [NotNullWhen(true)] out DeviceNetworkComponent connection)
-        {
-            var connections = ConnectionsForFrequency(netId, frequency);
-
-            var result = connections.Find(dvc => dvc.Address == address);
-
-            if(result != null)
+            if (!_networks.TryGetValue(device.DeviceNetId, out var network))
             {
-                connection = result;
-                return true;
+                return false;
             }
 
-            connection = default!;
-            return false;
+            if (!network.Remove(device))
+                return false;
+
+            if (network.Devices.Count == 0)
+                _networks.Remove(device.DeviceNetId);
+
+            return true;
         }
 
-        private List<DeviceNetworkComponent> ConnectionsWithReceiveAll(DeviceNetworkComponent.ConnectionType netId, int frequency)
+        #region Get Device
+        /// <summary>
+        /// Get a list of devices listening on a given frequency on some network.
+        /// </summary>
+        private HashSet<DeviceNetworkComponent> GetListeningDevices(ConnectionType netId, uint frequency)
         {
-            if (!_connections.ContainsKey(netId))
-                return new List<DeviceNetworkComponent>();
+            if (_networks.TryGetValue(netId, out var network) && network.ListeningDevices.TryGetValue(frequency, out var devices))
+                return devices;
 
-            var result = _connections[netId].FindAll(device => device.Frequency == frequency && device.ReceiveAll);
-
-            return result;
+            return new();
         }
 
-        private void SendPacket(NetworkPacket packet)
+        /// <summary>
+        /// Get a list of devices listening for ANY transmission on a given frequency, rather than just broadcast & addressed events.
+        /// </summary>
+        private HashSet<DeviceNetworkComponent> GetRecieveAllDevices(ConnectionType netId, uint frequency)
         {
-            if (!TryGetConnectionWithAddress(packet.NetId, packet.Frequency, packet.Address, out var connection))
-                return;
+            if (_networks.TryGetValue(netId, out var network) && network.ReceiveAllDevices.TryGetValue(frequency, out var devices))
+                return devices;
 
-            var receivers = ConnectionsWithReceiveAll(packet.NetId, packet.Frequency);
-            receivers.Add(connection);
-
-            SendToConnections(receivers, packet);
+            return new();
         }
 
-        private void BroadcastPacket(NetworkPacket packet)
+        /// <summary>
+        ///     Try to find a device on a network using its address.
+        /// </summary>
+        private bool TryGetDevice(ConnectionType netId, string address, [NotNullWhen(true)] out DeviceNetworkComponent? device)
         {
-            var receivers = ConnectionsForFrequency(packet.NetId, packet.Frequency);
-            SendToConnections(receivers, packet);
+            if (!_networks.TryGetValue(netId, out var network))
+            {
+                device = null;
+                return false;
+            }
+
+            return network.Devices.TryGetValue(address, out device);
+        }
+        #endregion
+
+        #region Packet Sending
+        private void SendPacket(DeviceNetworkPacketEvent packet)
+        {
+            HashSet<DeviceNetworkComponent> recipients;
+
+            if (packet.Address == null)
+            {
+                // Broadcast to all listening devices
+                recipients = GetListeningDevices(packet.NetId, packet.Frequency);
+            }
+            else
+            {
+                // Add devices listening to all messages
+                recipients = new(GetRecieveAllDevices(packet.NetId, packet.Frequency));
+
+                // add the intended recipient (if they are even listening).
+                if (TryGetDevice(packet.NetId, packet.Address, out var device) && device.ReceiveFrequency == packet.Frequency)
+                    recipients.Add(device);
+            }
+
+            SendToConnections(recipients, packet);
         }
 
-        private void SendToConnections(List<DeviceNetworkComponent> connections, NetworkPacket packet)
+        private void SendToConnections(HashSet<DeviceNetworkComponent> connections, DeviceNetworkPacketEvent packet)
         {
+            var xform = Transform(packet.Sender);
+
+            BeforePacketSentEvent beforeEv = new(packet.Sender, xform, _transformSystem.GetWorldPosition(xform));
+
             foreach (var connection in connections)
             {
-                var beforeEvent = new BeforePacketSentEvent(packet.Sender.Owner);
-                RaiseLocalEvent(connection.Owner, beforeEvent, false);
+                if (connection.Owner == packet.Sender)
+                    continue;
 
-                if (!beforeEvent.Cancelled)
-                {
-                    RaiseLocalEvent(connection.Owner, new PacketSentEvent(connection.Frequency, packet.Sender.Address, packet.Data, packet.Broadcast) , false);
-                }
+                RaiseLocalEvent(connection.Owner, beforeEv, false);
+
+                if (!beforeEv.Cancelled)
+                    RaiseLocalEvent(connection.Owner, packet, false);
+                else
+                    beforeEv.Uncancel();
+            }
+        }
+        #endregion
+
+        #region Component Setter Functions
+        public void SetReceiveFrequency(EntityUid uid, uint? frequency, DeviceNetworkComponent? device = null)
+        {
+            if (!Resolve(uid, ref device, false))
+                return;
+
+            if (device.ReceiveFrequency == frequency)
+                return;
+
+            if (!_networks.TryGetValue(device.DeviceNetId, out var deviceNet) || !deviceNet.UpdateReceiveFrequency(device.Address, frequency))
+                device.ReceiveFrequency = frequency;
+        }
+
+        public void SetTransmitFrequency(EntityUid uid, uint? frequency, DeviceNetworkComponent? device = null)
+        {
+            if (Resolve(uid, ref device, false))
+                device.TransmitFrequency = frequency;
+        }
+
+        public void SetReceiveAll(EntityUid uid, bool receiveAll, DeviceNetworkComponent? device = null)
+        {
+            if (!Resolve(uid, ref device, false))
+                return;
+
+            if (device.ReceiveAll == receiveAll)
+                return;
+
+            if (!_networks.TryGetValue(device.DeviceNetId, out var deviceNet) || !deviceNet.UpdateReceiveAll(device.Address, receiveAll))
+                device.ReceiveAll = receiveAll;
+        }
+
+        public void SetAddress(EntityUid uid, string address, DeviceNetworkComponent? device = null)
+        {
+            if (!Resolve(uid, ref device, false))
+                return;
+
+            if (device.Address == address)
+            {
+                device.CustomAddress = true;
+                return;
+            }
+
+            if (!_networks.TryGetValue(device.DeviceNetId, out var deviceNet) || !deviceNet.UpdateAddress(device.Address, address))
+            {
+                device.Address = address;
+                device.CustomAddress = true;
             }
         }
 
-        internal struct NetworkPacket
+        public void RandomizeAddress(EntityUid uid, string address, DeviceNetworkComponent? device = null)
         {
-            public DeviceNetworkComponent.ConnectionType NetId;
-            public int Frequency;
-            public string Address;
-            public bool Broadcast;
-            public NetworkPayload Data;
-            public DeviceNetworkComponent Sender;
+            if (!Resolve(uid, ref device, false))
+                return;
+
+            if (!_networks.TryGetValue(device.DeviceNetId, out var deviceNet) || !deviceNet.RandomizeAddress(device.Address, address))
+            {
+                var prefix = string.IsNullOrWhiteSpace(device.Prefix) ? null : Loc.GetString(device.Prefix);
+                device.Address = $"{prefix}{_random.Next():x}";
+                device.CustomAddress = false;
+            }
         }
+        #endregion
     }
 
     /// <summary>
@@ -241,46 +294,66 @@ namespace Content.Server.DeviceNetwork.Systems
         /// <summary>
         /// The EntityUid of the entity the packet was sent from.
         /// </summary>
-        public EntityUid Sender;
+        public readonly EntityUid Sender;
 
-        public BeforePacketSentEvent(EntityUid sender)
+        public readonly TransformComponent SenderTransform;
+
+        /// <summary>
+        ///     The senders current position in world coordinates.
+        /// </summary>
+        public readonly Vector2 SenderPosition;
+
+        public BeforePacketSentEvent(EntityUid sender, TransformComponent xform, Vector2 senderPosition)
         {
             Sender = sender;
+            SenderTransform = xform;
+            SenderPosition = senderPosition;
         }
     }
 
     /// <summary>
     /// Event raised when a device network packet gets sent.
     /// </summary>
-    public sealed class PacketSentEvent : EntityEventArgs
+    public sealed class DeviceNetworkPacketEvent : EntityEventArgs
     {
+        /// <summary>
+        /// The type of network that this packet is being sent on.
+        /// </summary>
+        public ConnectionType NetId;
+
         /// <summary>
         /// The frequency the packet is sent on.
         /// </summary>
-        public int Frequency;
+        public readonly uint Frequency;
+
+        /// <summary>
+        /// Address of the intended recipient. Null if the message was broadcast.
+        /// </summary>
+        public string? Address;
 
         /// <summary>
         /// The device network address of the sending entity.
         /// </summary>
-        public string SenderAddress;
+        public readonly string SenderAddress;
 
         /// <summary>
-        /// The data that is beeing sent.
+        /// The entity that sent the packet.
         /// </summary>
-        public NetworkPayload Data;
+        public EntityUid Sender;
 
         /// <summary>
-        /// Whether the packet was broadcasted.
+        /// The data that is being sent.
         /// </summary>
-        public bool Broadcast;
+        public readonly NetworkPayload Data;
 
-        public PacketSentEvent(int frequency, string senderAddress, NetworkPayload data, bool broadcast)
+        public DeviceNetworkPacketEvent(ConnectionType netId, string? address, uint frequency, string senderAddress, EntityUid sender, NetworkPayload data)
         {
+            NetId = netId;
+            Address = address;
             Frequency = frequency;
             SenderAddress = senderAddress;
+            Sender = sender;
             Data = data;
-            Broadcast = broadcast;
         }
-
     }
 }

Content.Server/DeviceNetwork/Systems/Devices/ApcNetSwitchSystem.cs

@@ -1,4 +1,4 @@
-using Content.Server.DeviceNetwork.Components;
+using Content.Server.DeviceNetwork.Components;
 using Content.Server.DeviceNetwork.Components.Devices;
 using Content.Shared.Interaction;
 using Robust.Shared.GameObjects;
@@ -15,7 +15,7 @@ namespace Content.Server.DeviceNetwork.Systems.Devices
             base.Initialize();
 
             SubscribeLocalEvent<ApcNetSwitchComponent, InteractHandEvent>(OnInteracted);
-            SubscribeLocalEvent<ApcNetSwitchComponent, PacketSentEvent>(OnPackedReceived);
+            SubscribeLocalEvent<ApcNetSwitchComponent, DeviceNetworkPacketEvent>(OnPackedReceived);
         }
 
         /// <summary>
@@ -28,13 +28,16 @@ namespace Content.Server.DeviceNetwork.Systems.Devices
 
             component.State = !component.State;
 
+            if (networkComponent.TransmitFrequency == null)
+                return;
+
             var payload = new NetworkPayload
             {
                 [DeviceNetworkConstants.Command] = DeviceNetworkConstants.CmdSetState,
                 [DeviceNetworkConstants.StateEnabled] = component.State,
             };
 
-            _deviceNetworkSystem.QueuePacket(uid, DeviceNetworkConstants.NullAddress, networkComponent.Frequency, payload, true);
+            _deviceNetworkSystem.QueuePacket(uid, null, payload, device: networkComponent);
 
             args.Handled = true;
         }
@@ -42,7 +45,7 @@ namespace Content.Server.DeviceNetwork.Systems.Devices
         /// <summary>
         /// Listens to the <see cref="DeviceNetworkConstants.CmdSetState"/> command of other switches to sync state
         /// </summary>
-        private void OnPackedReceived(EntityUid uid, ApcNetSwitchComponent component, PacketSentEvent args)
+        private void OnPackedReceived(EntityUid uid, ApcNetSwitchComponent component, DeviceNetworkPacketEvent args)
         {
             if (!EntityManager.TryGetComponent(uid, out DeviceNetworkComponent? networkComponent) || args.SenderAddress == networkComponent.Address) return;
             if (!args.Data.TryGetValue(DeviceNetworkConstants.Command, out string? command) || command != DeviceNetworkConstants.CmdSetState) return;

Content.Server/DeviceNetwork/Systems/WiredNetworkSystem.cs

@@ -18,7 +18,7 @@ namespace Content.Server.DeviceNetwork.Systems
         /// </summary>
         private void OnBeforePacketSent(EntityUid uid, WiredNetworkComponent component, BeforePacketSentEvent args)
         {
-            if (EntityManager.GetComponent<TransformComponent>(uid).GridID != EntityManager.GetComponent<TransformComponent>(args.Sender).GridID)
+            if (EntityManager.GetComponent<TransformComponent>(uid).GridID != args.SenderTransform.GridID)
             {
                 args.Cancel();
             }

Content.Server/DeviceNetwork/Systems/WirelessNetworkSystem.cs

@@ -18,11 +18,12 @@ namespace Content.Server.DeviceNetwork.Systems
         /// </summary>
         private void OnBeforePacketSent(EntityUid uid, WirelessNetworkComponent component, BeforePacketSentEvent args)
         {
-            var ownPosition = EntityManager.GetComponent<TransformComponent>(component.Owner).WorldPosition;
-            var position = EntityManager.GetComponent<TransformComponent>(args.Sender).WorldPosition;
-            var distance = (ownPosition - position).Length;
+            var ownPosition = args.SenderPosition;
+            var xform = Transform(uid);
 
-            if (EntityManager.TryGetComponent<WirelessNetworkComponent?>(args.Sender, out var sendingComponent) && distance > sendingComponent.Range)
+            if (xform.MapID != args.SenderTransform.MapID
+                || !TryComp<WirelessNetworkComponent?>(args.Sender, out var sendingComponent)
+                || (ownPosition - xform.WorldPosition).Length > sendingComponent.Range)
             {
                 args.Cancel();
             }