sharp-chat/SharpChatCommon/ConnectionsContext.cs

using System;
using System.Collections.Generic;
using System.Linq;

namespace SharpChat.SockChat {
    public class ConnectionsContext {
        public static readonly TimeSpan TimeOut = TimeSpan.FromMinutes(5);

        private readonly HashSet<ConnectionInfo> Connections = new();
        private readonly HashSet<ConnectionInfo> AuthedConnections = new();
        private readonly Dictionary<long, HashSet<ConnectionInfo>> UserConnections = new();

        public ConnectionInfo[] All => Connections.ToArray();
        public ConnectionInfo[] Authed => AuthedConnections.ToArray();

        public void WithAll(Action<ConnectionInfo> body) {
            foreach(ConnectionInfo conn in Connections)
                body(conn);
        }

        public void WithAuthed(Action<ConnectionInfo> body) {
            foreach(ConnectionInfo conn in AuthedConnections)
                body(conn);
        }

        public ConnectionInfo[] GetTimedOut() {
            List<ConnectionInfo> conns = new();

            foreach(ConnectionInfo conn in Connections)
                if(DateTimeOffset.UtcNow - conn.LastPing > TimeOut)
                    conns.Add(conn);

            return conns.ToArray();
        }

        public int GetCountForUser(long userId) {
            return UserConnections.ContainsKey(userId) ? UserConnections.Count : 0;
        }

        public int GetCountForUser(UserInfo userInfo) {
            return GetCountForUser(userInfo.UserId);
        }

        public bool HasUser(long userId) {
            return GetCountForUser(userId) > 0;
        }

        public bool HasUser(UserInfo userInfo) {
            return HasUser(userInfo.UserId);
        }

        public ConnectionInfo[] GetUser(long userId) {
            if(!UserConnections.ContainsKey(userId))
                return Array.Empty<ConnectionInfo>();

            return UserConnections[userId].ToArray();
        }

        public ConnectionInfo[] GetUser(UserInfo userInfo) {
            return GetUser(userInfo.UserId);
        }

        public void WithUser(long userId, Action<ConnectionInfo> body) {
            if(!UserConnections.ContainsKey(userId))
                return;

            foreach(ConnectionInfo conn in UserConnections[userId])
                body(conn);
        }

        public void WithUser(UserInfo userInfo, Action<ConnectionInfo> body) {
            WithUser(userInfo.UserId, body);
        }

        public string[] GetAllRemoteAddresses() {
            HashSet<string> addrs = new();

            foreach(ConnectionInfo conn in Connections)
                addrs.Add(conn.RemoteAddress);

            return addrs.ToArray();
        }

        public string[] GetAuthedRemoteAddresses() {
            HashSet<string> addrs = new();

            foreach(ConnectionInfo conn in AuthedConnections)
                addrs.Add(conn.RemoteAddress);

            return addrs.ToArray();
        }

        public string[] GetUserRemoteAddresses(long userId) {
            if(!UserConnections.ContainsKey(userId))
                return Array.Empty<string>();

            HashSet<string> addrs = new();

            foreach(ConnectionInfo conn in UserConnections[userId])
                addrs.Add(conn.RemoteAddress);

            return addrs.ToArray();
        }

        public string[] GetUserRemoteAddresses(UserInfo userInfo) {
            return GetUserRemoteAddresses(userInfo.UserId);
        }

        public void Add(ConnectionInfo conn) {
            if(Connections.Contains(conn))
                return;

            Connections.Add(conn);

            if(conn.UserId > 0) {
                AuthedConnections.Add(conn);

                if(UserConnections.ContainsKey(conn.UserId))
                    UserConnections[conn.UserId].Add(conn);
                else
                    UserConnections.Add(conn.UserId, new() { conn });
            }
        }

        public void Remove(ConnectionInfo conn) {
            if(Connections.Contains(conn))
                Connections.Remove(conn);

            if(AuthedConnections.Contains(conn))
                AuthedConnections.Remove(conn);

            if(conn.UserId > 0 && UserConnections.ContainsKey(conn.UserId)) {
                UserConnections[conn.UserId].Remove(conn);

                if(UserConnections[conn.UserId].Count < 1)
                    UserConnections.Remove(conn.UserId);
            }
        }

        public void SetUser(ConnectionInfo conn, long userId) {
            if(!Connections.Contains(conn))
                return;

            // so yeah this is implemented but SetUserId will throw an exception if we're trying to re-auth a connection
            // will just leave this here but i'm not sure how to go forward with this
            if(conn.UserId > 0) {
                if(UserConnections[conn.UserId].Contains(conn))
                    UserConnections[conn.UserId].Remove(conn);

                if(UserConnections[conn.UserId].Count < 1)
                    UserConnections.Remove(conn.UserId);
            }

            conn.SetUserId(userId);

            if(conn.UserId > 0) {
                if(UserConnections.ContainsKey(conn.UserId))
                    UserConnections[conn.UserId].Add(conn);
                else
                    UserConnections.Add(conn.UserId, new() { conn });
            }

            if(conn.UserId > 0) {
                if(!AuthedConnections.Contains(conn))
                    AuthedConnections.Add(conn);
            } else {
                if(AuthedConnections.Contains(conn))
                    AuthedConnections.Remove(conn);
            }
        }

        public void SetUser(ConnectionInfo conn, UserInfo userInfo) {
            SetUser(conn, userInfo.UserId);
        }
    }
}