Unified: Use getAnArgument() to check argument of a MaD token

unified/ql/lib/codeql/unified/internal/mad/LegacyMaD.qll

@@ -224,7 +224,8 @@ private DataFlow::Node getSinkFromCall(CallExpr call, AccessPathToken token) {
     arg.getName() + ":" = token.getAnArgument()
   )
   or
-  token = "Argument[self]" and
+  token.getName() = "Argument" and
+  token.getAnArgument() = "self" and
   result.isReceiverArgument(call)
 }
 
@@ -263,7 +264,8 @@ private DataFlow::Node getSourceFromCallable(Callable callable, AccessPathToken
     param.getExternalName() + ":" = token.getAnArgument()
   )
   or
-  token = "Parameter[self]" and
+  token.getName() = "Parameter" and
+  token.getAnArgument() = "self" and
   result.isReceiverParameter(callable)
 }
 

unified/ql/test/library-tests/mad/test.expected

@@ -16,6 +16,8 @@ isSink
 | test.swift:101:18:101:24 | fileURL | path-injection |
 | test.swift:107:23:107:34 | seedFilePath | path-injection |
 | test.swift:119:29:119:34 | string | path-injection |
+| test.swift:137:5:139:5 | [self arg] ... .sinkArgs(...) | shorthand-sink |
+| test.swift:138:9:138:14 | string | shorthand-sink |
 isSource
 | test.swift:4:5:4:27 | String(...) | remote |
 | test.swift:5:5:5:44 | String(...) | remote |
@@ -25,3 +27,5 @@ isSource
 | test.swift:10:5:12:24 | String(...) | local |
 | test.swift:13:5:15:28 | String(...) | local |
 | test.swift:29:29:29:32 | text | local |
+| test.swift:129:5:133:5 | [self param] func ShorthandTestSub.sourceParams | shorthand-source |
+| test.swift:130:11:130:11 | x | shorthand-source |

unified/ql/test/library-tests/mad/test.ext.yml

@@ -0,0 +1,11 @@
+extensions:
+  - addsTo:
+      pack: codeql/unified-all
+      extensible: legacySourceModel
+    data: # namespace, type, subtypes, name, signature, ext, output, kind, provenance
+      - ["", "ShorthandTest", true, "sourceParams(_:)", "", "", "Parameter[self,0]", "shorthand-source", "manual"]
+  - addsTo:
+      pack: codeql/unified-all
+      extensible: legacySinkModel
+    data: # namespace, type, subtypes, name, signature, ext, input, kind, provenance
+      - ["", "ShorthandTest", true, "sinkArgs(_:)", "", "", "Argument[self,0]", "shorthand-sink", "manual"]

unified/ql/test/library-tests/mad/test.swift

@@ -118,3 +118,23 @@ class Connection {
 func testConnectionLocation(string: String) {
     Connection.Location.uri(string, parameters: "")  // $ isSink=path-injection
 }
+
+class ShorthandTest {
+    func sourceParams(_ x: String) {}
+
+    func sinkArgs(_ x: String) {}
+}
+
+class ShorthandTestSub: ShorthandTest {
+    override func sourceParams(
+        _ x: String  // $ isSource=shorthand-source
+    ) {
+        _ = self
+    }  // $ isSource=shorthand-source
+}
+
+func testShorthandTokens(obj: ShorthandTest, string: String) {
+    obj.sinkArgs(
+        string  // $ isSink=shorthand-sink
+    )  // $ isSink=shorthand-sink
+}