fix: ensure websocket close messages are truncated to 123 bytes

coadler · coadler · commit ca133c82c4de · 2022-03-31T12:58:37.000-05:00
It's possible for websocket close messages to be too long, which cause
them to silently fail without a proper close message. See error below:

```
2022-03-31 17:08:34.862 [INFO]	(stdlib)	&lt;close_notjs.go:72&gt;	"2022/03/31 17:08:34 websocket: failed to marshal close frame: reason string max is 123 but got \"insert provisioner daemon:Cannot encode []database.ProvisionerType into oid 19098 - []database.ProvisionerType must implement Encoder or be converted to a string\" with length 161"
```
diff --git a/coderd/provisionerdaemons.go b/coderd/provisionerdaemons.go
@@ -56,7 +56,7 @@ func (api *api) provisionerDaemonsListen(rw http.ResponseWriter, r *http.Request
 		Provisioners: []database.ProvisionerType{database.ProvisionerTypeEcho, database.ProvisionerTypeTerraform},
 	})
 	if err != nil {
-		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("insert provisioner daemon:% s", err))
+		_ = conn.Close(websocket.StatusInternalError, fmtWebsocketCloseMsg("insert provisioner daemon: %s", err))
 		return
 	}
 
@@ -67,7 +67,7 @@ func (api *api) provisionerDaemonsListen(rw http.ResponseWriter, r *http.Request
 	config.LogOutput = io.Discard
 	session, err := yamux.Server(websocket.NetConn(r.Context(), conn, websocket.MessageBinary), config)
 	if err != nil {
-		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("multiplex server: %s", err))
+		_ = conn.Close(websocket.StatusInternalError, fmtWebsocketCloseMsg("multiplex server: %s", err))
 		return
 	}
 	mux := drpcmux.New()
@@ -80,13 +80,13 @@ func (api *api) provisionerDaemonsListen(rw http.ResponseWriter, r *http.Request
 		Logger:       api.Logger.Named(fmt.Sprintf("provisionerd-%s", daemon.Name)),
 	})
 	if err != nil {
-		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("drpc register provisioner daemon: %s", err))
+		_ = conn.Close(websocket.StatusInternalError, fmtWebsocketCloseMsg("drpc register provisioner daemon: %s", err))
 		return
 	}
 	server := drpcserver.New(mux)
 	err = server.Serve(r.Context(), session)
 	if err != nil {
-		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("serve: %s", err))
+		_ = conn.Close(websocket.StatusInternalError, fmtWebsocketCloseMsg("serve: %s", err))
 		return
 	}
 	_ = conn.Close(websocket.StatusGoingAway, "")
diff --git a/coderd/workspaceresources.go b/coderd/workspaceresources.go
@@ -108,7 +108,7 @@ func (api *api) workspaceResourceDial(rw http.ResponseWriter, r *http.Request) {
 		Pubsub:    api.Pubsub,
 	})
 	if err != nil {
-		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("serve: %s", err))
+		_ = conn.Close(websocket.StatusInternalError, fmtWebsocketCloseMsg("serve: %s", err))
 		return
 	}
 }
diff --git a/coderd/ws.go b/coderd/ws.go
@@ -0,0 +1,40 @@
+package coderd
+
+import (
+	"fmt"
+	"strings"
+)
+
+const websocketCloseMaxLen = 123
+
+// fmtWebsocketCloseMsg formats a websocket close message and ensures it is
+// truncated to the maximum allowed length.
+func fmtWebsocketCloseMsg(format string, vars ...any) string {
+	msg := fmt.Sprintf(format, vars...)
+
+	// Cap msg length at 123 bytes. nhooyr/websocket only allows close messages
+	// of this length.
+	if len(msg) > websocketCloseMaxLen {
+		return truncateString(msg, websocketCloseMaxLen)
+	}
+
+	return msg
+}
+
+// truncateString safely truncates a string to a maximum size of byteLen. It
+// writes whole runes until a single rune would increase the string size above
+// byteLen.
+func truncateString(str string, byteLen int) string {
+	sb := strings.Builder{}
+	sb.Grow(byteLen)
+
+	for _, char := range str {
+		if sb.Len()+len(string(char)) > byteLen {
+			break
+		}
+
+		sb.WriteRune(char)
+	}
+
+	return sb.String()
+}
diff --git a/coderd/ws_test.go b/coderd/ws_test.go
@@ -0,0 +1,28 @@
+package coderd
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_websocketCloseMsg(t *testing.T) {
+	t.Parallel()
+
+	t.Run("TruncateSingleByteCharacters", func(t *testing.T) {
+		t.Parallel()
+
+		msg := strings.Repeat("d", 255)
+		trunc := fmtWebsocketCloseMsg(msg)
+		assert.LessOrEqual(t, len(trunc), 123)
+	})
+
+	t.Run("TruncateMultiByteCharacters", func(t *testing.T) {
+		t.Parallel()
+
+		msg := strings.Repeat("こんにちは", 10)
+		trunc := fmtWebsocketCloseMsg(msg)
+		assert.LessOrEqual(t, len(trunc), 123)
+	})
+}

Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ func (api api) workspaceResourceDial(rw http.ResponseWriter, r http.Request) {`
`108`	`108`	`Pubsub: api.Pubsub,`
`109`	`109`	`})`
`110`	`110`	`if err != nil {`
`111`		`- _ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("serve: %s", err))`
	`111`	`+ _ = conn.Close(websocket.StatusInternalError, fmtWebsocketCloseMsg("serve: %s", err))`
`112`	`112`	`return`
`113`	`113`	`}`
`114`	`114`	`}`