From 236ab225a3be29049f7940d6a6610eaec92530a1 Mon Sep 17 00:00:00 2001
From: Bryan Phelps <bryan@coder.com>
Date: Wed, 2 Feb 2022 17:41:59 +0000
Subject: [PATCH 1/3] Add proxy 'Write' that clones the bytes prior to
 forwarding

---
 coderd/provisionerdaemons.go | 35 +++++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/coderd/provisionerdaemons.go b/coderd/provisionerdaemons.go
index c98c3261a9825..633ccb85778be 100644
--- a/coderd/provisionerdaemons.go
+++ b/coderd/provisionerdaemons.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"reflect"
 	"time"
@@ -46,6 +47,33 @@ func (api *api) provisionerDaemons(rw http.ResponseWriter, r *http.Request) {
 	render.JSON(rw, r, daemons)
 }
 
+type proxiedConn struct {
+	conn io.ReadWriteCloser
+}
+
+func (c *proxiedConn) Close() error {
+	return c.conn.Close()
+}
+func (c *proxiedConn) Write(p []byte) (int, error) {
+	// Copy the data to avoid a race condition
+	cpy := make([]byte, len(p))
+	copy(cpy, p)
+	return c.conn.Write(cpy)
+}
+
+func (c *proxiedConn) Read(p []byte) (int, error) {
+	// In theory, this could be deep-copied too - but it actually causes a failure.
+	//cpy := make([]byte, len(p))
+	//copy(cpy, p)
+	return c.conn.Read(p)
+}
+
+func createCopyOnReadProxy(pipesToProxy io.ReadWriteCloser) io.ReadWriteCloser {
+	return &proxiedConn{
+		conn: pipesToProxy,
+	}
+}
+
 // Serves the provisioner daemon protobuf API over a WebSocket.
 func (api *api) provisionerDaemonsServe(rw http.ResponseWriter, r *http.Request) {
 	conn, err := websocket.Accept(rw, r, nil)
@@ -67,10 +95,13 @@ func (api *api) provisionerDaemonsServe(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
+	netConn := websocket.NetConn(r.Context(), conn, websocket.MessageBinary)
+	proxiedConn := createCopyOnReadProxy(netConn)
+
 	// Multiplexes the incoming connection using yamux.
 	// This allows multiple function calls to occur over
 	// the same connection.
-	session, err := yamux.Server(websocket.NetConn(r.Context(), conn, websocket.MessageBinary), nil)
+	session, err := yamux.Server(proxiedConn, nil)
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("multiplex server: %s", err))
 		return
@@ -84,7 +115,7 @@ func (api *api) provisionerDaemonsServe(rw http.ResponseWriter, r *http.Request)
 	})
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("drpc register provisioner daemon: %s", err))
-		return
+		// return
 	}
 	server := drpcserver.New(mux)
 	err = server.Serve(r.Context(), session)

From 2a5ba89d89554daf5925fec848bf3fd2cd5fc62d Mon Sep 17 00:00:00 2001
From: Bryan Phelps <bryan@coder.com>
Date: Wed, 2 Feb 2022 17:44:55 +0000
Subject: [PATCH 2/3] Bring back return

---
 coderd/provisionerdaemons.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/coderd/provisionerdaemons.go b/coderd/provisionerdaemons.go
index 633ccb85778be..662a07a74c856 100644
--- a/coderd/provisionerdaemons.go
+++ b/coderd/provisionerdaemons.go
@@ -115,7 +115,7 @@ func (api *api) provisionerDaemonsServe(rw http.ResponseWriter, r *http.Request)
 	})
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("drpc register provisioner daemon: %s", err))
-		// return
+		return
 	}
 	server := drpcserver.New(mux)
 	err = server.Serve(r.Context(), session)

From eef4d3af9536226efb9b0a19e632bf2728a38c02 Mon Sep 17 00:00:00 2001
From: Bryan Phelps <bryan@coder.com>
Date: Wed, 2 Feb 2022 20:19:03 +0000
Subject: [PATCH 3/3] Turn off compression to avoid race

---
 coderd/provisionerdaemons.go | 38 +++++-------------------------------
 1 file changed, 5 insertions(+), 33 deletions(-)

diff --git a/coderd/provisionerdaemons.go b/coderd/provisionerdaemons.go
index 662a07a74c856..1a315402f08fc 100644
--- a/coderd/provisionerdaemons.go
+++ b/coderd/provisionerdaemons.go
@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"net/http"
 	"reflect"
 	"time"
@@ -47,36 +46,12 @@ func (api *api) provisionerDaemons(rw http.ResponseWriter, r *http.Request) {
 	render.JSON(rw, r, daemons)
 }
 
-type proxiedConn struct {
-	conn io.ReadWriteCloser
-}
-
-func (c *proxiedConn) Close() error {
-	return c.conn.Close()
-}
-func (c *proxiedConn) Write(p []byte) (int, error) {
-	// Copy the data to avoid a race condition
-	cpy := make([]byte, len(p))
-	copy(cpy, p)
-	return c.conn.Write(cpy)
-}
-
-func (c *proxiedConn) Read(p []byte) (int, error) {
-	// In theory, this could be deep-copied too - but it actually causes a failure.
-	//cpy := make([]byte, len(p))
-	//copy(cpy, p)
-	return c.conn.Read(p)
-}
-
-func createCopyOnReadProxy(pipesToProxy io.ReadWriteCloser) io.ReadWriteCloser {
-	return &proxiedConn{
-		conn: pipesToProxy,
-	}
-}
-
 // Serves the provisioner daemon protobuf API over a WebSocket.
 func (api *api) provisionerDaemonsServe(rw http.ResponseWriter, r *http.Request) {
-	conn, err := websocket.Accept(rw, r, nil)
+	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
+		// Need to disable compression to avoid a data-race
+		CompressionMode: websocket.CompressionDisabled,
+	})
 	if err != nil {
 		httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
 			Message: fmt.Sprintf("accept websocket: %s", err),
@@ -95,13 +70,10 @@ func (api *api) provisionerDaemonsServe(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	netConn := websocket.NetConn(r.Context(), conn, websocket.MessageBinary)
-	proxiedConn := createCopyOnReadProxy(netConn)
-
 	// Multiplexes the incoming connection using yamux.
 	// This allows multiple function calls to occur over
 	// the same connection.
-	session, err := yamux.Server(proxiedConn, nil)
+	session, err := yamux.Server(websocket.NetConn(r.Context(), conn, websocket.MessageBinary), nil)
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("multiplex server: %s", err))
 		return