Fix the issue of dispatching to custom remote functions if onMessage …

…is present
ballerina-platform · Apr 21, 2023 · c938c08 · c938c08
1 parent eccdd2b
commit c938c08
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 3 deletions.
diff --git a/ballerina/tests/custom_remote_functions.bal b/ballerina/tests/custom_remote_functions.bal
@@ -123,6 +123,24 @@ service class LotOfUnderscoreService {
     }
 }
 
+@ServiceConfig {dispatcherKey: "type"}
+service /onMessageWithCustom on customDispatchingLis {
+    resource function get .() returns Service|Error {
+        return new OnMessageService();
+    }
+}
+
+service class OnMessageService {
+    *Service;
+    remote function onMessage(Caller caller, json data) returns Error? {
+        check caller->writeMessage({"type": "onMessage"});
+    }
+
+    remote function onConnectionInit(Caller caller, string text) returns Error? {
+        check caller->writeMessage({"type": "onConnectionInit"});
+    }
+}
+
 @test:Config {}
 public function testPingMessage() returns Error? {
     Client cl = check new("ws://localhost:21401");
@@ -187,3 +205,11 @@ public function testUnderscoresAndSpaces() returns Error? {
     json resp = check cl->readMessage();
     test:assertEquals(resp, {"event": "onMessages"});
 }
+
+@test:Config {}
+public function testOnMessageAtTheBeginning() returns Error? {
+    Client cl = check new("ws://localhost:21401/onMessageWithCustom");
+    check cl->writeMessage({'type: "connection_init"});
+    json resp = check cl->readMessage();
+    test:assertEquals(resp, {"type": "onConnectionInit"});
+}
diff --git a/native/src/main/java/io/ballerina/stdlib/websocket/WebSocketResourceDispatcher.java b/native/src/main/java/io/ballerina/stdlib/websocket/WebSocketResourceDispatcher.java
@@ -420,12 +420,14 @@ public static void dispatchOnText(WebSocketConnectionInfo connectionInfo, WebSoc
             MethodType[] remoteFunctions = ((ServiceType) (((BValue) dispatchingService).getType())).getMethods();
             for (MethodType remoteFunc : remoteFunctions) {
                 String funcName = remoteFunc.getName();
-                if (funcName.equals(methodName) ||
-                        funcName.equals(WebSocketConstants.RESOURCE_NAME_ON_TEXT_MESSAGE) ||
-                        funcName.equals(WebSocketConstants.RESOURCE_NAME_ON_MESSAGE)) {
+                if (funcName.equals(methodName)) {
                     onTextMessageResource = remoteFunc;
                     break;
                 }
+                if (funcName.equals(WebSocketConstants.RESOURCE_NAME_ON_TEXT_MESSAGE) ||
+                        funcName.equals(WebSocketConstants.RESOURCE_NAME_ON_MESSAGE)) {
+                    onTextMessageResource = remoteFunc;
+                }
             }
             if (onTextMessageResource == null) {
                 stringAggregator.resetAggregateString();