wasm: Limit read buffer

nhooyr · nhooyr · commit 3edd28bcf09f · 2019-09-23T17:50:33.000-05:00
diff --git a/conn.go b/conn.go
@@ -120,12 +120,6 @@ func (c *Conn) Subprotocol() string {
 	return c.subprotocol
 }
 
-func (c *Conn) setCloseErr(err error) {
-	c.closeErrOnce.Do(func() {
-		c.closeErr = fmt.Errorf("websocket closed: %w", err)
-	})
-}
-
 func (c *Conn) close(err error) {
 	c.closeOnce.Do(func() {
 		runtime.SetFinalizer(c, nil)
diff --git a/conn_common.go b/conn_common.go
@@ -202,3 +202,9 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
 func (c *Conn) SetReadLimit(n int64) {
 	c.msgReadLimit = n
 }
+
+func (c *Conn) setCloseErr(err error) {
+	c.closeErrOnce.Do(func() {
+		c.closeErr = fmt.Errorf("websocket closed: %w", err)
+	})
+}
diff --git a/websocket_js.go b/websocket_js.go
@@ -23,24 +23,26 @@ type Conn struct {
 
 	msgReadLimit int64
 
-	readClosed int64
-	closeOnce  sync.Once
-	closed     chan struct{}
-	closeErr   error
+	readClosed   int64
+	closeOnce    sync.Once
+	closed       chan struct{}
+	closeErrOnce sync.Once
+	closeErr     error
 
 	releaseOnClose   func()
 	releaseOnMessage func()
 
 	readSignal chan struct{}
 	readBufMu  sync.Mutex
-	readBuf    []wsjs.MessageEvent
+	// Max size of readBuf is 32.
+	readBuf []wsjs.MessageEvent
 }
 
 func (c *Conn) close(err error) {
 	c.closeOnce.Do(func() {
 		runtime.SetFinalizer(c, nil)
 
-		c.closeErr = fmt.Errorf("websocket closed: %w", err)
+		c.setCloseErr(err)
 		close(c.closed)
 	})
 }
@@ -49,6 +51,8 @@ func (c *Conn) init() {
 	c.closed = make(chan struct{})
 	c.readSignal = make(chan struct{}, 1)
 	c.msgReadLimit = 32768
+	// Capacity limited to 32 messages.
+	c.readBuf = make([]wsjs.MessageEvent, 0, 32)
 
 	c.releaseOnClose = c.ws.OnClose(func(e wsjs.CloseEvent) {
 		cerr := CloseError{
@@ -66,6 +70,12 @@ func (c *Conn) init() {
 		c.readBufMu.Lock()
 		defer c.readBufMu.Unlock()
 
+		if len(c.readBuf) == cap(c.readBuf) {
+			c.setCloseErr(fmt.Errorf("too many messages in buffer, cannot keep up: %v", len(c.readBuf)))
+			c.Close(StatusPolicyViolation, "unable to read fast enough")
+			return
+		}
+
 		c.readBuf = append(c.readBuf, e)
 
 		// Lets the read goroutine know there is definitely something in readBuf.
@@ -76,11 +86,15 @@ func (c *Conn) init() {
 	})
 
 	runtime.SetFinalizer(c, func(c *Conn) {
-		c.ws.Close(int(StatusInternalError), "")
-		c.close(errors.New("connection garbage collected"))
+		c.setCloseErr(errors.New("connection garbage collected"))
+		c.closeWithInternal()
 	})
 }
 
+func (c *Conn) closeWithInternal() {
+	c.Close(StatusInternalError, "something went wrong")
+}
+
 // Read attempts to read a message from the connection.
 // The maximum time spent waiting is bounded by the context.
 func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
@@ -113,11 +127,8 @@ func (c *Conn) read(ctx context.Context) (MessageType, []byte, error) {
 	defer c.readBufMu.Unlock()
 
 	me := c.readBuf[0]
-	// Ensures GC can collect the message event.
-	c.readBuf[0] = wsjs.MessageEvent{}
-	// We do not shrink the array since it will be resized
-	// as appropriate by append in the OnMessage callback.
-	c.readBuf = c.readBuf[1:]
+	copy(c.readBuf, c.readBuf[1:])
+	c.readBuf = c.readBuf[:len(c.readBuf)-1]
 
 	if len(c.readBuf) > 0 {
 		// Next time we read, we'll grab the message.
@@ -146,8 +157,10 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
 		// to match the Go API. It can only error if the message type
 		// is unexpected or the passed bytes contain invalid UTF-8 for
 		// MessageText.
-		c.Close(StatusInternalError, "something went wrong")
-		return fmt.Errorf("failed to write: %w", err)
+		err := fmt.Errorf("failed to write: %w", err)
+		c.setCloseErr(err)
+		c.closeWithInternal()
+		return err
 	}
 	return nil
 }
