diff --git a/close.go b/close.go index fcc68065..36ef0a63 100644 --- a/close.go +++ b/close.go @@ -92,11 +92,29 @@ func CloseStatus(err error) StatusCode { // The connection can only be closed once. Additional calls to Close // are no-ops. // -// The maximum length of reason must be 125 bytes. Avoid sending a dynamic reason. +// Close replaces invalid UTF-8 in reason with U+FFFD, then truncates it to fit +// 123 bytes without splitting a rune. If it truncated the reason and the close +// had no other error, Close returns an error wrapping ErrCloseReasonTruncated. +// +// Valid codes are 1000 to 1003, 1007 to 1014, and 3000 to 4999. Close sends +// any other code as StatusInternalError with no reason and returns an error. +// The exception is StatusNoStatusRcvd, which sends a close frame with no code +// and no reason. +// +// When the first call to Close returns, the connection is closed, even on error. // // Close will unblock all goroutines interacting with the connection once // complete. -func (c *Conn) Close(code StatusCode, reason string) (err error) { +func (c *Conn) Close(code StatusCode, reason string) error { + ce, wireErr := wireCloseError(code, reason) + err := c.exportedClose(ce.Code, ce.Reason) + if err != nil { + return err + } + return wireErr +} + +func (c *Conn) exportedClose(code StatusCode, reason string) (err error) { defer errd.Wrap(&err, "failed to close WebSocket") if c.casClosing() { @@ -168,24 +186,14 @@ func (c *Conn) closeHandshake(code StatusCode, reason string) error { } func (c *Conn) writeClose(code StatusCode, reason string) error { - ce := CloseError{ - Code: code, - Reason: reason, - } - - var p []byte - var err error - if ce.Code != StatusNoStatusRcvd { - p, err = ce.bytes() - if err != nil { - return err - } - } + // The error is dropped: Close reports it from its own wireCloseError call, + // and writeError and the reply to a peer's close frame do not report it. + ce, _ := wireCloseError(code, reason) ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) defer cancel() - err = c.writeControl(ctx, opClose, p) + err := c.writeControl(ctx, opClose, ce.bytes()) // If the connection closed as we're writing we ignore the error as we might // have written the close frame, the peer responded and then someone else read it // and closed the connection. @@ -292,33 +300,29 @@ func validWireCloseCode(code StatusCode) bool { return false } -func (ce CloseError) bytes() ([]byte, error) { - p, err := ce.bytesErr() - if err != nil { - err = fmt.Errorf("failed to marshal close frame: %w", err) - ce = CloseError{ - Code: StatusInternalError, - } - p, _ = ce.bytesErr() +// wireCloseError returns the CloseError to send for code and reason, also when +// the error is non-nil. The error reports a replaced code or a truncated reason. +func wireCloseError(code StatusCode, reason string) (CloseError, error) { + if code == StatusNoStatusRcvd { + return CloseError{Code: code}, nil } - return p, err -} - -const maxCloseReason = maxControlPayload - 2 - -func (ce CloseError) bytesErr() ([]byte, error) { - if len(ce.Reason) > maxCloseReason { - return nil, fmt.Errorf("reason string max is %v but got %q with length %v", maxCloseReason, ce.Reason, len(ce.Reason)) + if !validWireCloseCode(code) { + return CloseError{Code: StatusInternalError}, fmt.Errorf("websocket: invalid close status code %v, sent %v instead", code, StatusInternalError) } + reason, err := wireCloseReason(reason) + return CloseError{Code: code, Reason: reason}, err +} - if !validWireCloseCode(ce.Code) { - return nil, fmt.Errorf("status code %v cannot be set", ce.Code) +// bytes returns the close frame payload for ce, which must come from +// wireCloseError. +func (ce CloseError) bytes() []byte { + if ce.Code == StatusNoStatusRcvd { + return nil } - buf := make([]byte, 2+len(ce.Reason)) binary.BigEndian.PutUint16(buf, uint16(ce.Code)) copy(buf[2:], ce.Reason) - return buf, nil + return buf } func (c *Conn) casClosing() bool { diff --git a/close_shared.go b/close_shared.go new file mode 100644 index 00000000..8b0b4c1e --- /dev/null +++ b/close_shared.go @@ -0,0 +1,33 @@ +package websocket + +import ( + "fmt" + "strings" + "unicode/utf8" +) + +// maxControlPayload is the maximum length of a control frame payload. +// See https://tools.ietf.org/html/rfc6455#section-5.5. +const maxControlPayload = 125 + +// maxCloseReason is maxControlPayload minus the 2-byte status code. +const maxCloseReason = maxControlPayload - 2 + +// wireCloseReason returns a reason that can always be sent, and an error if it +// truncated reason. +func wireCloseReason(reason string) (string, error) { + // RFC 6455 section 5.5.1 requires a UTF-8 reason. + valid := strings.ToValidUTF8(reason, "\uFFFD") + if len(valid) <= maxCloseReason { + return valid, nil + } + n := maxCloseReason + for !utf8.RuneStart(valid[n]) { + n-- + } + var replaced string + if len(valid) != len(reason) { + replaced = " after replacing invalid UTF-8" + } + return valid[:n], fmt.Errorf("%w from %d to %d bytes%s", ErrCloseReasonTruncated, len(valid), n, replaced) +} diff --git a/close_test.go b/close_test.go index 1e04807e..2808c252 100644 --- a/close_test.go +++ b/close_test.go @@ -7,6 +7,7 @@ import ( "math" "strings" "testing" + "unicode/utf8" "github.com/coder/websocket/internal/test/assert" ) @@ -17,6 +18,7 @@ func TestCloseError(t *testing.T) { testCases := []struct { name string ce CloseError + want CloseError success bool }{ { @@ -25,6 +27,10 @@ func TestCloseError(t *testing.T) { Code: StatusNormalClosure, Reason: strings.Repeat("x", maxCloseReason), }, + want: CloseError{ + Code: StatusNormalClosure, + Reason: strings.Repeat("x", maxCloseReason), + }, success: true, }, { @@ -33,6 +39,10 @@ func TestCloseError(t *testing.T) { Code: StatusNormalClosure, Reason: strings.Repeat("x", maxCloseReason+1), }, + want: CloseError{ + Code: StatusNormalClosure, + Reason: strings.Repeat("x", maxCloseReason), + }, success: false, }, { @@ -41,8 +51,18 @@ func TestCloseError(t *testing.T) { Code: math.MaxUint16, Reason: strings.Repeat("x", maxCloseReason), }, + want: CloseError{Code: StatusInternalError}, success: false, }, + { + name: "noStatus", + ce: CloseError{ + Code: StatusNoStatusRcvd, + Reason: "dropped", + }, + want: CloseError{Code: StatusNoStatusRcvd}, + success: true, + }, } for _, tc := range testCases { @@ -50,12 +70,13 @@ func TestCloseError(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() - _, err := tc.ce.bytesErr() + ce, err := wireCloseError(tc.ce.Code, tc.ce.Reason) if tc.success { assert.Success(t, err) } else { assert.Error(t, err) } + assert.Equal(t, "wire close error", tc.want, ce) }) } @@ -69,6 +90,52 @@ func TestCloseError(t *testing.T) { }) } +func Test_wireCloseReason(t *testing.T) { + t.Parallel() + + ascii := strings.Repeat("x", maxCloseReason) + testCases := []struct { + name string + reason string + want string + wantErrSuffix string + }{ + {name: "empty", reason: "", want: ""}, + {name: "max", reason: ascii, want: ascii}, + {name: "maxPlusOne", reason: ascii + "x", want: ascii, wantErrSuffix: "from 124 to 123 bytes"}, + {name: "twoByteFits", reason: ascii[:121] + "é", want: ascii[:121] + "é"}, + {name: "twoByteSplit", reason: ascii[:122] + "é", want: ascii[:122], wantErrSuffix: "from 124 to 122 bytes"}, + {name: "threeByteAligned", reason: strings.Repeat("€", 42), want: strings.Repeat("€", 41), wantErrSuffix: "from 126 to 123 bytes"}, + {name: "threeByteSplit", reason: ascii[:122] + "€", want: ascii[:122], wantErrSuffix: "from 125 to 122 bytes"}, + {name: "fourByteSplit", reason: strings.Repeat("😀", 31), want: strings.Repeat("😀", 30), wantErrSuffix: "from 124 to 120 bytes"}, + {name: "replacementCharSplit", reason: ascii[:122] + "\uFFFD", want: ascii[:122], wantErrSuffix: "from 125 to 122 bytes"}, + {name: "invalidByte", reason: "a\x80b", want: "a\uFFFDb"}, + {name: "invalidRun", reason: "a\x80\x80\xffb", want: "a\uFFFDb"}, + {name: "incompleteRune", reason: "a\xe2\x82", want: "a\uFFFD"}, + {name: "invalidRunOverMax", reason: strings.Repeat("\x80", 200), want: "\uFFFD"}, + {name: "growthFromMax", reason: ascii[:122] + "\x80", want: ascii[:122], wantErrSuffix: "from 125 to 122 bytes after replacing invalid UTF-8"}, + {name: "growthIsolated", reason: strings.Repeat("\x80x", 61), want: strings.Repeat("\uFFFDx", 30) + "\uFFFD", wantErrSuffix: "from 244 to 123 bytes after replacing invalid UTF-8"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, err := wireCloseReason(tc.reason) + assert.Equal(t, "reason", tc.want, got) + if tc.wantErrSuffix == "" { + assert.Success(t, err) + } else { + assert.ErrorIs(t, ErrCloseReasonTruncated, err) + assert.Equal(t, "error", "websocket: close reason truncated "+tc.wantErrSuffix, err.Error()) + } + if !utf8.ValidString(got) || len(got) > maxCloseReason { + t.Fatalf("reason %q is not valid UTF-8 within %d bytes", got, maxCloseReason) + } + }) + } +} + func Test_parseClosePayload(t *testing.T) { t.Parallel() diff --git a/conn_test.go b/conn_test.go index 01af8cd4..c955e86f 100644 --- a/conn_test.go +++ b/conn_test.go @@ -64,10 +64,104 @@ func TestConn(t *testing.T) { t.Run("badClose", func(t *testing.T) { tt, c1, c2 := newConnTest(t, nil, nil) - c2.CloseRead(tt.ctx) + readErr := xsync.Go(func() error { + _, _, err := c2.Read(tt.ctx) + return err + }) err := c1.Close(-1, "") - assert.Contains(t, err, "failed to marshal close frame: status code StatusCode(-1) cannot be set") + assert.Equal(t, "close error", "websocket: invalid close status code StatusCode(-1), sent StatusInternalError instead", err.Error()) + + assert.Success(t, assertCloseStatus(websocket.StatusInternalError, <-readErr)) + }) + + t.Run("noStatusClose", func(t *testing.T) { + tt, c1, c2 := newConnTest(t, nil, nil) + + readErr := xsync.Go(func() error { + _, _, err := c2.Read(tt.ctx) + return err + }) + + assert.Success(t, c1.Close(websocket.StatusNoStatusRcvd, "dropped")) + peerErr := <-readErr + ce, ok := errors.AsType[websocket.CloseError](peerErr) + if !ok { + t.Fatalf("peer did not receive a close frame: %v", peerErr) + } + assert.Equal(t, "peer close error", websocket.CloseError{Code: websocket.StatusNoStatusRcvd}, ce) + }) + + t.Run("closeReason", func(t *testing.T) { + t.Parallel() + + const code websocket.StatusCode = 4000 + ascii := strings.Repeat("x", 123) + testCases := []struct { + name string + reason string + want string + truncated bool + }{ + {name: "max", reason: ascii, want: ascii}, + {name: "overMax", reason: ascii + "xyz", want: ascii, truncated: true}, + {name: "twoByteSplit", reason: ascii[:122] + "é", want: ascii[:122], truncated: true}, + {name: "fourByteSplit", reason: strings.Repeat("😀", 40), want: strings.Repeat("😀", 30), truncated: true}, + {name: "invalidUTF8", reason: "bad \xff\xfe reason", want: "bad \uFFFD reason"}, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + tt, c1, c2 := newConnTest(t, nil, nil) + + readErr := xsync.Go(func() error { + _, _, err := c2.Read(tt.ctx) + return err + }) + + err := c1.Close(code, tc.reason) + if !tc.truncated { + assert.Success(t, err) + } else { + assert.ErrorIs(t, websocket.ErrCloseReasonTruncated, err) + assert.Equal(t, "close error", fmt.Sprintf("websocket: close reason truncated from %d to %d bytes", len(tc.reason), len(tc.want)), err.Error()) + } + + peerErr := <-readErr + ce, ok := errors.AsType[websocket.CloseError](peerErr) + if !ok { + t.Fatalf("peer did not receive a close frame: %v", peerErr) + } + assert.Equal(t, "peer close error", websocket.CloseError{Code: code, Reason: tc.want}, ce) + }) + } + }) + + t.Run("closeReasonTruncatedWithWriteError", func(t *testing.T) { + _, c1, c2 := newConnTest(t, nil, nil) + + assert.Success(t, c2.CloseNow()) + + err := c1.Close(websocket.StatusNormalClosure, strings.Repeat("x", 124)) + assert.ErrorIs(t, io.ErrClosedPipe, err) + assert.Equal(t, "reports truncation", false, errors.Is(err, websocket.ErrCloseReasonTruncated)) + }) + + t.Run("writeErrorReasonTruncated", func(t *testing.T) { + tt, c1, c2 := newConnTest(t, nil, nil) + + _ = c1.CloseRead(tt.ctx) + readErr := xsync.Go(func() error { + _, _, err := c2.Read(tt.ctx) + return err + }) + + c1.WriteError(websocket.StatusProtocolError, errors.New(strings.Repeat("x", 124))) + peerErr := <-readErr + ce, ok := errors.AsType[websocket.CloseError](peerErr) + if !ok { + t.Fatalf("peer did not receive a close frame: %v", peerErr) + } + assert.Equal(t, "peer close error", websocket.CloseError{Code: websocket.StatusProtocolError, Reason: strings.Repeat("x", 123)}, ce) }) t.Run("ping", func(t *testing.T) { diff --git a/errors.go b/errors.go index bf4fc2b0..71f36e67 100644 --- a/errors.go +++ b/errors.go @@ -6,3 +6,7 @@ import ( // ErrMessageTooBig is returned when a message exceeds the read limit. var ErrMessageTooBig = errors.New("websocket: message too big") + +// ErrCloseReasonTruncated indicates that Close shortened the reason to fit +// 123 bytes after replacing invalid UTF-8. +var ErrCloseReasonTruncated = errors.New("websocket: close reason truncated") diff --git a/export_test.go b/export_test.go index e0071922..d211e2ac 100644 --- a/export_test.go +++ b/export_test.go @@ -37,3 +37,7 @@ var ( ) var CompressionModeOpts = CompressionMode.opts + +func (c *Conn) WriteError(code StatusCode, err error) { + c.writeError(code, err) +} diff --git a/frame.go b/frame.go index e7ab76be..a6c8a10e 100644 --- a/frame.go +++ b/frame.go @@ -101,10 +101,6 @@ func readFrameHeader(r *bufio.Reader, readBuf []byte) (h header, err error) { return h, nil } -// maxControlPayload is the maximum length of a control frame payload. -// See https://tools.ietf.org/html/rfc6455#section-5.5. -const maxControlPayload = 125 - // writeFrameHeader writes the bytes of the header to w. // See https://tools.ietf.org/html/rfc6455#section-5.2 func writeFrameHeader(h header, w *bufio.Writer, buf []byte) (err error) { diff --git a/ws_js.go b/ws_js.go index 026b75fc..d8bed4ba 100644 --- a/ws_js.go +++ b/ws_js.go @@ -228,12 +228,18 @@ func (c *Conn) write(typ MessageType, p []byte) error { // It will wait until the peer responds with a close frame // or the connection is closed. // It thus performs the full WebSocket close handshake. +// +// Close replaces invalid UTF-8 in reason with U+FFFD, then truncates it to fit +// 123 bytes without splitting a rune. If it truncated the reason and the close +// had no other error, Close returns an error wrapping ErrCloseReasonTruncated. func (c *Conn) Close(code StatusCode, reason string) error { + // Browsers throw on a reason over 123 bytes and leave the connection open. + reason, truncErr := wireCloseReason(reason) err := c.exportedClose(code, reason) if err != nil { return fmt.Errorf("failed to close WebSocket: %w", err) } - return nil + return truncErr } // CloseNow closes the WebSocket connection without attempting a close handshake. diff --git a/ws_js_test.go b/ws_js_test.go index 1fa242f4..a1d846fa 100644 --- a/ws_js_test.go +++ b/ws_js_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "os" + "strings" "testing" "time" @@ -37,6 +38,40 @@ func TestWasm(t *testing.T) { assert.Success(t, err) } +func TestWasmCloseReasonTruncated(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + c, _, err := websocket.Dial(ctx, os.Getenv("WS_ECHO_SERVER_URL"), &websocket.DialOptions{ + Subprotocols: []string{"echo"}, + }) + assert.Success(t, err) + + // The echo server checks the close code. + err = c.Close(websocket.StatusNormalClosure, strings.Repeat("é", 62)) + assert.ErrorIs(t, websocket.ErrCloseReasonTruncated, err) + assert.Equal(t, "close error", "websocket: close reason truncated from 124 to 122 bytes", err.Error()) +} + +func TestWasmCloseReasonInvalidUTF8(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + c, _, err := websocket.Dial(ctx, os.Getenv("WS_ECHO_SERVER_URL"), &websocket.DialOptions{ + Subprotocols: []string{"echo"}, + }) + assert.Success(t, err) + + // Without the U+FFFD replacement, these 50 bytes become more than 123 bytes + // when converted to a JavaScript string. The echo server checks the close code. + err = c.Close(websocket.StatusNormalClosure, strings.Repeat("\x80", 50)) + assert.Success(t, err) +} + func TestWasmDialTimeout(t *testing.T) { t.Parallel()