Repository: ably/ably-go (go)

Issue:
Allow configuring ably/websocket read limits
The messages we are receiving through the Ably real-time client are somewhat large (seem to max out ~175k), but the default real limit for messages in `nhooyr.io/websocket` is 32,678 bytes. If a message is received larger than this limit, the connection is closed.

Ideally, we'd be able to set this message limit in the `ClientOptions`, but this seems like an implementation-specific configuration and may not apply to other connection mechanisms. 

I also see that the `Dial` is a client option. I'd like to use all of the `ably.websocketConn` (but increase the read limit) without needing to implement my own websocket connection/client functionality.

I'm happy to do some work here too and open a PR, but I want to make sure that 1) this is a change that would be welcome in the package and 2) the option I implement is the right path.

Thanks so much!!

-- Chris

Pull request:
[SDK-3731] Fix websocket read limits
Fixed #610 

Patch:
diff --git a/README.md b/README.md
index 5696751..7d1159f 100644
--- a/README.md
+++ b/README.md
@@ -196,6 +196,22 @@ if err != nil {
 }
 ```
 
+#### Update MaxMessageSize/read limit for realtime message subscription
+- The default `MaxMessageSize` is automatically configured by Ably when connection is established with Ably.
+- This value defaults to 64kb, please [get in touch](https://ably.com/support) if you would like to request a higher limit for your account.
+- Upgrading your account to higher limit will automatically update `MaxMessageSize` property and should accordingly set the client side connection read limit.
+- If you are still facing issues when receiving large messages or intentionally want to reduce the limit, you can explicitly update the connection read limit:
+
+```go
+client, err = ably.NewRealtime(ably.WithKey("xxx:xxx"))
+if err != nil {
+        panic(err)
+}
+client.Connection.SetReadLimit(131072) // Set read limit to 128kb, overriding default ConnectionDetails.MaxMessageSize
+```
+
+Note - If connection read limit is less than size of received message, the client will throw an error "failed to read: read limited at {READ_LIMIT + 1} bytes" and will close the connection.
+
 ### Using the REST API
 
 #### Introduction
diff --git a/ably/options.go b/ably/options.go
index c5bc29d..3ab536d 100644
--- a/ably/options.go
+++ b/ably/options.go
@@ -24,9 +24,10 @@ const (
 	// restHost is the primary ably host.
 	restHost = "rest.ably.io"
 	// realtimeHost is the primary ably host.
-	realtimeHost = "realtime.ably.io"
-	Port         = 80
-	TLSPort      = 443
+	realtimeHost   = "realtime.ably.io"
+	Port           = 80
+	TLSPort        = 443
+	maxMessageSize = 65536 // 64kb, default value TO3l8
 )
 
 var defaultOptions = clientOptions{
diff --git a/ably/realtime_conn.go b/ably/realtime_conn.go
index 2b7ba8c..3959c8f 100644
--- a/ably/realtime_conn.go
+++ b/ably/realtime_conn.go
@@ -83,6 +83,9 @@ type Connection struct {
 	// after a reauthorization, to avoid re-reauthorizing.
 	reauthorizing bool
 	arg           connArgs
+
+	readLimit                int64
+	isReadLimitSetExternally bool
 }
 
 type connCallbacks struct {
@@ -107,6 +110,7 @@ func newConn(opts *clientOptions, auth *Auth, callbacks connCallbacks) *Connecti
 		pending:   newPendingEmitter(auth.log()),
 		auth:      auth,
 		callbacks: callbacks,
+		readLimit: maxMessageSize,
 	}
 	auth.onExplicitAuthorize = c.onClientAuthorize
 	c.queue = newMsgQueue(c)
@@ -176,6 +180,15 @@ func (c *Connection) Connect() {
 	}()
 }
 
+// SetReadLimit is used to override internal websocket connection read limit.
+// It sets the max number of bytes to read for a single message.
+// By default, the connection has a message read limit of [ably.maxMessageSize] or 65536 bytes.
+// When the limit is hit, the connection will be closed with StatusMessageTooBig.
+func (c *Connection) SetReadLimit(readLimit int64) {
+	c.readLimit = readLimit
+	c.isReadLimitSetExternally = true
+}
+
 // Close causes the connection to close, entering the [ably.ConnectionStateClosing] state.
 // Once connection state is [ably.ConnectionStateClosed], the library does not attempt to re-establish
 // the connection without an explicit call to Connection.Connect (RTN12).
@@ -792,6 +805,16 @@ func (c *Connection) eventloop() {
 				c.connStateTTL = connDetails.ConnectionStateTTL
 				// Spec RSA7b3, RSA7b4, RSA12a
 				c.auth.updateClientID(connDetails.ClientID)
+				if !c.isReadLimitSetExternally {
+					c.readLimit = connDetails.MaxMessageSize // set MaxMessageSize limit as per TO3l8
+				}
+			}
+			err := setConnectionReadLimit(c.conn, c.readLimit)
+			if err != nil {
+				c.readLimit = -1
+				c.log().Errorf("%v, using default readlimit instead", err)
+			} else {
+				c.log().Verbosef("connection readlimit set to %v", c.readLimit)
 			}
 			reconnecting := c.reconnecting
 			if reconnecting {
diff --git a/ably/websocket.go b/ably/websocket.go
index 02483cf..9a9370a 100644
--- a/ably/websocket.go
+++ b/ably/websocket.go
@@ -113,3 +113,16 @@ func dialWebsocketTimeout(uri, origin string, timeout time.Duration, agents map[
 
 	return c, nil
 }
+
+func setConnectionReadLimit(c conn, readLimit int64) error {
+	verboseConn, ok := c.(verboseConn)
+	if !ok {
+		return errors.New("cannot set readlimit for connection, connection does not use verboseConn")
+	}
+	websocketConn, ok := verboseConn.conn.(*websocketConn)
+	if !ok {
+		return errors.New("cannot set readlimit for connection, connection does not use nhooyr.io/websocket")
+	}
+	websocketConn.conn.SetReadLimit(readLimit)
+	return nil
+}


Test patch:
diff --git a/ably/export_test.go b/ably/export_test.go
index 4bdd5e2..1da81e6 100644
--- a/ably/export_test.go
+++ b/ably/export_test.go
@@ -136,6 +136,18 @@ func (c *Connection) MsgSerial() int64 {
 	return c.msgSerial
 }
 
+func (c *Connection) IsReadLimitSetExternally() bool {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	return c.isReadLimitSetExternally
+}
+
+func (c *Connection) ReadLimit() int64 {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	return c.readLimit
+}
+
 func WithTrace(trace *httptrace.ClientTrace) ClientOption {
 	return func(os *clientOptions) {
 		os.Trace = trace
diff --git a/ably/realtime_channel_integration_test.go b/ably/realtime_channel_integration_test.go
index 5859a69..88dadd5 100644
--- a/ably/realtime_channel_integration_test.go
+++ b/ably/realtime_channel_integration_test.go
@@ -268,3 +268,64 @@ func TestRealtimeChannel_AttachWhileDisconnected(t *testing.T) {
 	err = ablytest.Wait(<-res, nil)
 	assert.NoError(t, err)
 }
+
+func TestRealtimeChannel_ShouldSetAblySandboxDefaultReadLimit(t *testing.T) {
+	app, client := ablytest.NewRealtime(ably.WithEchoMessages(false))
+	defer safeclose(t, ablytest.FullRealtimeCloser(client), app)
+	assert.Equal(t, int64(65536), client.Connection.ReadLimit()) // Default read limit when not connected
+
+	err := ablytest.Wait(ablytest.ConnWaiter(client, client.Connect, ably.ConnectionEventConnected), nil)
+	assert.NoError(t, err)
+
+	assert.Equal(t, int64(16384), client.Connection.ReadLimit()) // sandbox read limit
+	assert.False(t, client.Connection.IsReadLimitSetExternally())
+}
+
+func TestRealtimeChannel_ShouldSetProvidedReadLimit(t *testing.T) {
+	app, client := ablytest.NewRealtime(ably.WithEchoMessages(false))
+	defer safeclose(t, ablytest.FullRealtimeCloser(client), app)
+	client.Connection.SetReadLimit(2048)
+
+	err := ablytest.Wait(ablytest.ConnWaiter(client, client.Connect, ably.ConnectionEventConnected), nil)
+	assert.NoError(t, err)
+
+	assert.True(t, client.Connection.IsReadLimitSetExternally())
+	assert.Equal(t, int64(2048), client.Connection.ReadLimit())
+}
+
+func TestRealtimeChannel_ShouldReturnErrorIfReadLimitExceeded(t *testing.T) {
+	app, client1 := ablytest.NewRealtime(ably.WithEchoMessages(false))
+	defer safeclose(t, ablytest.FullRealtimeCloser(client1), app)
+
+	client2 := app.NewRealtime(ably.WithEchoMessages(false))
+	client2.Connection.SetReadLimit(1024) // set read limit explicitly to 1 mb.
+
+	err := ablytest.Wait(ablytest.ConnWaiter(client1, client1.Connect, ably.ConnectionEventConnected), nil)
+	assert.NoError(t, err)
+	err = ablytest.Wait(ablytest.ConnWaiter(client2, client2.Connect, ably.ConnectionEventConnected), nil)
+	assert.NoError(t, err)
+
+	assert.True(t, client2.Connection.IsReadLimitSetExternally())
+	assert.Equal(t, int64(1024), client2.Connection.ReadLimit())
+
+	channel1 := client1.Channels.Get("test")
+	channel2 := client2.Channels.Get("test")
+
+	err = channel1.Attach(context.Background())
+	assert.NoError(t, err, "client1: Attach()=%v", err)
+	err = channel2.Attach(context.Background())
+	assert.NoError(t, err, "client2: Attach()=%v", err)
+
+	_, unsub2, err := ablytest.ReceiveMessages(channel2, "")
+	assert.NoError(t, err, "client2:.Subscribe(context.Background())=%v", err)
+	defer unsub2()
+
+	messageWith2MbSize := ablytest.GenerateRandomString(2048)
+	err = channel1.Publish(context.Background(), "hello", messageWith2MbSize)
+	assert.NoError(t, err, "client1: Publish()=%v", err)
+
+	err = ablytest.Wait(ablytest.ConnWaiter(client2, nil, ably.ConnectionEventClosed), nil)
+	assert.NotNil(t, err)
+	errorInfo := err.(*ably.ErrorInfo)
+	assert.Equal(t, "failed to read: read limited at 1025 bytes", errorInfo.Unwrap().Error())
+}
diff --git a/ablytest/ablytest.go b/ablytest/ablytest.go
index fc7b836..dba905a 100644
--- a/ablytest/ablytest.go
+++ b/ablytest/ablytest.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"math/rand"
 	"net/http"
 	"os"
 	"reflect"
@@ -209,3 +210,15 @@ func TimeFuncs(afterCalls chan<- AfterCall) (
 
 	return now, after
 }
+
+const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+var seededRand *rand.Rand = rand.New(rand.NewSource(time.Now().UnixNano()))
+
+func GenerateRandomString(length int) string {
+	b := make([]byte, length)
+	for i := range b {
+		b[i] = charset[seededRand.Intn(len(charset))]
+	}
+	return string(b)
+}