move to coderd.go

coadler · coadler · commit 51f2225701dc · 2022-03-31T13:42:59.000-05:00
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -1,8 +1,10 @@
 package coderd
 
 import (
+	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
 	"sync"
 	"time"
 
@@ -197,3 +199,37 @@ type api struct {
 	websocketWaitMutex sync.Mutex
 	websocketWaitGroup sync.WaitGroup
 }
+
+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 {
+		// 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.
+		truncateString := func(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()
+		}
+
+		return truncateString(msg, websocketCloseMaxLen)
+	}
+
+	return msg
+}
diff --git a/coderd/coderd_test.go b/coderd/coderd_test.go
@@ -1,11 +1,34 @@
 package coderd_test
 
 import (
+	"strings"
 	"testing"
 
+	"github.com/coder/coder/coderd"
+	"github.com/stretchr/testify/assert"
 	"go.uber.org/goleak"
 )
 
 func TestMain(m *testing.M) {
 	goleak.VerifyTestMain(m)
 }
+
+func TestFmtWebsocketCloseMsg(t *testing.T) {
+	t.Parallel()
+
+	t.Run("TruncateSingleByteCharacters", func(t *testing.T) {
+		t.Parallel()
+
+		msg := strings.Repeat("d", 255)
+		trunc := coderd.FmtWebsocketCloseMsg(msg)
+		assert.LessOrEqual(t, len(trunc), 123)
+	})
+
+	t.Run("TruncateMultiByteCharacters", func(t *testing.T) {
+		t.Parallel()
+
+		msg := strings.Repeat("こんにちは", 10)
+		trunc := coderd.FmtWebsocketCloseMsg(msg)
+		assert.LessOrEqual(t, len(trunc), 123)
+	})
+}
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, fmtWebsocketCloseMsg("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, fmtWebsocketCloseMsg("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, fmtWebsocketCloseMsg("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, fmtWebsocketCloseMsg("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, fmtWebsocketCloseMsg("serve: %s", err))
+		_ = conn.Close(websocket.StatusInternalError, FmtWebsocketCloseMsg("serve: %s", err))
 		return
 	}
 }
diff --git a/coderd/ws.go b/coderd/ws.go
diff --git a/coderd/ws_test.go b/coderd/ws_test.go

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, fmtWebsocketCloseMsg("serve: %s", err))`
	`111`	`+ _ = conn.Close(websocket.StatusInternalError, FmtWebsocketCloseMsg("serve: %s", err))`
`112`	`112`	`return`
`113`	`113`	`}`
`114`	`114`	`}`