declarative websockets: parameter binding

Author: kliushnichenko

modules/jooby-apt/src/main/java/io/jooby/internal/apt/MvcParameter.java

@@ -183,4 +183,8 @@ private List<? extends AnnotationMirror> annotationFromAnnotationType(Element el
   public boolean isRequireBeanValidation() {
     return requireBeanValidation;
   }
+
+  public VariableElement variableElement() {
+    return parameter;
+  }
 }

modules/jooby-apt/src/main/java/io/jooby/internal/apt/ws/WsRoute.java

@@ -58,30 +58,6 @@ public TypeDefinition getReturnType() {
     return new TypeDefinition(types, method.getReturnType());
   }
 
-  private String wsParameterName(MvcParameter parameter) {
-    String rawParamType = parameter.getType().getRawType().toString();
-    var name = WsParamTypes.generateArgumentName(rawParamType);
-    if (name != null) {
-      return name;
-    }
-
-    getContext()
-        .error("Unsupported websocket handler parameter type: %s.", rawParamType);
-    return "null";
-  }
-
-  private String paramList() {
-    var joiner = new StringJoiner(", ", "(", ")");
-    for (var param : getParameters(true)) {
-      joiner.add(wsParameterName(param));
-    }
-    return joiner.toString();
-  }
-
-  public String invocation(boolean kt) {
-    return makeCall(kt, paramList(), false, false);
-  }
-
   public void appendBody(boolean kt, StringBuilder buffer, String indent) {
     buffer.append(statement(indent, var(kt), "c = this.factory.apply(ctx)", semicolon(kt)));
 
@@ -112,6 +88,54 @@ public void appendBody(boolean kt, StringBuilder buffer, String indent) {
     }
   }
 
+  public String invocation(boolean kt) {
+    return makeCall(kt, paramList(kt), false, false);
+  }
+
+  private String paramList(boolean kt) {
+    var joiner = new StringJoiner(", ", "(", ")");
+    for (var param : getParameters(true)) {
+      joiner.add(websocketArgumentExpression(param, kt));
+    }
+    return joiner.toString();
+  }
+
+  private String websocketArgumentExpression(MvcParameter parameter, boolean kt) {
+    String rawParamType = parameter.getType().getRawType().toString();
+    if (wsLifecycle == WsLifecycle.MESSAGE) {
+      if (WsParamTypes.getAllowedTypes(WsLifecycle.MESSAGE).contains(rawParamType)) {
+        return WsParamTypes.generateArgumentName(rawParamType);
+      }
+
+      var mvcBodyExpr =
+          ParameterGenerator.BodyParam.toSourceCode(
+              kt,
+              this,
+              null,
+              parameter.getType(),
+              parameter.variableElement(),
+              parameter.getName(),
+              parameter.isNullable(kt));
+      return mvcExprToWsExpr(mvcBodyExpr);
+    }
+
+    String expr = WsParamTypes.generateArgumentName(rawParamType);
+    if (expr != null) {
+      return expr;
+    }
+
+    getContext()
+        .error("Unsupported websocket handler parameter type: %s.", rawParamType);
+    return "null";
+  }
+
+  private static String mvcExprToWsExpr(String mvcBodyExpr) {
+    String s = mvcBodyExpr;
+    s = s.replace("ctx.body().", "message.");
+    s = s.replace("ctx.body(", "message.to(");
+    return s;
+  }
+
   private void validateLifecycleParameters(MvcContext context, ExecutableElement method) {
     var env = context.getProcessingEnvironment();
     var types = env.getTypeUtils();
@@ -131,6 +155,10 @@ private void validateLifecycleParameters(MvcContext context, ExecutableElement m
         continue;
       }
 
+      if (wsLifecycle == WsLifecycle.MESSAGE) {
+        continue;
+      }
+
       context.error(
           "Illegal parameter type %s on websocket %s method %s#%s",
           raw,

modules/jooby-apt/src/test/java/tests/ws/WebsocketBeanMessage.java

@@ -0,0 +1,23 @@
+/*
+ * Jooby https://jooby.io
+ * Apache License Version 2.0 https://jooby.io/LICENSE.txt
+ * Copyright 2014 Edgar Espina
+ */
+package tests.ws;
+
+import io.jooby.Context;
+import io.jooby.WebSocket;
+import io.jooby.annotation.ws.OnMessage;
+
+import java.util.Map;
+
+public class WebsocketBeanMessage {
+
+  public record Incoming(String text) {
+  }
+
+  @OnMessage
+  public Map<String, String> onMessage(WebSocket ws, Context ctx, Incoming msg) {
+    return Map.of("echo", msg.text());
+  }
+}

modules/jooby-apt/src/test/java/tests/ws/WebsocketGeneratorTest.java

@@ -26,6 +26,20 @@ public void chatWebsocketMatchesGeneratedSource() throws Exception {
         );
   }
 
+  @Test
+  public void beanMessageWebsocketMatchesGeneratedSource() throws Exception {
+    var expected = new String(
+        getClass()
+            .getResourceAsStream("/tests/ws/WebsocketBeanMessageWs_expected.java")
+            .readAllBytes()
+    );
+
+    new ProcessorRunner(new WebsocketBeanMessage())
+        .withWsCode(source -> assertThat(normalize(source))
+            .isEqualTo(normalize(expected))
+        );
+  }
+
   private static String normalize(String source) {
     return source.replace("\r\n", "\n").replace('\r', '\n')
         .stripTrailing();

modules/jooby-apt/src/test/resources/tests/ws/WebsocketBeanMessageWs_expected.java

@@ -0,0 +1,40 @@
+package tests.ws;
+
+@io.jooby.annotation.Generated(WebsocketBeanMessage.class)
+public class WebsocketBeanMessageWs_ implements io.jooby.Extension {
+    protected java.util.function.Function<io.jooby.Context, WebsocketBeanMessage> factory;
+
+    public WebsocketBeanMessageWs_() {
+      this(io.jooby.SneakyThrows.singleton(WebsocketBeanMessage::new));
+    }
+
+    public WebsocketBeanMessageWs_(WebsocketBeanMessage instance) {
+      setup(ctx -> instance);
+    }
+
+    public WebsocketBeanMessageWs_(io.jooby.SneakyThrows.Supplier<WebsocketBeanMessage> provider) {
+      setup(ctx -> provider.get());
+    }
+
+    public WebsocketBeanMessageWs_(io.jooby.SneakyThrows.Function<Class<WebsocketBeanMessage>, WebsocketBeanMessage> provider) {
+      setup(ctx -> provider.apply(WebsocketBeanMessage.class));
+    }
+
+    private void setup(java.util.function.Function<io.jooby.Context, WebsocketBeanMessage> factory) {
+      this.factory = factory;
+    }
+
+    public void install(io.jooby.Jooby app) throws Exception {
+      app.ws("/", this::wsInit);
+    }
+
+    private void wsInit(io.jooby.Context ctx, io.jooby.WebSocketConfigurer configurer) {
+      /** See {@link WebsocketBeanMessage#onMessage(io.jooby.WebSocket, io.jooby.Context, tests.ws.WebsocketBeanMessage.Incoming)} */
+      configurer.onMessage((ws, message) -> {
+        var c = this.factory.apply(ctx);
+        var __wsReturn = c.onMessage(ws, ctx, message.to(tests.ws.WebsocketBeanMessage.Incoming.class));
+        ws.render(__wsReturn);
+      });
+
+    }
+}