fix: ensure websocket close messages are truncated to 123 bytes

coadler · coadler · commit 947271714476 · 2022-03-31T13:44:23.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/cli/configssh_test.go b/cli/configssh_test.go
@@ -4,10 +4,11 @@ import (
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/pty/ptytest"
-	"github.com/stretchr/testify/require"
 )
 
 func TestConfigSSH(t *testing.T) {
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -13,11 +13,12 @@ import (
 	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/xerrors"
 
+	"golang.org/x/crypto/ssh/terminal"
+
 	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/codersdk"
-	"golang.org/x/crypto/ssh/terminal"
 )
 
 func ssh() *cobra.Command {
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 {
+	builder := strings.Builder{}
+	builder.Grow(byteLen)
+
+	for _, char := range str {
+		if builder.Len()+len(string(char)) > byteLen {
+			break
+		}
+
+		_, _ = builder.WriteRune(char)
+	}
+
+	return builder.String()
+}
diff --git a/coderd/ws_test.go b/coderd/ws_test.go
@@ -0,0 +1,30 @@
+// This file tests an internal function.
+//nolint:testpackage
+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)
+	})
+}
diff --git a/provisioner/terraform/serve.go b/provisioner/terraform/serve.go
@@ -10,6 +10,7 @@ import (
 	"cdr.dev/slog"
 
 	"github.com/cli/safeexec"
+
 	"github.com/coder/coder/provisionersdk"
 
 	"github.com/hashicorp/hc-install/product"

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`	`}`