Add WebSockets telemetry (#1237)

* Add WebSockets telemetry

* Simplify test

* Add IClock, don't count control frames

* Add WebSocketsParser unit tests

* Only inspect WebSocket streams

* Test control frames with payloads

Author: MihaZupan

samples/ReverseProxy.Metrics.Sample/Startup.cs

@@ -44,6 +44,8 @@ public void ConfigureServices(IServiceCollection services)
             // Registration of a consumer to events for HttpClient telemetry
             // Note: this depends on changes implemented in .NET 5
             services.AddTelemetryConsumer<HttpClientTelemetryConsumer>();
+
+            services.AddTelemetryConsumer<WebSocketsTelemetryConsumer>();
         }
 
         /// <summary>
@@ -55,6 +57,9 @@ public void Configure(IApplicationBuilder app)
             // Placed at the beginning so it is the first and last thing run for each request
             app.UsePerRequestMetricCollection();
 
+            // Middleware used to intercept the WebSocket connection and collect telemetry exposed to WebSocketsTelemetryConsumer
+            app.UseWebSocketsTelemetry();
+
             app.UseRouting();
             app.UseEndpoints(endpoints =>
             {

samples/ReverseProxy.Metrics.Sample/WebSocketsTelemetryConsumer.cs

@@ -0,0 +1,21 @@
+using System;
+using Microsoft.Extensions.Logging;
+using Yarp.Telemetry.Consumption;
+
+namespace Yarp.Sample
+{
+    public sealed class WebSocketsTelemetryConsumer : IWebSocketsTelemetryConsumer
+    {
+        private readonly ILogger<WebSocketsTelemetryConsumer> _logger;
+
+        public WebSocketsTelemetryConsumer(ILogger<WebSocketsTelemetryConsumer> logger)
+        {
+            _logger = logger ?? throw new ArgumentNullException(nameof(logger));
+        }
+
+        public void OnWebSocketClosed(DateTime timestamp, DateTime establishedTime, WebSocketCloseReason closeReason, long messagesRead, long messagesWritten)
+        {
+            _logger.LogInformation($"WebSocket connection closed ({closeReason}) after reading {messagesRead} and writing {messagesWritten} messages over {(timestamp - establishedTime).TotalSeconds:N2} seconds.");
+        }
+    }
+}

src/ReverseProxy/Forwarder/HttpForwarder.cs

@@ -585,7 +585,7 @@ private async ValueTask<ForwarderError> HandleUpgradedResponse(HttpContext conte
                 var (secondResult, secondException) = await secondTask;
                 if (secondResult != StreamCopyResult.Success)
                 {
-                    error = ReportResult(context, requestFinishedFirst, secondResult, secondException!);
+                    error = ReportResult(context, !requestFinishedFirst, secondResult, secondException!);
                 }
                 else
                 {

test/ReverseProxy.Tests/Common/DelegatingStream.cs renamed to src/ReverseProxy/Utilities/DelegatingStream.cs

@@ -1,14 +1,15 @@
-// Licensed to the .NET Foundation under one or more agreements.
-// The .NET Foundation licenses this file to you under the MIT license.
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
 
 using System;
 using System.Diagnostics;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
 
-namespace Yarp.Tests.Common
+namespace Yarp.ReverseProxy.Utilities
 {
+    // Taken from https://github.com/dotnet/runtime/blob/00f37bc13b4edbba1afca9e98d74432a94f5192f/src/libraries/Common/src/System/IO/DelegatingStream.cs
     // Forwards all calls to an inner stream except where overridden in a derived class.
     internal abstract class DelegatingStream : Stream
     {
@@ -113,9 +114,9 @@ public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken
             return _innerStream.ReadAsync(buffer, cancellationToken);
         }
 
-        public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback callback, object state)
+        public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
         {
-            return _innerStream.BeginRead(buffer, offset, count, callback, state);
+            return _innerStream.BeginRead(buffer, offset, count, callback!, state);
         }
 
         public override int EndRead(IAsyncResult asyncResult)
@@ -167,9 +168,9 @@ public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationTo
             return _innerStream.WriteAsync(buffer, cancellationToken);
         }
 
-        public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback callback, object state)
+        public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
         {
-            return _innerStream.BeginWrite(buffer, offset, count, callback, state);
+            return _innerStream.BeginWrite(buffer, offset, count, callback!, state);
         }
 
         public override void EndWrite(IAsyncResult asyncResult)

src/ReverseProxy/WebSocketsTelemetry/WebSocketCloseReason.cs

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+namespace Yarp.ReverseProxy.WebSocketsTelemetry
+{
+    internal enum WebSocketCloseReason : int
+    {
+        Unknown,
+        ClientGracefulClose,
+        ServerGracefulClose,
+        ClientDisconnect,
+        ServerDisconnect,
+    }
+}

src/ReverseProxy/WebSocketsTelemetry/WebSocketsParser.cs

@@ -0,0 +1,142 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.Diagnostics;
+using Yarp.ReverseProxy.Utilities;
+
+namespace Yarp.ReverseProxy.WebSocketsTelemetry
+{
+    internal unsafe struct WebSocketsParser
+    {
+        private const int MaskLength = 4;
+        private const int MinHeaderSize = 2;
+        private const int MaxHeaderSize = MinHeaderSize + MaskLength + sizeof(ulong);
+
+        private fixed byte _leftoverBuffer[MaxHeaderSize - 1];
+        private readonly byte _minHeaderSize;
+        private byte _leftover;
+        private ulong _bytesToSkip;
+        private long _closeTime;
+        private readonly IClock _clock;
+
+        public long MessageCount { get; private set; }
+
+        public DateTime? CloseTime => _closeTime == 0 ? null : new DateTime(_closeTime, DateTimeKind.Utc);
+
+        public WebSocketsParser(IClock clock, bool isServer)
+        {
+            _minHeaderSize = (byte)(MinHeaderSize + (isServer ? MaskLength : 0));
+            _leftover = 0;
+            _bytesToSkip = 0;
+            _closeTime = 0;
+            _clock = clock;
+            MessageCount = 0;
+        }
+
+        // The WebSocket Protocol: https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
+        //  0                   1                   2                   3
+        //  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+        // +-+-+-+-+-------+-+-------------+-------------------------------+
+        // |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
+        // |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
+        // |N|V|V|V|       |S|             |   (if payload len==126/127)   |
+        // | |1|2|3|       |K|             |                               |
+        // +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
+        // |     Extended payload length continued, if payload len == 127  |
+        // + - - - - - - - - - - - - - - - +-------------------------------+
+        // |                               |Masking-key, if MASK set to 1  |
+        // +-------------------------------+-------------------------------+
+        // | Masking-key (continued)       |          Payload Data         |
+        // +-------------------------------- - - - - - - - - - - - - - - - +
+        // :                     Payload Data continued ...                :
+        // +---------------------------------------------------------------+
+        //
+        // The header can be 2-10 bytes long, followed by a 4 byte mask if the message was sent by the client.
+        // We have to read the first 2 bytes to know how long the frame header will be.
+        // Since the buffer may not contain the full frame, we make use of a leftoverBuffer
+        // where we store leftover bytes that don't represent a complete frame header.
+        // On the next call to Consume, we interpret the leftover bytes as the beginning of the frame.
+        // As we are not interested in the actual payload data, we skip over (payload length + mask length) bytes after each header.
+        public void Consume(ReadOnlySpan<byte> buffer)
+        {
+            int leftover = _leftover;
+            var bytesToSkip = _bytesToSkip;
+
+            while (true)
+            {
+                var toSkip = Math.Min(bytesToSkip, (ulong)buffer.Length);
+                buffer = buffer.Slice((int)toSkip);
+                bytesToSkip -= toSkip;
+
+                var available = leftover + buffer.Length;
+                int headerSize = _minHeaderSize;
+
+                if (available < headerSize)
+                {
+                    break;
+                }
+
+                var length = (leftover > 1 ? _leftoverBuffer[1] : buffer[1 - leftover]) & 0x7FUL;
+
+                if (length > 125)
+                {
+                    // The actual length will be encoded in 2 or 8 bytes, based on whether the length was 126 or 127
+                    var lengthBytes = 2 << (((int)length & 1) << 1);
+                    headerSize += lengthBytes;
+                    Debug.Assert(leftover < headerSize);
+
+                    if (available < headerSize)
+                    {
+                        break;
+                    }
+
+                    lengthBytes += MinHeaderSize;
+
+                    length = 0;
+                    for (var i = MinHeaderSize; i < lengthBytes; i++)
+                    {
+                        length <<= 8;
+                        length |= i < leftover ? _leftoverBuffer[i] : buffer[i - leftover];
+                    }
+                }
+
+                Debug.Assert(leftover < headerSize);
+                bytesToSkip = length;
+
+                const int NonReservedBitsMask = 0b_1000_1111;
+                var header = (leftover > 0 ? _leftoverBuffer[0] : buffer[0]) & NonReservedBitsMask;
+
+                // Don't count control frames under MessageCount
+                if ((uint)(header - 0x80) <= 0x02)
+                {
+                    // Has FIN (0x80) and is a Continuation (0x00) / Text (0x01) / Binary (0x02) opcode
+                    MessageCount++;
+                }
+                else if ((header & 0xF) == 0x8) // CLOSE
+                {
+                    if (_closeTime == 0)
+                    {
+                        _closeTime = _clock.GetUtcNow().Ticks;
+                    }
+                }
+
+                // Advance the buffer by the number of bytes read for the header,
+                // accounting for any bytes we may have read from the leftoverBuffer
+                buffer = buffer.Slice(headerSize - leftover);
+                leftover = 0;
+            }
+
+            Debug.Assert(bytesToSkip == 0 || buffer.Length == 0);
+            _bytesToSkip = bytesToSkip;
+
+            Debug.Assert(leftover + buffer.Length < MaxHeaderSize);
+            for (var i = 0; i < buffer.Length; i++, leftover++)
+            {
+                _leftoverBuffer[leftover] = buffer[i];
+            }
+
+            _leftover = (byte)leftover;
+        }
+    }
+}

src/ReverseProxy/WebSocketsTelemetry/WebSocketsTelemetry.cs

@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Diagnostics.Tracing;
+
+namespace Yarp.ReverseProxy.WebSocketsTelemetry
+{
+    [EventSource(Name = "Yarp.ReverseProxy.WebSockets")]
+    internal sealed class WebSocketsTelemetry : EventSource
+    {
+        public static readonly WebSocketsTelemetry Log = new();
+
+        [Event(1, Level = EventLevel.Informational)]
+        public void WebSocketClosed(long establishedTime, WebSocketCloseReason closeReason, long messagesRead, long messagesWritten)
+        {
+            if (IsEnabled(EventLevel.Informational, EventKeywords.All))
+            {
+                WriteEvent(eventId: 1, establishedTime, closeReason, messagesRead, messagesWritten);
+            }
+        }
+    }
+}

src/ReverseProxy/WebSocketsTelemetry/WebSocketsTelemetryExtensions.cs

@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Linq;
+using Microsoft.Extensions.DependencyInjection;
+using Yarp.ReverseProxy.Utilities;
+using Yarp.ReverseProxy.WebSocketsTelemetry;
+
+namespace Microsoft.AspNetCore.Builder
+{
+    /// <summary>
+    /// <see cref="IApplicationBuilder"/> extension methods to add the <see cref="WebSocketsTelemetryMiddleware"/>.
+    /// </summary>
+    public static class WebSocketsTelemetryExtensions
+    {
+        /// <summary>
+        /// Adds a <see cref="WebSocketsTelemetryMiddleware"/> to the request pipeline.
+        /// Must be added before <see cref="WebSockets.WebSocketMiddleware"/>.
+        /// </summary>
+        public static IApplicationBuilder UseWebSocketsTelemetry(this IApplicationBuilder app)
+        {
+            return app.Use(next =>
+            {
+                // Avoid exposing another extension method (AddWebSocketsTelemetry) just because of IClock
+                var clock = app.ApplicationServices.GetServices<IClock>().FirstOrDefault() ?? new Clock();
+                return new WebSocketsTelemetryMiddleware(next, clock).InvokeAsync;
+            });
+        }
+    }
+}