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.closeAndWait(ce)
+	if err != nil {
+		return err
+	}
+	return wireErr
+}
+
+func (c *Conn) closeAndWait(ce CloseError) (err error) {
 	defer errd.Wrap(&err, "failed to close WebSocket")
 
 	if c.casClosing() {
@@ -112,7 +130,7 @@ func (c *Conn) Close(code StatusCode, reason string) (err error) {
 		}
 	}()
 
-	err = c.closeHandshake(code, reason)
+	err = c.closeHandshake(ce)
 
 	err2 := c.close()
 	if err == nil && err2 != nil {
@@ -154,38 +172,24 @@ func (c *Conn) CloseNow() (err error) {
 	return err
 }
 
-func (c *Conn) closeHandshake(code StatusCode, reason string) error {
-	err := c.writeClose(code, reason)
+func (c *Conn) closeHandshake(ce CloseError) error {
+	err := c.writeClose(ce)
 	if err != nil {
 		return err
 	}
 
 	err = c.waitCloseHandshake()
-	if CloseStatus(err) != code {
+	if CloseStatus(err) != ce.Code {
 		return err
 	}
 	return nil
 }
 
-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
-		}
-	}
-
+func (c *Conn) writeClose(ce CloseError) error {
 	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.
@@ -262,16 +266,17 @@ func parseClosePayload(p []byte) (CloseError, error) {
 		return CloseError{}, fmt.Errorf("close payload %q too small, cannot even contain the 2 byte status code", p)
 	}
 
-	ce := CloseError{
-		Code:   StatusCode(binary.BigEndian.Uint16(p)),
-		Reason: string(p[2:]),
-	}
-
-	if !validWireCloseCode(ce.Code) {
-		return CloseError{}, fmt.Errorf("invalid status code %v", ce.Code)
+	code := StatusCode(binary.BigEndian.Uint16(p))
+	if !validWireCloseCode(code) {
+		return CloseError{}, fmt.Errorf("invalid close status code %v", code)
 	}
 
-	return ce, nil
+	// The reason is normalized as for Close, because the reply to this frame
+	// sends the returned CloseError. The truncation error is dropped, because a
+	// received reason fits in a close frame and only replacing invalid UTF-8
+	// can make it longer.
+	reason, _ := wireCloseReason(string(p[2:]))
+	return CloseError{Code: code, Reason: reason}, nil
 }
 
 // See http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
@@ -292,33 +297,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, 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, replaced with %v", 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 or parseClosePayload.
+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 the reason to send, also when the error is non-nil.
+// The error reports a 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()
 
@@ -77,6 +144,7 @@ func Test_parseClosePayload(t *testing.T) {
 		p       []byte
 		success bool
 		ce      CloseError
+		err     string
 	}{
 		{
 			name:    "normal",
@@ -103,6 +171,26 @@ func Test_parseClosePayload(t *testing.T) {
 			name:    "badStatusCode",
 			p:       []byte{0x17, 0x70},
 			success: false,
+			err:     "invalid close status code StatusCode(6000)",
+		},
+		{
+			name:    "noStatusCode",
+			p:       []byte{0x03, 0xED},
+			success: false,
+			err:     "invalid close status code StatusNoStatusRcvd",
+		},
+		{
+			name:    "invalidUTF8Reason",
+			p:       []byte{0x03, 0xE8, 0xFF, 'a'},
+			success: true,
+			ce:      CloseError{Code: StatusNormalClosure, Reason: "\uFFFDa"},
+		},
+		{
+			// Each "\xFFa" becomes 4 bytes, so the reason is truncated.
+			name:    "invalidUTF8ReasonTooLong",
+			p:       append([]byte{0x03, 0xE8}, strings.Repeat("\xFFa", 61)+"a"...),
+			success: true,
+			ce:      CloseError{Code: StatusNormalClosure, Reason: strings.Repeat("\uFFFDa", 30) + "\uFFFD"},
 		},
 	}
 
@@ -117,6 +205,9 @@ func Test_parseClosePayload(t *testing.T) {
 				assert.Equal(t, "close payload", tc.ce, ce)
 			} else {
 				assert.Error(t, err)
+				if tc.err != "" {
+					assert.Equal(t, "error", tc.err, err.Error())
+				}
 			}
 		})
 	}