diff --git a/accept.go b/accept.go index 25c1be7d..4a86686a 100644 --- a/accept.go +++ b/accept.go @@ -79,6 +79,17 @@ 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. err may contain server details, so ErrorHandler decides + // what to send to the client. 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 +106,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,21 +119,23 @@ 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 { + defaultBody := err.Error() if errors.Is(err, path.ErrBadPattern) { log.Printf("websocket: %v", err) - err = errors.New(http.StatusText(http.StatusForbidden)) + defaultBody = http.StatusText(http.StatusForbidden) } - http.Error(w, err.Error(), http.StatusForbidden) + opts.writeRejection(w, r, http.StatusForbidden, err, defaultBody) return nil, err } } @@ -128,7 +143,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 +196,14 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con }), nil } +func (opts *AcceptOptions) writeRejection(w http.ResponseWriter, r *http.Request, status int, err error, defaultBody string) { + if opts.ErrorHandler != nil { + opts.ErrorHandler(w, r, status, err) + return + } + http.Error(w, defaultBody, 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..746b9cf0 100644 --- a/accept_test.go +++ b/accept_test.go @@ -5,9 +5,11 @@ package websocket import ( "bufio" "errors" + "io" "net" "net/http" "net/http/httptest" + "path" "strings" "sync" "testing" @@ -207,6 +209,287 @@ 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 + // wantErrIs, if set, must be wrapped by the error Accept returns and the + // error ErrorHandler receives. + wantErrIs error + // wantHeader is the response header when the handler runs. + wantHeader http.Header + // wantDefaultBody is the body without an ErrorHandler, if it differs + // from wantErr. + wantDefaultBody string +}{ + { + 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: "invalidOrigin", + modify: func(r *http.Request) { r.Header.Set("Origin", "http://[::1") }, + wantStatus: http.StatusForbidden, + wantErr: `failed to parse Origin header "http://[::1": parse "http://[::1": missing ']' in host`, + wantHeader: http.Header{}, + }, + { + name: "noOriginHost", + modify: func(r *http.Request) { r.Header.Set("Origin", "null") }, + wantStatus: http.StatusForbidden, + wantErr: `request Origin "null" is not a valid URL with a host`, + 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: `failed to parse path pattern "[": syntax error in pattern`, + wantErrIs: path.ErrBadPattern, + wantHeader: http.Header{}, + wantDefaultBody: "Forbidden", + }, + { + name: "noHijacker", + wantStatus: http.StatusNotImplemented, + wantErr: "http.ResponseWriter does not implement http.Hijacker", + wantHeader: http.Header{}, + wantDefaultBody: "Not Implemented", + }, +} + +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 + gotReq *http.Request + gotStatus int + gotErr error + gotHeader http.Header + ) + opts := tc.opts + opts.ErrorHandler = func(w http.ResponseWriter, r *http.Request, status int, err error) { + calls++ + gotReq, gotStatus, gotErr, gotHeader = r, 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) + if gotReq != r { + t.Fatalf("handler request: expected %p but got %p", r, gotReq) + } + if tc.wantErrIs != nil { + assert.ErrorIs(t, tc.wantErrIs, err) + assert.ErrorIs(t, tc.wantErrIs, gotErr) + } + 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()) + }) + } + + t.Run("hijackFails", func(t *testing.T) { + t.Parallel() + + errHijack := errors.New("hijack error") + w := mockHijacker{ + ResponseWriter: httptest.NewRecorder(), + hijack: func() (net.Conn, *bufio.ReadWriter, error) { + return nil, nil, errHijack + }, + } + var calls int + _, err := Accept(w, newHandshakeRequest(), &AcceptOptions{ + ErrorHandler: func(http.ResponseWriter, *http.Request, int, error) { calls++ }, + }) + assert.ErrorIs(t, errHijack, err) + assert.Equal(t, "handler calls", 0, calls) + }) + + // The handler response reaches a client over a real connection. + 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() + + // A raw connection avoids http.Client transport goroutines, + // which outlive s.Close and fail the TestMain leak check. + c, err := net.Dial("tcp", s.Listener.Addr().String()) + assert.Success(t, err) + defer c.Close() + req := newHandshakeRequest() + req.Header.Set("Connection", "Upgrade, close") + req.Header.Set("Origin", "https://other.com") + assert.Success(t, req.Write(c)) + res, err := http.ReadResponse(bufio.NewReader(c), 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 + } + + for _, tc := range acceptRejectionTests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + wantBody := tc.wantErr + if tc.wantDefaultBody != "" { + wantBody = tc.wantDefaultBody + } + + r := newHandshakeRequest() + if tc.modify != nil { + tc.modify(r) + } + w := httptest.NewRecorder() + + _, err := Accept(w, r, &tc.opts) + if tc.wantErrIs != nil { + assert.ErrorIs(t, tc.wantErrIs, err) + } + 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", plain(tc.wantHeader.Clone()), w.Header()) + assert.Equal(t, "body", wantBody+"\n", 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.