CatStation/Content.Server/DeviceNetwork/Systems/StationLimitedNetworkSystem.cs

using Content.Server.DeviceNetwork.Components;
using Content.Server.Station.Systems;
using Content.Shared.DeviceNetwork.Events;
using JetBrains.Annotations;

namespace Content.Server.DeviceNetwork.Systems
{
    /// <summary>
    /// This system requires the StationLimitedNetworkComponent to be on the the sending entity as well as the receiving entity
    /// </summary>
    [UsedImplicitly]
    public sealed partial class StationLimitedNetworkSystem : EntitySystem
    {
        [Dependency] private StationSystem _stationSystem = default!;
        public override void Initialize()
        {
            base.Initialize();
            SubscribeLocalEvent<StationLimitedNetworkComponent, MapInitEvent>(OnMapInit);
            SubscribeLocalEvent<StationLimitedNetworkComponent, BeforePacketSentEvent>(OnBeforePacketSent);
        }

        /// <summary>
        /// Sets the station id the device is limited to.
        /// </summary>
        public void SetStation(EntityUid uid, EntityUid? stationId, StationLimitedNetworkComponent? component = null)
        {
            if (!Resolve(uid, ref component))
                return;

            component.StationId = stationId;
        }

        /// <summary>
        /// Tries to set the station id to the current station if the device is currently on a station
        /// </summary>
        public bool TrySetStationId(EntityUid uid, StationLimitedNetworkComponent? component = null)
        {
            if (!Resolve(uid, ref component) || !Transform(uid).GridUid.HasValue)
                return false;

            component.StationId = _stationSystem.GetOwningStation(uid);
            return component.StationId.HasValue;
        }

        /// <summary>
        /// Set the station id to the one the entity is on when the station limited component is added
        /// </summary>
        private void OnMapInit(EntityUid uid, StationLimitedNetworkComponent networkComponent, MapInitEvent args)
        {
            networkComponent.StationId = _stationSystem.GetOwningStation(uid);
        }

        /// <summary>
        /// Checks if both devices are limited to the same station
        /// </summary>
        private void OnBeforePacketSent(Entity<StationLimitedNetworkComponent> ent, ref BeforePacketSentEvent args)
        {
            var (uid, component) = ent;
            if (!component.StationId.HasValue)
                TrySetStationId(uid, component);

            if (!CheckStationId(args.Sender, component.AllowNonStationPackets, component.StationId))
            {
                args.Cancelled = true;
            }
        }

        /// <summary>
        /// Compares the station IDs of the sending and receiving network components.
        /// Returns false if either of them doesn't have a station ID or if their station ID isn't equal.
        /// Returns true even when the sending entity isn't tied to a station if `allowNonStationPackets` is set to true.
        /// </summary>
        private bool CheckStationId(EntityUid senderUid, bool allowNonStationPackets, EntityUid? receiverStationId, StationLimitedNetworkComponent? sender = null)
        {
            if (!receiverStationId.HasValue)
                return false;

            if (!Resolve(senderUid, ref sender, false))
                return allowNonStationPackets;

            if (!sender.StationId.HasValue)
                TrySetStationId(senderUid, sender);

            return sender.StationId == receiverStationId;
        }
    }
}