From 12cbf54e585c3e8a1f93e29183ebb5508d2fc981 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Thu, 15 Aug 2024 20:26:19 +0000
Subject: [PATCH 1/2] Implement http.ResponseWriter unwrapping like
 http.ResponseController

Since we rely on the connection not being hijacked too early (i.e.
detecting the presence of http.Hijacker) to set headers, we must
manually implement the unwrapping of the http.ResponseController. By
doing so, we also retain Go 1.19 compatibility without build tags.

Closes #455
---
 accept.go |  2 +-
 hijack.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 hijack.go

diff --git a/accept.go b/accept.go
index f672a730..c27a721e 100644
--- a/accept.go
+++ b/accept.go
@@ -105,7 +105,7 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con
 		}
 	}
 
-	hj, ok := w.(http.Hijacker)
+	hj, ok := hijacker(w)
 	if !ok {
 		err = errors.New("http.ResponseWriter does not implement http.Hijacker")
 		http.Error(w, http.StatusText(http.StatusNotImplemented), http.StatusNotImplemented)
diff --git a/hijack.go b/hijack.go
new file mode 100644
index 00000000..9cce45ca
--- /dev/null
+++ b/hijack.go
@@ -0,0 +1,33 @@
+//go:build !js
+
+package websocket
+
+import (
+	"net/http"
+)
+
+type rwUnwrapper interface {
+	Unwrap() http.ResponseWriter
+}
+
+// hijacker returns the Hijacker interface of the http.ResponseWriter.
+// It follows the Unwrap method of the http.ResponseWriter if available,
+// matching the behavior of http.ResponseController. If the Hijacker
+// interface is not found, it returns false.
+//
+// Since the http.ResponseController is not available in Go 1.19, and
+// does not support checking the presence of the Hijacker interface,
+// this function is used to provide a consistent way to check for the
+// Hijacker interface across Go versions.
+func hijacker(rw http.ResponseWriter) (http.Hijacker, bool) {
+	for {
+		switch t := rw.(type) {
+		case http.Hijacker:
+			return t, true
+		case rwUnwrapper:
+			rw = t.Unwrap()
+		default:
+			return nil, false
+		}
+	}
+}

From ee7b9af6e3430cd1f84c47882b7ae517cba08439 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 21 Aug 2024 10:07:31 +0000
Subject: [PATCH 2/2] Add tests

---
 accept_test.go       | 38 ++++++++++++++++++++++++++++++++++++++
 hijack_go120_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)
 create mode 100644 hijack_go120_test.go

diff --git a/accept_test.go b/accept_test.go
index 4f799126..3b45ac5c 100644
--- a/accept_test.go
+++ b/accept_test.go
@@ -143,6 +143,33 @@ func TestAccept(t *testing.T) {
 		_, err := Accept(w, r, nil)
 		assert.Contains(t, err, `failed to hijack connection`)
 	})
+
+	t.Run("wrapperHijackerIsUnwrapped", func(t *testing.T) {
+		t.Parallel()
+
+		rr := httptest.NewRecorder()
+		w := mockUnwrapper{
+			ResponseWriter: rr,
+			unwrap: func() http.ResponseWriter {
+				return mockHijacker{
+					ResponseWriter: rr,
+					hijack: func() (conn net.Conn, writer *bufio.ReadWriter, err error) {
+						return nil, nil, errors.New("haha")
+					},
+				}
+			},
+		}
+
+		r := httptest.NewRequest("GET", "/", nil)
+		r.Header.Set("Connection", "Upgrade")
+		r.Header.Set("Upgrade", "websocket")
+		r.Header.Set("Sec-WebSocket-Version", "13")
+		r.Header.Set("Sec-WebSocket-Key", xrand.Base64(16))
+
+		_, err := Accept(w, r, nil)
+		assert.Contains(t, err, "failed to hijack connection")
+	})
+
 	t.Run("closeRace", func(t *testing.T) {
 		t.Parallel()
 
@@ -534,3 +561,14 @@ var _ http.Hijacker = mockHijacker{}
 func (mj mockHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	return mj.hijack()
 }
+
+type mockUnwrapper struct {
+	http.ResponseWriter
+	unwrap func() http.ResponseWriter
+}
+
+var _ rwUnwrapper = mockUnwrapper{}
+
+func (mu mockUnwrapper) Unwrap() http.ResponseWriter {
+	return mu.unwrap()
+}
diff --git a/hijack_go120_test.go b/hijack_go120_test.go
new file mode 100644
index 00000000..0f0673a9
--- /dev/null
+++ b/hijack_go120_test.go
@@ -0,0 +1,38 @@
+//go:build !js && go1.20
+
+package websocket
+
+import (
+	"bufio"
+	"errors"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/coder/websocket/internal/test/assert"
+)
+
+func Test_hijackerHTTPResponseControllerCompatibility(t *testing.T) {
+	t.Parallel()
+
+	rr := httptest.NewRecorder()
+	w := mockUnwrapper{
+		ResponseWriter: rr,
+		unwrap: func() http.ResponseWriter {
+			return mockHijacker{
+				ResponseWriter: rr,
+				hijack: func() (conn net.Conn, writer *bufio.ReadWriter, err error) {
+					return nil, nil, errors.New("haha")
+				},
+			}
+		},
+	}
+
+	_, _, err := http.NewResponseController(w).Hijack()
+	assert.Contains(t, err, "haha")
+	hj, ok := hijacker(w)
+	assert.Equal(t, "hijacker found", ok, true)
+	_, _, err = hj.Hijack()
+	assert.Contains(t, err, "haha")
+}