Test fixes + more tests

Author: joefarebrother

java/ql/lib/ext/org.springframework.web.socket.model.yml

@@ -11,7 +11,7 @@ extensions:
      - ["org.springframework.web.socket.handler", "AbstractWebSocketHandler", True, "handleBinaryMessage", "", "", "Parameter[0]", "remote", "manual"]
      - ["org.springframework.web.socket.handler", "AbstractWebSocketHandler", True, "handleBinaryMessage", "", "", "Parameter[1]", "remote", "manual"]
      - ["org.springframework.web.socket.handler", "AbstractWebSocketHandler", True, "handleTextMessage", "", "", "Parameter[0]", "remote", "manual"]
-     - ["org.springframework.web.socket.handler", "AbstractWebSocketHandler", True, "handleTextMessage", "", "", "Parameter[0]", "remote", "manual"]
+     - ["org.springframework.web.socket.handler", "AbstractWebSocketHandler", True, "handleTextMessage", "", "", "Parameter[1]", "remote", "manual"]
   - addsTo:
       pack: codeql/java-all
       extensible: summaryModel

java/ql/test/library-tests/frameworks/spring/websocket/Test.java

@@ -2,21 +2,56 @@
 import org.springframework.web.socket.handler.TextWebSocketHandler;
 import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.WebSocketMessage;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.BinaryMessage;
+import org.springframework.web.socket.CloseStatus;
 
 
-public class Test extends TextWebSocketHandler {
+public class Test  {
     void sink(Object o) {}
 
-    @Override
-    public void handleMessage(WebSocketSession s, WebSocketMessage<?> m) {
-        sink(s); // $hasTaintFlow
-        sink(s.getAcceptedProtocol()); // $hasTaintFlow
-        sink(s.getHandshakeHeaders()); // $hasTaintFlow
-        sink(s.getPrincipal()); // $hasTaintFlow
-        sink(s.getUri()); // $hasTaintFlow
+    public class A extends TextWebSocketHandler {
+        @Override
+        public void handleMessage(WebSocketSession s, WebSocketMessage<?> m) {
+            sink(s); // $hasTaintFlow
+            sink(s.getAcceptedProtocol()); // $hasTaintFlow
+            sink(s.getHandshakeHeaders()); // $hasTaintFlow
+            sink(s.getPrincipal()); // $hasTaintFlow
+            sink(s.getUri()); // $hasTaintFlow
 
-        sink(m); // $hasTaintFlow
-        sink(m.getPayload()); // $hasTaintFlow
+            sink(m); // $hasTaintFlow
+            sink(m.getPayload()); // $hasTaintFlow
+
+        }
+
+        @Override 
+        protected void handleTextMessage(WebSocketSession s, TextMessage m) {
+            sink(s);  // $hasTaintFlow
+            sink(m);  // $hasTaintFlow
+            sink(m.asBytes()); // $hasTaintFlow
+        }
+
+        @Override 
+        protected void handleBinaryMessage(WebSocketSession s, BinaryMessage m) {
+            sink(s); // $hasTaintFlow
+            sink(m); // $hasTaintFlow
+        }
+
+        @Override
+        public void afterConnectionEstablished(WebSocketSession s) {
+            sink(s); // $hasTaintFlow
+        }
+
+        @Override 
+        public void afterConnectionClosed(WebSocketSession s, CloseStatus c) {
+            sink(s); // $ hasTaintFlow
+        }
+
+        @Override 
+        public void handleTransportError(WebSocketSession s, Throwable exc) { 
+            sink(s);  // $ hasTaintFlow
+        }
 
     }
+
 }

java/ql/test/library-tests/frameworks/spring/websocket/test.ql

@@ -1,4 +1,16 @@
 import java
+import semmle.code.java.dataflow.DataFlow
+import semmle.code.java.dataflow.FlowSources
 import utils.test.InlineFlowTest
-import DefaultFlowTest
-import TaintFlow::PathGraph
+
+module Config implements DataFlow::ConfigSig {
+  predicate isSource(DataFlow::Node node) {
+    DefaultFlowConfig::isSource(node)
+    or
+    node instanceof ActiveThreatModelSource
+  }
+
+  predicate isSink = DefaultFlowConfig::isSink/1;
+}
+
+import FlowTest<DefaultFlowConfig, Config>