Skip to content

Commit 4a60130

Browse files
paldepindpaldepind
committed
Rust: Make await call async blocks as lambdas in dataflow
1 parent d099129 commit 4a60130

4 files changed

Lines changed: 61 additions & 8 deletions

File tree

‎rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,9 @@ final class DataFlowCall extends TDataFlowCall {
6262
/** Gets the underlying call, if any. */
6363
Call asCall() { this = TCall(result) }
6464

65+
/** Gets the underlying await expression, if any. */
66+
AwaitExpr asAwaitExpr() { this = TAwaitExpr(result) }
67+
6568
predicate isImplicitDerefCall(Expr e, DerefChain derefChain, int i, Function target) {
6669
this = TImplicitDerefCall(e, derefChain, i, target)
6770
}
@@ -73,15 +76,16 @@ final class DataFlowCall extends TDataFlowCall {
7376
}
7477

7578
DataFlowCallable getEnclosingCallable() {
76-
result.asCfgScope() = this.asCall().getEnclosingCfgScope()
79+
result.asCfgScope() =
80+
[this.asCall().getEnclosingCfgScope(), this.asAwaitExpr().getEnclosingCfgScope()]
7781
or
7882
result.asCfgScope() = any(Expr e | this.isImplicitDerefCall(e, _, _, _)).getEnclosingCfgScope()
7983
or
8084
this.isSummaryCall(result.asSummarizedCallable(), _)
8185
}
8286

8387
string toString() {
84-
result = this.asCall().toString()
88+
result = [this.asCall().toString(), this.asAwaitExpr().toString()]
8589
or
8690
exists(Expr e, DerefChain derefChain, int i |
8791
this.isImplicitDerefCall(e, derefChain, i, _) and
@@ -97,7 +101,7 @@ final class DataFlowCall extends TDataFlowCall {
97101
}
98102

99103
Location getLocation() {
100-
result = this.asCall().getLocation()
104+
result = [this.asCall().getLocation(), this.asAwaitExpr().getLocation()]
101105
or
102106
result = any(Expr e | this.isImplicitDerefCall(e, _, _, _)).getLocation()
103107
}
@@ -1004,7 +1008,11 @@ module RustDataFlowGen<RustDataFlowInputSig Input> implements InputSig<Location>
10041008
*/
10051009
predicate lambdaCall(DataFlowCall call, LambdaCallKind kind, Node receiver) {
10061010
(
1007-
receiver.asExpr() = call.asCall().(CallExprImpl::DynamicCallExpr).getFunction()
1011+
receiver.asExpr() =
1012+
[
1013+
call.asCall().(CallExprImpl::DynamicCallExpr).getFunction(),
1014+
call.asAwaitExpr().getExpr()
1015+
]
10081016
or
10091017
call.isSummaryCall(_, receiver.(FlowSummaryNode).getSummaryNode())
10101018
) and
@@ -1176,6 +1184,7 @@ private module Cached {
11761184
Stages::DataFlowStage::ref() and
11771185
call.hasEnclosingCfgScope()
11781186
} or
1187+
TAwaitExpr(AwaitExpr await) { await.hasEnclosingCfgScope() } or
11791188
TImplicitDerefCall(Expr e, DerefChain derefChain, int i, Function target) {
11801189
TypeInference::implicitDerefChainBorrow(e, derefChain, _) and
11811190
target = derefChain.getElement(i).getDerefFunction() and

‎rust/ql/lib/codeql/rust/dataflow/internal/Node.qll‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -435,6 +435,20 @@ final class ClosureArgumentNode extends ArgumentNode, ExprNode {
435435
}
436436
}
437437

438+
/**
439+
* A data flow node that represents the run-time representation of an async
440+
* block passed into its body when awaited.
441+
*/
442+
final class AsyncBlockArgumentNode extends ArgumentNode, ExprNode {
443+
private AwaitExpr await;
444+
445+
AsyncBlockArgumentNode() { this.asExpr() = await.getExpr() }
446+
447+
override predicate isArgumentOf(DataFlowCall call, RustDataFlow::ArgumentPosition pos) {
448+
call.asAwaitExpr() = await and pos.isClosureSelf()
449+
}
450+
}
451+
438452
/** An SSA node. */
439453
class SsaNode extends Node, TSsaNode {
440454
SsaImpl::DataFlowIntegration::SsaNode node;
@@ -485,12 +499,17 @@ final private class ExprOutNode extends ExprNode, OutNode {
485499
not call instanceof DerefExpr and // Handled by `DerefOutNode`
486500
not call instanceof IndexExpr // Handled by `IndexOutNode`
487501
)
502+
or
503+
this.asExpr() instanceof AwaitExpr
488504
}
489505

490506
/** Gets the underlying call node that includes this out node. */
491507
override DataFlowCall getCall(ReturnKind kind) {
492508
result.asCall() = n and
493509
kind = TNormalReturnKind()
510+
or
511+
result.asAwaitExpr() = n and
512+
kind = TNormalReturnKind()
494513
}
495514
}
496515

‎rust/ql/test/library-tests/dataflow/lambdas/inline-flow.expected‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,14 @@ edges
6363
| main.rs:111:26:111:36 | \|...\| ... : ... => .. [captured a] | main.rs:111:35:111:35 | a | provenance | heuristic-callback |
6464
| main.rs:112:26:112:36 | \|...\| ... : ... => .. [captured a] | main.rs:112:35:112:35 | a | provenance | heuristic-callback |
6565
| main.rs:113:26:113:36 | \|...\| ... : ... => .. [captured a] | main.rs:113:35:113:35 | a | provenance | heuristic-callback |
66+
| main.rs:120:30:120:40 | source(...) | main.rs:121:14:121:25 | await future | provenance | |
67+
| main.rs:125:20:125:30 | source(...) | main.rs:126:9:128:9 | { ... } : ... => .. [captured data] | provenance | |
68+
| main.rs:126:9:128:9 | { ... } : ... => .. [captured data] | main.rs:127:18:127:21 | data | provenance | |
69+
| main.rs:133:20:133:30 | source(...) | main.rs:135:14:135:19 | future [captured data] | provenance | |
70+
| main.rs:135:14:135:19 | future [captured data] | main.rs:134:30:134:33 | data | provenance | |
71+
| main.rs:135:14:135:19 | future [captured data] | main.rs:135:14:135:25 | await future | provenance | |
72+
| main.rs:141:9:143:9 | [post] { ... } : ... => .. [captured captured] | main.rs:145:14:145:21 | captured | provenance | |
73+
| main.rs:142:24:142:34 | source(...) | main.rs:141:9:143:9 | [post] { ... } : ... => .. [captured captured] | provenance | |
6674
nodes
6775
| main.rs:10:20:10:52 | if cond {...} else {...} | semmle.label | if cond {...} else {...} |
6876
| main.rs:10:30:10:39 | source(...) | semmle.label | source(...) |
@@ -138,9 +146,22 @@ nodes
138146
| main.rs:112:35:112:35 | a | semmle.label | a |
139147
| main.rs:113:26:113:36 | \|...\| ... : ... => .. [captured a] | semmle.label | \|...\| ... : ... => .. [captured a] |
140148
| main.rs:113:35:113:35 | a | semmle.label | a |
149+
| main.rs:120:30:120:40 | source(...) | semmle.label | source(...) |
150+
| main.rs:121:14:121:25 | await future | semmle.label | await future |
151+
| main.rs:125:20:125:30 | source(...) | semmle.label | source(...) |
152+
| main.rs:126:9:128:9 | { ... } : ... => .. [captured data] | semmle.label | { ... } : ... => .. [captured data] |
153+
| main.rs:127:18:127:21 | data | semmle.label | data |
154+
| main.rs:133:20:133:30 | source(...) | semmle.label | source(...) |
155+
| main.rs:134:30:134:33 | data | semmle.label | data |
156+
| main.rs:135:14:135:19 | future [captured data] | semmle.label | future [captured data] |
157+
| main.rs:135:14:135:25 | await future | semmle.label | await future |
158+
| main.rs:141:9:143:9 | [post] { ... } : ... => .. [captured captured] | semmle.label | [post] { ... } : ... => .. [captured captured] |
159+
| main.rs:142:24:142:34 | source(...) | semmle.label | source(...) |
160+
| main.rs:145:14:145:21 | captured | semmle.label | captured |
141161
subpaths
142162
| main.rs:29:21:29:21 | a | main.rs:27:20:27:23 | ... | main.rs:27:26:27:52 | if cond {...} else {...} | main.rs:29:13:29:22 | f(...) |
143163
| main.rs:77:21:77:21 | a | main.rs:66:24:66:32 | ...: i64 | main.rs:66:42:72:1 | { ... } | main.rs:77:13:77:22 | f(...) |
164+
| main.rs:135:14:135:19 | future [captured data] | main.rs:134:30:134:33 | data | main.rs:134:30:134:33 | data | main.rs:135:14:135:25 | await future |
144165
testFailures
145166
#select
146167
| main.rs:11:10:11:16 | f(...) | main.rs:10:30:10:39 | source(...) | main.rs:11:10:11:16 | f(...) | $@ | main.rs:10:30:10:39 | source(...) | source(...) |
@@ -161,3 +182,7 @@ testFailures
161182
| main.rs:111:35:111:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:111:35:111:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
162183
| main.rs:112:35:112:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:112:35:112:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
163184
| main.rs:113:35:113:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:113:35:113:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
185+
| main.rs:121:14:121:25 | await future | main.rs:120:30:120:40 | source(...) | main.rs:121:14:121:25 | await future | $@ | main.rs:120:30:120:40 | source(...) | source(...) |
186+
| main.rs:127:18:127:21 | data | main.rs:125:20:125:30 | source(...) | main.rs:127:18:127:21 | data | $@ | main.rs:125:20:125:30 | source(...) | source(...) |
187+
| main.rs:135:14:135:25 | await future | main.rs:133:20:133:30 | source(...) | main.rs:135:14:135:25 | await future | $@ | main.rs:133:20:133:30 | source(...) | source(...) |
188+
| main.rs:145:14:145:21 | captured | main.rs:142:24:142:34 | source(...) | main.rs:145:14:145:21 | captured | $@ | main.rs:142:24:142:34 | source(...) | source(...) |

‎rust/ql/test/library-tests/dataflow/lambdas/main.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -118,21 +118,21 @@ mod async_blocks {
118118

119119
async fn async_block_flow_out() {
120120
let future = async { source(206) };
121-
sink(future.await); // $ MISSING: hasValueFlow=206
121+
sink(future.await); // $ hasValueFlow=206
122122
}
123123

124124
async fn async_block_flow_in() {
125125
let data = source(207);
126126
async {
127-
sink(data); // $ MISSING: hasValueFlow=207
127+
sink(data); // $ hasValueFlow=207
128128
}
129129
.await;
130130
}
131131

132132
async fn async_block_flow_through() {
133133
let data = source(208);
134134
let future = async { data };
135-
sink(future.await); // $ MISSING: hasValueFlow=208
135+
sink(future.await); // $ hasValueFlow=208
136136
}
137137

138138
async fn async_block_captured_variable() {
@@ -142,7 +142,7 @@ mod async_blocks {
142142
captured = source(209);
143143
}
144144
.await;
145-
sink(captured); // $ MISSING: hasValueFlow=209
145+
sink(captured); // $ hasValueFlow=209
146146
}
147147

148148
async fn boxed_async_block_captured_variable() {

0 commit comments

Comments
 (0)