Python: Add test for function

python/ql/test/query-tests/Security/CWE-943-NoSqlInjection/PoC/server.py

@@ -46,6 +46,32 @@ def by_where():
     post = posts.find_one({'$where': 'this.author === "'+author+'"'}) # $ result=BAD
     return show_post(post, author)
 
+
+@app.route('/byFunction', methods=['GET'])
+def by_function():
+    author = request.args['author']
+    search = {
+        "body": 'function(author) { return(author === "'+author+'") }',
+        "args": [ "$author" ],
+        "lang": "js"
+    }
+    # Use `" | "a" === "a` as author
+    # making the query `this.author === "" | "a" === "a"`
+    # Found by http://127.0.0.1:5000/byFunction?author=%22%20|%20%22a%22%20===%20%22a
+    post = posts.find_one({'$expr': {'$function': search}}) # $ MISING: result=BAD
+    return show_post(post, author)
+
+@app.route('/byFunctionArg', methods=['GET'])
+def by_function_arg():
+    author = request.args['author']
+    search = {
+        "body": 'function(author, target) { return(author === target) }',
+        "args": [ "$author", author ],
+        "lang": "js"
+    }
+    post = posts.find_one({'$expr': {'$function': search}}) # $ result=OK
+    return show_post(post, author)
+
 @app.route('/', methods=['GET'])
 def show_routes():
     links = []