WebSocket: Allow to configure websocket size

jknack · jknack · commit 2c40dee977ec · 2022-04-26T20:56:23.000-03:00
- Fix jooby-project#2566
diff --git a/docs/asciidoc/websocket.adoc b/docs/asciidoc/websocket.adoc
@@ -1,7 +1,6 @@
 == Web Sockets
 
-Jooby supports https://developer.mozilla.org/es/docs/Web/API/WebSockets_API[WebSockets].
-A javadoc:WebSocket[] is registered like any other handler:
+Adding a javadoc:WebSocket[]:
 
 .WebSocket
 [source,java,role="primary"]
@@ -170,7 +169,9 @@ import io.jooby.json.JacksonModule;
 Structure messages depends/requires a javadoc:MessageDecoder[] and jadoc:MessageEncoder[]. In this
 example both are provided by the JacksonModule.
 
-=== Connection Timeouts
+=== Options
+
+==== Connection Timeouts
 By default Jooby will timeout idle connections that have no activity after 5 minutes. You can 
 control this behaviour by setting the `websocket.idleTimeout` property:
 
@@ -181,3 +182,15 @@ websocket.idleTimeout = 1h
 ----
 
 See https://github.com/lightbend/config/blob/master/HOCON.md#duration-format[duration format]
+
+==== Max size
+
+Max size is set to `128K` you can override it like:
+
+.application.conf
+[source, properties]
+----
+websocket.maxSize = 128K
+----
+
+See https://github.com/lightbend/config/blob/master/HOCON.md#size-in-bytes-format[sizes in bytes]
diff --git a/jooby/src/main/java/io/jooby/WebSocket.java b/jooby/src/main/java/io/jooby/WebSocket.java
@@ -108,7 +108,7 @@ interface OnError {
     void onError(@Nonnull WebSocket ws, @Nonnull Throwable cause);
   }
 
-  /** Max message size for websocket (13k). */
+  /** Max message size for websocket (128K). */
   int MAX_BUFFER_SIZE = 131072;
 
   /**
diff --git a/modules/jooby-jetty/src/main/java/io/jooby/jetty/Jetty.java b/modules/jooby-jetty/src/main/java/io/jooby/jetty/Jetty.java
@@ -186,11 +186,14 @@ public class Jetty extends io.jooby.Server.Base {
         context.setHandler(handler);
       }
       /* ********************************* WebSocket *************************************/
+      Config conf = application.getConfig();
+      int maxSize = conf.hasPath("websocket.maxSize")
+          ? conf.getBytes("websocket.maxSize").intValue()
+          : WebSocket.MAX_BUFFER_SIZE;
       context.setAttribute(DecoratedObjectFactory.ATTR, new DecoratedObjectFactory());
       WebSocketPolicy policy = new WebSocketPolicy(WebSocketBehavior.SERVER);
-      policy.setMaxTextMessageBufferSize(WebSocket.MAX_BUFFER_SIZE);
-      policy.setMaxTextMessageSize(WebSocket.MAX_BUFFER_SIZE);
-      Config conf = application.getConfig();
+      policy.setMaxTextMessageBufferSize(maxSize);
+      policy.setMaxTextMessageSize(maxSize);
       long timeout = conf.hasPath("websocket.idleTimeout")
           ? conf.getDuration("websocket.idleTimeout", TimeUnit.MILLISECONDS)
           : TimeUnit.MINUTES.toMillis(5);
diff --git a/modules/jooby-netty/src/main/java/io/jooby/internal/netty/NettyContext.java b/modules/jooby-netty/src/main/java/io/jooby/internal/netty/NettyContext.java
@@ -373,12 +373,16 @@ boolean isHttpGet() {
   @Nonnull @Override public Context upgrade(WebSocket.Initializer handler) {
     try {
       responseStarted = true;
+      Config conf = getRouter().getConfig();
+      int maxSize = conf.hasPath("websocket.maxSize")
+          ? conf.getBytes("websocket.maxSize").intValue()
+          : WebSocket.MAX_BUFFER_SIZE;
       String webSocketURL = getProtocol() + "://" + req.headers().get(HttpHeaderNames.HOST) + path;
       WebSocketDecoderConfig config = WebSocketDecoderConfig.newBuilder()
           .allowExtensions(true)
           .allowMaskMismatch(false)
           .withUTF8Validator(false)
-          .maxFramePayloadLength(WebSocket.MAX_BUFFER_SIZE)
+          .maxFramePayloadLength(maxSize)
           .build();
       webSocket = new NettyWebSocket(this);
       handler.init(Context.readOnly(this), webSocket);
@@ -389,7 +393,6 @@ boolean isHttpGet() {
       WebSocketServerHandshaker handshaker = factory.newHandshaker(webSocketRequest);
       handshaker.handshake(ctx.channel(), webSocketRequest);
       webSocket.fireConnect();
-      Config conf = getRouter().getConfig();
       long timeout = conf.hasPath("websocket.idleTimeout")
           ? conf.getDuration("websocket.idleTimeout", MILLISECONDS)
           : MINUTES.toMillis(5);
diff --git a/modules/jooby-utow/src/main/java/io/jooby/internal/utow/UtowWebSocket.java b/modules/jooby-utow/src/main/java/io/jooby/internal/utow/UtowWebSocket.java
@@ -56,20 +56,26 @@ public class UtowWebSocket extends AbstractReceiveListener
   private String key;
   private CountDownLatch ready = new CountDownLatch(1);
   private AtomicBoolean open = new AtomicBoolean(false);
+  private int maxSize;
 
   public UtowWebSocket(UtowContext ctx, WebSocketChannel channel) {
     this.ctx = ctx;
     this.channel = channel;
     this.dispatch = !ctx.isInIoThread();
     this.key = ctx.getRoute().getPattern();
+
+    Config conf = ctx.getRouter().getConfig();
+    maxSize = conf.hasPath("websocket.maxSize")
+        ? conf.getBytes("websocket.maxSize").intValue()
+        : WebSocket.MAX_BUFFER_SIZE;
   }
 
   @Override protected long getMaxTextBufferSize() {
-    return MAX_BUFFER_SIZE;
+    return maxSize;
   }
 
   @Override protected long getMaxBinaryBufferSize() {
-    return MAX_BUFFER_SIZE;
+    return maxSize;
   }
 
   @Nonnull @Override public Context getContext() {

Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ interface OnError {`
`108`	`108`	`void onError(@Nonnull WebSocket ws, @Nonnull Throwable cause);`
`109`	`109`	`}`
`110`	`110`
`111`		`- /** Max message size for websocket (13k). */`
	`111`	`+ /** Max message size for websocket (128K). */`
`112`	`112`	`int MAX_BUFFER_SIZE = 131072;`
`113`	`113`
`114`	`114`	`/**`