Added support for WebSockets

pkgs/shelf_proxy/lib/shelf_proxy.dart

@@ -7,6 +7,9 @@ import 'dart:async';
 import 'package:http/http.dart' as http;
 import 'package:path/path.dart' as p;
 import 'package:shelf/shelf.dart';
+import 'package:shelf_web_socket/shelf_web_socket.dart';
+import 'package:web_socket_channel/io.dart';
+import 'package:web_socket_channel/web_socket_channel.dart';
 
 /// A handler that proxies requests to [url].
 ///
@@ -35,11 +38,42 @@ Handler proxyHandler(Object url, {http.Client? client, String? proxyName}) {
   proxyName ??= 'shelf_proxy';
 
   return (serverRequest) async {
-    // TODO(nweiz): Support WebSocket requests.
+    final requestUrl = uri.resolve(serverRequest.url.toString());
+
+    if (serverRequest.headers['Upgrade'] == 'websocket') {
+      final isSecure =
+          requestUrl.isScheme('https') || requestUrl.isScheme('wss');
+
+      final wsRequestUrl = Uri(
+        scheme: isSecure ? 'wss' : 'ws',
+        userInfo: requestUrl.userInfo,
+        host: requestUrl.host,
+        port: requestUrl.port,
+        path: requestUrl.path,
+        query: requestUrl.query,
+      );
+
+      final handler = webSocketHandler((WebSocketChannel serverChannel) {
+        final headers = Map<String, String>.from(serverRequest.headers);
+
+        // Add a Via header. See
+        // http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.45
+        _addHeader(
+            headers, 'via', '${serverRequest.protocolVersion} $proxyName');
+
+        final clientChannel = IOWebSocketChannel.connect(
+          wsRequestUrl,
+          headers: headers,
+        );
+        clientChannel.stream.pipe(serverChannel.sink);
+        serverChannel.stream.pipe(clientChannel.sink);
+      });
+
+      return handler(serverRequest);
+    }
 
     // TODO(nweiz): Handle TRACE requests correctly. See
     // http://www.w3.org/Protocols/rfc2616/rfc2616-sec9.html#sec9.8
-    final requestUrl = uri.resolve(serverRequest.url.toString());
     final clientRequest = http.StreamedRequest(serverRequest.method, requestUrl)
       ..followRedirects = false
       ..headers.addAll(serverRequest.headers)

pkgs/shelf_proxy/pubspec.yaml

@@ -14,6 +14,8 @@ dependencies:
   http: '>=0.13.0 <2.0.0'
   path: ^1.8.0
   shelf: ^1.0.0
+  shelf_web_socket: ^1.0.4
+  web_socket_channel: ^2.4.0
 
 dev_dependencies:
   dart_flutter_team_lints: ^2.0.0

pkgs/shelf_proxy/test/shelf_proxy_test.dart

@@ -9,7 +9,9 @@ import 'package:http/testing.dart';
 import 'package:shelf/shelf.dart' as shelf;
 import 'package:shelf/shelf_io.dart' as shelf_io;
 import 'package:shelf_proxy/shelf_proxy.dart';
+import 'package:shelf_web_socket/shelf_web_socket.dart';
 import 'package:test/test.dart';
+import 'package:web_socket_channel/web_socket_channel.dart';
 
 /// The URI of the server the current proxy server is proxying to.
 late Uri targetUri;
@@ -167,6 +169,28 @@ void main() {
     expect(response.headers,
         containsPair('warning', '214 shelf_proxy "GZIP decoded"'));
   });
+
+  test('proxy websockets echo', () async {
+    await createWebSocketServer((webSocket) {
+      webSocket.stream.listen((message) {
+        webSocket.sink.add('echo $message');
+      });
+    });
+
+    final client = WebSocketChannel.connect(proxyUri);
+
+    final completer = Completer<String>();
+
+    client.sink.add('hello');
+
+    client.stream.listen((event) {
+      completer.complete(event as String);
+    });
+
+    final message = await completer.future;
+
+    expect(message, 'echo hello');
+  });
 }
 
 /// Creates a proxy server proxying to a server running [handler].
@@ -206,3 +230,25 @@ Future<http.Response> get({Map<String, String>? headers}) {
   request.followRedirects = false;
   return request.send().then(http.Response.fromStream);
 }
+
+Future<void> createWebSocketServer(
+    void Function(WebSocketChannel webSocket) listener,
+    {String? targetPath}) async {
+  final wshandler = webSocketHandler(listener);
+
+  final handler = expectAsync1(wshandler, reason: 'target server handler');
+
+  final targetServer = await shelf_io.serve(handler, 'localhost', 0);
+  targetUri = Uri.parse('ws://localhost:${targetServer.port}');
+  if (targetPath != null) targetUri = targetUri.resolve(targetPath);
+  final proxyServerHandler =
+      expectAsync1(proxyHandler(targetUri), reason: 'proxy server handler');
+
+  final proxyServer = await shelf_io.serve(proxyServerHandler, 'localhost', 0);
+  proxyUri = Uri.parse('ws://localhost:${proxyServer.port}');
+
+  addTearDown(() {
+    proxyServer.close(force: true);
+    targetServer.close(force: true);
+  });
+}