insight/src/Agent/Insight.Agent/Network/AgentSession.cs

using Insight.Agent.Interfaces;
using Insight.Agent.Messages;
using Microsoft.Extensions.Logging;
using Vaitr.Network;

namespace Insight.Agent.Network
{
    public class AgentSession : TcpSession<IAgentMessage>
    {
        private readonly IEnumerable<IAgentMessageHandler<AgentSession>> _handlers;

        public AgentSession(IEnumerable<IAgentMessageHandler<AgentSession>> handlers, ISerializer<IAgentMessage> serializer, ILogger<AgentSession> logger) : base(serializer, logger)
        {
            _handlers = handlers;
        }

        protected override ValueTask OnConnectedAsync(CancellationToken cancellationToken)
        {
            _logger.LogInformation("Agent ({ep?}) connected", RemoteEndPoint);
            return default;
        }

        protected override ValueTask OnDisconnectedAsync(CancellationToken cancellationToken)
        {
            _logger.LogInformation("Agent ({ep?}) disconnected", RemoteEndPoint);
            return default;
        }

        protected override ValueTask OnSentAsync(IPacketContext<IAgentMessage> context, CancellationToken cancellationToken)
        {
            return base.OnSentAsync(context, cancellationToken);
        }

        protected override async ValueTask OnReceivedAsync(IPacketContext<IAgentMessage> context, CancellationToken cancellationToken)
        {
            await base.OnReceivedAsync(context, cancellationToken);

            foreach (var handler in _handlers)
            {
                try
                {
                    await handler.HandleAsync(this, context.Packet, cancellationToken);
                }
                catch (Exception ex)
                {
                    _logger.LogWarning("Agent ({ep?}) {ex}", RemoteEndPoint, ex.ToString());
                }
            }
        }

        protected override ValueTask OnHeartbeatAsync(CancellationToken cancellationToken)
        {
            _logger.LogInformation("Agent ({ep?}) Heartbeat", RemoteEndPoint);
            return default;
        }
    }
}