fix: missing sec-websocket-protocol response header in andUpgradeToWebSocket()

Author: Donnerbart

src/main/java/io/fabric8/mockwebserver/internal/SimpleResponse.java

@@ -28,6 +28,8 @@
 
 public class SimpleResponse implements ServerResponse {
 
+  private static final String HTTP_HEADER_SEC_WEBSOCKET_PROTOCOL = "sec-websocket-protocol";
+
   private final ResponseProvider<String> bodyProvider;
 
   private final WebSocketSession webSocketSession;
@@ -67,6 +69,14 @@ public MockResponse toMockResponse(RecordedRequest request) {
 
     if (webSocketSession != null) {
       mockResponse.withWebSocketUpgrade(webSocketSession);
+      // see https://developer.mozilla.org/en-US/docs/Web/HTTP/Protocol_upgrade_mechanism
+      // see https://github.com/netty/netty/blob/4.1/codec-http/src/main/java/io/netty/handler/codec/http/websocketx/WebSocketClientHandshaker.java#L366
+      String requestWebsocketProtocol = request.getHeaders().get(HTTP_HEADER_SEC_WEBSOCKET_PROTOCOL);
+      if (requestWebsocketProtocol != null
+        // only add the response header if it's not set, to prevent changing custom response headers
+        && mockResponse.getHeaders().get(HTTP_HEADER_SEC_WEBSOCKET_PROTOCOL) == null) {
+        mockResponse.addHeader(HTTP_HEADER_SEC_WEBSOCKET_PROTOCOL, requestWebsocketProtocol);
+      }
     } else {
       mockResponse.setBody(bodyProvider.getBody(request));
     }

src/test/groovy/io/fabric8/mockwebserver/DefaultMockServerWebSocketTest.groovy

@@ -17,6 +17,7 @@ package io.fabric8.mockwebserver
 
 import okhttp3.OkHttpClient
 import okhttp3.Request
+import okhttp3.Response
 import okhttp3.WebSocket
 import okhttp3.WebSocketListener
 import spock.lang.Shared
@@ -119,4 +120,48 @@ class DefaultMockServerWebSocketTest extends Specification {
     cleanup:
     wss.forEach(ws -> ws.close(1000, "Test finished"))
   }
+
+  // https://github.com/fabric8io/mockwebserver/issues/77
+  def "andUpgradeToWebSocket, with request header 'sec-websocket-protocol', should create response with matching header"() {
+    given:
+    server.expect()
+            .withPath("/websocket")
+            .andUpgradeToWebSocket().open().done().always()
+    def future = new CompletableFuture()
+    when:
+    def ws = client.newWebSocket(new Request.Builder().url(server.url("/websocket")).header("sec-websocket-protocol", "v4.channel.k8s.io").build(), new WebSocketListener() {
+      @Override
+      void onOpen(WebSocket webSocket, Response response) {
+        future.complete(response.header("sec-websocket-protocol"))
+      }
+    })
+    then:
+    assert future.get(100L, TimeUnit.MILLISECONDS) == "v4.channel.k8s.io"
+    cleanup:
+    ws.close(1000, "Test finished")
+  }
+
+  // https://github.com/fabric8io/mockwebserver/issues/77
+  def "andUpgradeToWebSocket, with request header 'sec-websocket-protocol', should not change existing response header"() {
+    given:
+    server.expect()
+            .withPath("/websocket")
+            .andUpgradeToWebSocket()
+            .open()
+            .done()
+            .withHeader("sec-websocket-protocol", "v3.channel.k8s.io,v4.channel.k8s.io")
+            .always()
+    def future = new CompletableFuture()
+    when:
+    def ws = client.newWebSocket(new Request.Builder().url(server.url("/websocket")).header("sec-websocket-protocol", "v4.channel.k8s.io").build(), new WebSocketListener() {
+      @Override
+      void onOpen(WebSocket webSocket, Response response) {
+        future.complete(response.header("sec-websocket-protocol"))
+      }
+    })
+    then:
+    assert future.get(100L, TimeUnit.MILLISECONDS) == "v3.channel.k8s.io,v4.channel.k8s.io"
+    cleanup:
+    ws.close(1000, "Test finished")
+  }
 }