TrafficMiddleware.cs 1.94 KB
using HHECS.DAQServer.Services;
using HHECS.DAQShared.Models;

namespace HHECS.DAQServer.Middlewares
{
    public class TrafficMiddleware
    {
        private readonly RequestDelegate _next;
        private readonly DataCacheService _dataCacheService;

        public TrafficMiddleware(RequestDelegate next, DataCacheService dataCacheService)
        {
            _next = next;
            _dataCacheService = dataCacheService;
        }

        public async Task InvokeAsync(HttpContext context)
        {
            _ = Guid.TryParse(context.Request.Headers["ClientId"], out var clientId);

            // 请求之前
            var requestSize = context.Request.ContentLength ?? 0;

            // 创建一个流来保存响应
            var originalBodyStream = context.Response.Body;
            using var newBodyStream = new MemoryStream();
            context.Response.Body = newBodyStream;

            // 继续处理请求
            await _next(context);

            // 响应完成后获取响应数据长度
            newBodyStream.Seek(0, SeekOrigin.Begin);
            var responseSize = newBodyStream.Length;

            // 记录流量数据
            var record = new TrafficRecord
            {
                Origin = $"{context.Request.Scheme}://{context.Request.Host.Value}",
                RequestPath = context.Request.Path,
                Method = context.Request.Method,
                RemoteIpAddress = context.Connection.RemoteIpAddress.ToString(),
                UserAgent = context.Request.Headers.UserAgent,
                RequestSize = requestSize,
                ResponseSize = responseSize,
                ClientId = clientId,
                Created = DateTime.Now,
            };
            _dataCacheService.TrafficRecords.Enqueue(record);

            // 将流复制回原始响应流
            newBodyStream.Seek(0, SeekOrigin.Begin);
            await newBodyStream.CopyToAsync(originalBodyStream);
        }
    }
}