Allow partial reads from a GorillaConn; add SetDeadline (from net.Conn)

The current implementation of GorillaConn will drop data if the
websocket frame isn't read 100%. For example, if a websocket frame is
size=3, and Read() is called with a []byte of len=1, the 2 other bytes
in the frame are lost forever.

This is currently masked by the fact that this is used primarily in
io.Copy to another socket (in ingress.Stream) - as long as the read buffer
used by io.Copy is big enough (it is 32*1024, so in theory we could see
this today?) then data is copied over to the other socket.

The client then can do partial reads just fine as the kernel will take
care of the buffer from here on out.

I hit this by trying to create my own tunnel and avoiding
ingress.Stream, but this could be a real bug today I think if a
websocket frame bigger than 32*1024 was received, although it is also
possible that we are lucky and the upstream size which I haven't checked
uses a smaller buffer than that always.

The test I added hangs before my change, succeeds after.

Also add SetDeadline so that GorillaConn fully implements net.Conn

websocket/connection.go

@@ -1,7 +1,9 @@
 package websocket
 
 import (
+	"bytes"
 	"context"
+	"fmt"
 	"io"
 	"time"
 
@@ -28,16 +30,24 @@ const (
 type GorillaConn struct {
 	*websocket.Conn
 	log     *zerolog.Logger
+	readBuf bytes.Buffer
 }
 
 // Read will read messages from the websocket connection
 func (c *GorillaConn) Read(p []byte) (int, error) {
+	// Intermediate buffer may contain unread bytes from the last read, start there before blocking on a new frame
+	if c.readBuf.Len() > 0 {
+		return c.readBuf.Read(p)
+	}
+
 	_, message, err := c.Conn.ReadMessage()
 	if err != nil {
 		return 0, err
 	}
-	return copy(p, message), nil
-
+	// Write into an intermediate buffer in case the websocket frame contains more data then len(p)
+	// Write returns nil error always and grows the buffer everything is always written or panic
+	c.readBuf.Write(message)
+	return c.readBuf.Read(p)
 }
 
 // Write will write messages to the websocket connection
@@ -49,6 +59,19 @@ func (c *GorillaConn) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+// SetDeadline sets both read and write deadlines, as per net.Conn interface docs:
+// "It is equivalent to calling both SetReadDeadline and SetWriteDeadline."
+// Note there is no synchronization here, but the gorilla implementation isn't thread safe anyway
+func (c *GorillaConn) SetDeadline(t time.Time) error {
+	if err := c.Conn.SetReadDeadline(t); err != nil {
+		return fmt.Errorf("error setting read deadline: %w", err)
+	}
+	if err := c.Conn.SetWriteDeadline(t); err != nil {
+		return fmt.Errorf("error setting write deadline: %w", err)
+	}
+	return nil
+}
+
 // pinger simulates the websocket connection to keep it alive
 func (c *GorillaConn) pinger(ctx context.Context) {
 	ticker := time.NewTicker(pingPeriod)

websocket/websocket.go

@@ -114,7 +114,7 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 	_ = conn.SetReadDeadline(time.Now().Add(pongWait))
 	conn.SetPongHandler(func(string) error { _ = conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
-	gorillaConn := &GorillaConn{conn, h.log}
+	gorillaConn := &GorillaConn{Conn: conn, log: h.log}
 	go gorillaConn.pinger(r.Context())
 	defer conn.Close()
 

websocket/websocket_test.go

@@ -1,18 +1,23 @@
 package websocket
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
-	"github.com/rs/zerolog"
+	"fmt"
 	"io"
 	"math/rand"
 	"net/http"
 	"testing"
+	"time"
+
+	"github.com/rs/zerolog"
 
 	"github.com/cloudflare/cloudflared/hello"
 	"github.com/cloudflare/cloudflared/tlsconfig"
 	gws "github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/net/websocket"
 )
 
@@ -102,6 +107,51 @@ func TestServe(t *testing.T) {
 	<-errC
 }
 
+func TestWebsocketWrapper(t *testing.T) {
+
+	listener, err := hello.CreateTLSListener("localhost:0")
+	require.NoError(t, err)
+
+	serverErrorChan := make(chan error)
+	helloSvrCtx, cancelHelloSvr := context.WithCancel(context.Background())
+	defer func() { <-serverErrorChan }()
+	defer cancelHelloSvr()
+	go func() {
+		log := zerolog.Nop()
+		serverErrorChan <- hello.StartHelloWorldServer(&log, listener, helloSvrCtx.Done())
+	}()
+
+	tlsConfig := websocketClientTLSConfig(t)
+	d := gws.Dialer{TLSClientConfig: tlsConfig, HandshakeTimeout: time.Minute}
+	testAddr := fmt.Sprintf("https://%s/ws", listener.Addr().String())
+	req := testRequest(t, testAddr, nil)
+	conn, resp, err := ClientConnect(req, &d)
+	require.NoError(t, err)
+	require.Equal(t, testSecWebsocketAccept, resp.Header.Get("Sec-WebSocket-Accept"))
+
+	// Websocket now connected to test server so lets check our wrapper
+	wrapper := GorillaConn{Conn: conn}
+	buf := make([]byte, 100)
+	wrapper.Write([]byte("abc"))
+	n, err := wrapper.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, n, 3)
+	require.Equal(t, "abc", string(buf[:n]))
+
+	// Test partial read, read 1 of 3 bytes in one read and the other 2 in another read
+	wrapper.Write([]byte("abc"))
+	buf = buf[:1]
+	n, err = wrapper.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, n, 1)
+	require.Equal(t, "a", string(buf[:n]))
+	buf = buf[:cap(buf)]
+	n, err = wrapper.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, n, 2)
+	require.Equal(t, "bc", string(buf[:n]))
+}
+
 // func TestStartProxyServer(t *testing.T) {
 // 	var wg sync.WaitGroup
 // 	remoteAddress := "localhost:1113"