feat: add AcceptOptions.ErrorHandler

accept.go

@@ -79,6 +79,16 @@ type AcceptOptions struct {
 	// Unlike OnPingReceived, this callback does not return a value because a pong frame
 	// is a response to a ping and does not trigger any further frame transmission.
 	OnPongReceived func(ctx context.Context, payload []byte)
+
+	// ErrorHandler writes the HTTP response when Accept rejects the handshake
+	// request. status is the HTTP status Accept chose, w.Header() holds the
+	// headers Accept set for the rejection, and Accept returns an error
+	// wrapping err. ErrorHandler must write a response; otherwise net/http
+	// sends 200 OK with an empty body. If nil, Accept writes a plain text
+	// response with http.Error. ErrorHandler is not called when hijacking the
+	// connection fails, because Accept has already written 101 Switching
+	// Protocols.
+	ErrorHandler func(w http.ResponseWriter, r *http.Request, status int, err error)
 }
 
 func (opts *AcceptOptions) cloneWithDefaults() *AcceptOptions {
@@ -95,7 +105,9 @@ func (opts *AcceptOptions) cloneWithDefaults() *AcceptOptions {
 // Accept will not allow cross origin requests by default.
 // See the InsecureSkipVerify and OriginPatterns options to allow cross origin requests.
 //
-// Accept will write a response to w on all errors.
+// Accept writes a response to w on all errors, through
+// AcceptOptions.ErrorHandler when it rejects the handshake request. Do not
+// write to w after Accept returns an error.
 //
 // Note that using the http.Request Context after Accept returns may lead to
 // unexpected behavior (see http.Hijacker).
@@ -106,29 +118,30 @@ func Accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (*Conn,
 func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Conn, err error) {
 	defer errd.Wrap(&err, "failed to accept WebSocket connection")
 
+	opts = opts.cloneWithDefaults()
+
 	errCode, err := verifyClientRequest(w, r)
 	if err != nil {
-		http.Error(w, err.Error(), errCode)
+		opts.writeRejection(w, r, errCode, err, err.Error())
 		return nil, err
 	}
 
-	opts = opts.cloneWithDefaults()
 	if !opts.InsecureSkipVerify {
 		err = authenticateOrigin(r, opts.OriginPatterns)
 		if err != nil {
 			if errors.Is(err, path.ErrBadPattern) {
 				log.Printf("websocket: %v", err)
 				err = errors.New(http.StatusText(http.StatusForbidden))
 			}
-			http.Error(w, err.Error(), http.StatusForbidden)
+			opts.writeRejection(w, r, http.StatusForbidden, err, err.Error())
 			return nil, err
 		}
 	}
 
 	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)
+		opts.writeRejection(w, r, http.StatusNotImplemented, err, http.StatusText(http.StatusNotImplemented))
 		return nil, err
 	}
 
@@ -181,6 +194,16 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con
 	}), nil
 }
 
+// writeRejection writes the response for a handshake request that Accept
+// rejects. body is the http.Error text used when ErrorHandler is nil.
+func (opts *AcceptOptions) writeRejection(w http.ResponseWriter, r *http.Request, status int, err error, body string) {
+	if opts.ErrorHandler != nil {
+		opts.ErrorHandler(w, r, status, err)
+		return
+	}
+	http.Error(w, body, status)
+}
+
 func verifyClientRequest(w http.ResponseWriter, r *http.Request) (errCode int, _ error) {
 	if !r.ProtoAtLeast(1, 1) {
 		return http.StatusUpgradeRequired, fmt.Errorf("WebSocket protocol violation: handshake request must be at least HTTP/1.1: %q", r.Proto)

accept_test.go

@@ -5,6 +5,7 @@ package websocket
 import (
 	"bufio"
 	"errors"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -207,6 +208,242 @@ func TestAccept(t *testing.T) {
 	})
 }
 
+func newHandshakeRequest() *http.Request {
+	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", "dGhlIHNhbXBsZSBub25jZQ==")
+	return r
+}
+
+// acceptRejectionTests covers every rejection before the 101 response. The
+// http.ResponseWriter is an httptest.ResponseRecorder, which does not
+// implement http.Hijacker.
+var acceptRejectionTests = []struct {
+	name       string
+	modify     func(r *http.Request)
+	opts       AcceptOptions
+	wantStatus int
+	wantErr    string
+	// wantHeader is the response header when the handler runs.
+	wantHeader http.Header
+}{
+	{
+		name: "http1.0",
+		modify: func(r *http.Request) {
+			r.Proto, r.ProtoMajor, r.ProtoMinor = "HTTP/1.0", 1, 0
+		},
+		wantStatus: http.StatusUpgradeRequired,
+		wantErr:    `WebSocket protocol violation: handshake request must be at least HTTP/1.1: "HTTP/1.0"`,
+		wantHeader: http.Header{},
+	},
+	{
+		name:       "noConnectionUpgrade",
+		modify:     func(r *http.Request) { r.Header.Set("Connection", "keep-alive") },
+		wantStatus: http.StatusUpgradeRequired,
+		wantErr:    `WebSocket protocol violation: Connection header "keep-alive" does not contain Upgrade`,
+		wantHeader: http.Header{"Connection": {"Upgrade"}, "Upgrade": {"websocket"}},
+	},
+	{
+		name:       "noUpgradeWebSocket",
+		modify:     func(r *http.Request) { r.Header.Set("Upgrade", "h2c") },
+		wantStatus: http.StatusUpgradeRequired,
+		wantErr:    `WebSocket protocol violation: Upgrade header "h2c" does not contain websocket`,
+		wantHeader: http.Header{"Connection": {"Upgrade"}, "Upgrade": {"websocket"}},
+	},
+	{
+		name:       "methodNotGet",
+		modify:     func(r *http.Request) { r.Method = "POST" },
+		wantStatus: http.StatusMethodNotAllowed,
+		wantErr:    `WebSocket protocol violation: handshake request method is not GET but "POST"`,
+		wantHeader: http.Header{},
+	},
+	{
+		name:       "unsupportedVersion",
+		modify:     func(r *http.Request) { r.Header.Set("Sec-WebSocket-Version", "8") },
+		wantStatus: http.StatusBadRequest,
+		wantErr:    `unsupported WebSocket protocol version (only 13 is supported): "8"`,
+		wantHeader: http.Header{"Sec-Websocket-Version": {"13"}},
+	},
+	{
+		name:       "missingKey",
+		modify:     func(r *http.Request) { r.Header.Del("Sec-WebSocket-Key") },
+		wantStatus: http.StatusBadRequest,
+		wantErr:    "WebSocket protocol violation: missing Sec-WebSocket-Key",
+		wantHeader: http.Header{},
+	},
+	{
+		name:       "multipleKeys",
+		modify:     func(r *http.Request) { r.Header.Add("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") },
+		wantStatus: http.StatusBadRequest,
+		wantErr:    "WebSocket protocol violation: multiple Sec-WebSocket-Key headers",
+		wantHeader: http.Header{},
+	},
+	{
+		name:       "invalidKey",
+		modify:     func(r *http.Request) { r.Header.Set("Sec-WebSocket-Key", "notbase64") },
+		wantStatus: http.StatusBadRequest,
+		wantErr:    `WebSocket protocol violation: invalid Sec-WebSocket-Key "notbase64", must be a 16 byte base64 encoded string`,
+		wantHeader: http.Header{},
+	},
+	{
+		name:       "unauthorizedOrigin",
+		modify:     func(r *http.Request) { r.Header.Set("Origin", "https://other.com") },
+		wantStatus: http.StatusForbidden,
+		wantErr:    `request Origin "other.com" is not authorized for Host "example.com"`,
+		wantHeader: http.Header{},
+	},
+	{
+		// The pattern error is logged and not sent to the client.
+		name:       "badOriginPattern",
+		modify:     func(r *http.Request) { r.Header.Set("Origin", "https://other.com") },
+		opts:       AcceptOptions{OriginPatterns: []string{"["}},
+		wantStatus: http.StatusForbidden,
+		wantErr:    "Forbidden",
+		wantHeader: http.Header{},
+	},
+	{
+		name:       "noHijacker",
+		wantStatus: http.StatusNotImplemented,
+		wantErr:    "http.ResponseWriter does not implement http.Hijacker",
+		wantHeader: http.Header{},
+	},
+}
+
+func TestAcceptErrorHandler(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range acceptRejectionTests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			r := newHandshakeRequest()
+			if tc.modify != nil {
+				tc.modify(r)
+			}
+			w := httptest.NewRecorder()
+
+			var (
+				calls     int
+				gotStatus int
+				gotErr    error
+				gotHeader http.Header
+			)
+			opts := tc.opts
+			opts.ErrorHandler = func(w http.ResponseWriter, _ *http.Request, status int, err error) {
+				calls++
+				gotStatus, gotErr, gotHeader = status, err, w.Header().Clone()
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(status)
+				_, _ = w.Write([]byte(`{"message":"custom"}`))
+			}
+
+			_, err := Accept(w, r, &opts)
+			assert.Equal(t, "handler calls", 1, calls)
+			assert.Equal(t, "handler status", tc.wantStatus, gotStatus)
+			assert.Equal(t, "handler err", tc.wantErr, gotErr.Error())
+			assert.ErrorIs(t, gotErr, err)
+			assert.Equal(t, "header when handler runs", tc.wantHeader, gotHeader)
+
+			assert.Equal(t, "response status", tc.wantStatus, w.Code)
+			assert.Equal(t, "response content type", "application/json", w.Header().Get("Content-Type"))
+			assert.Equal(t, "response body", `{"message":"custom"}`, w.Body.String())
+		})
+	}
+
+	// The handler response reaches a client over a real connection, and a
+	// handler that writes nothing leaves net/http's default 200 OK.
+	t.Run("server", func(t *testing.T) {
+		t.Parallel()
+
+		for _, tc := range []struct {
+			name       string
+			write      bool
+			wantStatus int
+			wantBody   string
+		}{
+			{name: "writes", write: true, wantStatus: http.StatusForbidden, wantBody: `{"message":"custom"}`},
+			{name: "writesNothing", write: false, wantStatus: http.StatusOK, wantBody: ""},
+		} {
+			t.Run(tc.name, func(t *testing.T) {
+				t.Parallel()
+
+				s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					_, _ = Accept(w, r, &AcceptOptions{
+						ErrorHandler: func(w http.ResponseWriter, _ *http.Request, status int, _ error) {
+							if tc.write {
+								w.WriteHeader(status)
+								_, _ = w.Write([]byte(`{"message":"custom"}`))
+							}
+						},
+					})
+				}))
+				defer s.Close()
+
+				req, err := http.NewRequest("GET", s.URL, nil)
+				assert.Success(t, err)
+				req.Header = newHandshakeRequest().Header
+				req.Header.Set("Origin", "https://other.com")
+				res, err := s.Client().Do(req)
+				assert.Success(t, err)
+				defer res.Body.Close()
+				body, err := io.ReadAll(res.Body)
+				assert.Success(t, err)
+
+				assert.Equal(t, "status", tc.wantStatus, res.StatusCode)
+				assert.Equal(t, "body", tc.wantBody, string(body))
+			})
+		}
+	})
+}
+
+// TestAcceptRejectionDefault pins the response Accept writes without an
+// ErrorHandler.
+func TestAcceptRejectionDefault(t *testing.T) {
+	t.Parallel()
+
+	plain := func(h http.Header) http.Header {
+		h.Set("Content-Type", "text/plain; charset=utf-8")
+		h.Set("X-Content-Type-Options", "nosniff")
+		return h
+	}
+	want := map[string]struct {
+		header http.Header
+		body   string
+	}{
+		"http1.0":             {plain(http.Header{}), `WebSocket protocol violation: handshake request must be at least HTTP/1.1: "HTTP/1.0"` + "\n"},
+		"noConnectionUpgrade": {plain(http.Header{"Connection": {"Upgrade"}, "Upgrade": {"websocket"}}), `WebSocket protocol violation: Connection header "keep-alive" does not contain Upgrade` + "\n"},
+		"noUpgradeWebSocket":  {plain(http.Header{"Connection": {"Upgrade"}, "Upgrade": {"websocket"}}), `WebSocket protocol violation: Upgrade header "h2c" does not contain websocket` + "\n"},
+		"methodNotGet":        {plain(http.Header{}), `WebSocket protocol violation: handshake request method is not GET but "POST"` + "\n"},
+		"unsupportedVersion":  {plain(http.Header{"Sec-Websocket-Version": {"13"}}), `unsupported WebSocket protocol version (only 13 is supported): "8"` + "\n"},
+		"missingKey":          {plain(http.Header{}), "WebSocket protocol violation: missing Sec-WebSocket-Key\n"},
+		"multipleKeys":        {plain(http.Header{}), "WebSocket protocol violation: multiple Sec-WebSocket-Key headers\n"},
+		"invalidKey":          {plain(http.Header{}), `WebSocket protocol violation: invalid Sec-WebSocket-Key "notbase64", must be a 16 byte base64 encoded string` + "\n"},
+		"unauthorizedOrigin":  {plain(http.Header{}), `request Origin "other.com" is not authorized for Host "example.com"` + "\n"},
+		"badOriginPattern":    {plain(http.Header{}), "Forbidden\n"},
+		"noHijacker":          {plain(http.Header{}), "Not Implemented\n"},
+	}
+
+	for _, tc := range acceptRejectionTests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			r := newHandshakeRequest()
+			if tc.modify != nil {
+				tc.modify(r)
+			}
+			w := httptest.NewRecorder()
+
+			_, err := Accept(w, r, &tc.opts)
+			assert.Equal(t, "err", "failed to accept WebSocket connection: "+tc.wantErr, err.Error())
+			assert.Equal(t, "status", tc.wantStatus, w.Code)
+			assert.Equal(t, "header", want[tc.name].header, w.Header())
+			assert.Equal(t, "body", want[tc.name].body, w.Body.String())
+		})
+	}
+}
+
 func Test_verifyClientHandshake(t *testing.T) {
 	t.Parallel()
 

ws_js.go

@@ -439,6 +439,7 @@ type AcceptOptions struct {
 	OriginPatterns       []string
 	CompressionMode      CompressionMode
 	CompressionThreshold int
+	ErrorHandler         func(w http.ResponseWriter, r *http.Request, status int, err error)
 }
 
 // Accept is stubbed out for Wasm.