Implement websockets models

python/ql/test/library-tests/frameworks/websockets/ConceptsTest.ql

@@ -0,0 +1,2 @@
+import python
+import experimental.meta.ConceptsTest

python/ql/test/library-tests/frameworks/websockets/InlineTaintTest.expected

@@ -0,0 +1,3 @@
+argumentToEnsureNotTaintedNotMarkedAsSpurious
+untaintedArgumentToEnsureTaintedNotMarkedAsMissing
+testFailures

python/ql/test/library-tests/frameworks/websockets/InlineTaintTest.ql

@@ -0,0 +1,2 @@
+import experimental.meta.InlineTaintTest
+import MakeInlineTaintTest<TestTaintTrackingConfig>

python/ql/test/library-tests/frameworks/websockets/response_test.py

@@ -0,0 +1,46 @@
+import websockets.sync.server 
+import websockets.sync.router
+from werkzeug.routing import Map, Rule
+
+def arg_handler(websocket): # $ requestHandler routedParameter=websocket
+    websocket.send("arg" + websocket.recv())
+
+s1 = websockets.sync.server.serve(arg_handler, "localhost", 8000)
+
+def kw_handler(websocket): # $ requestHandler routedParameter=websocket
+    websocket.send("kw" + websocket.recv())
+
+s2 = websockets.sync.server.serve(handler=kw_handler, host="localhost", port=8001)
+
+def route_handler(websocket, x):  # $ MISSING: requestHandler routedParameter=websocket routedParameter=x
+    websocket.send(f"route {x} {websocket.recv()}")
+
+s3 = websockets.sync.router.route(Map([
+    Rule("/<string:x>", endpoint=route_handler)
+]), "localhost", 8002)
+
+def unix_handler(websocket): # $ requestHandler routedParameter=websocket
+    websocket.send("unix" + websocket.recv())
+
+s4 = websockets.sync.server.unix_serve(unix_handler, path="/tmp/ws.sock")
+
+def unix_route_handler(websocket, x):  # $ MISSING: requestHandler routedParameter=websocket routedParameter=x
+    websocket.send(f"unix route {x} {websocket.recv()}")
+
+s5 = websockets.sync.router.unix_route(Map([
+    Rule("/<string:x>", endpoint=unix_route_handler)
+]), path="/tmp/ws2.sock")
+
+if __name__ == "__main__":
+    import sys
+    server = s1 
+    if len(sys.argv) > 1:
+        if sys.argv[1] == "kw":
+            server = s2
+        elif sys.argv[1] == "route":
+            server = s3
+        elif sys.argv[1] == "unix":
+            server = s4
+        elif sys.argv[1] == "unix_route":
+            server = s5
+    server.serve_forever()

python/ql/test/library-tests/frameworks/websockets/taint_test_asyncio.py

@@ -0,0 +1,30 @@
+import websockets.asyncio.server
+import asyncio
+
+def ensure_tainted(*args):
+    print("tainted", args)
+
+def ensure_not_tainted(*args):
+    print("not tainted", args)
+
+async def handler(websocket): # $ requestHandler routedParameter=websocket
+    ensure_tainted(
+        websocket, # $ tainted
+        await websocket.recv()  # $ tainted
+    )
+
+    async for msg in websocket:
+        ensure_tainted(msg)  # $ tainted
+        await websocket.send(msg) 
+
+    async for msg in websocket.recv_streaming():
+        ensure_tainted(msg)  # $ tainted
+        await websocket.send(msg)
+
+
+async def main():
+    server = await websockets.asyncio.server.serve(handler, "localhost", 8000)
+    await server.serve_forever()
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/ql/test/library-tests/frameworks/websockets/taint_test_sync.py

@@ -0,0 +1,26 @@
+import websockets.sync.server
+
+def ensure_tainted(*args):
+    print("tainted", args)
+
+def ensure_not_tainted(*args):
+    print("not tainted", args)
+
+def handler(websocket): # $ requestHandler routedParameter=websocket
+    ensure_tainted(
+        websocket, # $ tainted
+        websocket.recv()  # $ tainted
+    )
+
+    for msg in websocket:
+        ensure_tainted(msg)  # $ tainted
+        websocket.send(msg) 
+
+    for msg in websocket.recv_streaming():
+        ensure_tainted(msg)  # $ tainted
+        websocket.send(msg)
+
+
+if __name__ == "__main__":
+    server = websockets.sync.server.serve(handler, "localhost", 8000)
+    server.serve_forever()