diff --git a/accept.go b/accept.go index 25c1be7d..1d6a442b 100644 --- a/accept.go +++ b/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,13 +118,14 @@ 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 { @@ -120,7 +133,7 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con 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 } } @@ -128,7 +141,7 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con 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) diff --git a/accept_test.go b/accept_test.go index 92dbfcc7..83b3f2d0 100644 --- a/accept_test.go +++ b/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() diff --git a/ws_js.go b/ws_js.go index 026b75fc..b39ca4b5 100644 --- a/ws_js.go +++ b/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.