fix: send a close frame for over-long reasons and invalid codes

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 {

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)
+}

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,21 +51,32 @@ 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 {
 		tc := tc
 		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()
 

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) {

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")

export_test.go

@@ -37,3 +37,7 @@ var (
 )
 
 var CompressionModeOpts = CompressionMode.opts
+
+func (c *Conn) WriteError(code StatusCode, err error) {
+	c.writeError(code, err)
+}