Skip to content

Commit 4dcf801

Browse files
Jim8yJim8y
committed
Fix write-in-try analyzer for local storage syscalls
1 parent 8a2e070 commit 4dcf801

File tree

3 files changed

+111
-27
lines changed

3 files changed

+111
-27
lines changed

‎src/Neo.Compiler.CSharp/Optimizer/Analysers/TryCatchFinallyCoverage.cs‎

Lines changed: 79 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,8 @@ public TryCatchFinallySingleCoverage(ContractInBasicBlocks contractInBasicBlocks
7777

7878
public class TryCatchFinallyCoverage
7979
{
80+
private static readonly bool EnableDiagnostics = Environment.GetEnvironmentVariable("TRY_COVERAGE_DEBUG") == "1";
81+
8082
public ContractInBasicBlocks contractInBasicBlocks { get; protected set; }
8183
// key: start of try block. prevBlock.instructions.Last() is TRY
8284
public Dictionary<BasicBlock, TryCatchFinallySingleCoverage> allTry { get; protected set; }
@@ -103,7 +105,7 @@ public TryCatchFinallyCoverage(ContractInBasicBlocks contractInBasicBlocks)
103105
{
104106
Stack<(BasicBlock tryBlock, BasicBlock? endFinallyBlock, TryType tryType, bool continueAfterFinally)> tryStack = new();
105107
tryStack.Push((b, null, TryType.TRY, true));
106-
CoverSingleTry(b, tryStack);
108+
ContinueToBlock(b, tryStack);
107109
}
108110
}
109111

@@ -113,16 +115,17 @@ public BranchType CoverSingleTry(BasicBlock currentBlock,
113115
if (tryStack.Count <= 0)
114116
return BranchType.OK;
115117
tryStack = InstructionCoverage.CopyStack(tryStack);
116-
(BasicBlock tryBlock, BasicBlock? endFinallyBlock, TryType tryType, bool continueAfterFinally) = tryStack.Peek();
117-
HashSet<BasicBlock> handledBlocks = tryType switch
118-
{
119-
TryType.TRY => allTry[tryBlock].tryBlocks,
120-
TryType.CATCH => allTry[tryBlock].catchBlocks,
121-
TryType.FINALLY => allTry[tryBlock].finallyBlocks,
122-
_ => throw new ArgumentException($"Invalid {nameof(tryType)} {tryType}"),
123-
};
124118
while (true)
125119
{
120+
(BasicBlock tryBlock, BasicBlock? endFinallyBlock, TryType tryType, bool continueAfterFinally) = tryStack.Peek();
121+
HashSet<BasicBlock> handledBlocks = tryType switch
122+
{
123+
TryType.TRY => allTry[tryBlock].tryBlocks,
124+
TryType.CATCH => allTry[tryBlock].catchBlocks,
125+
TryType.FINALLY => allTry[tryBlock].finallyBlocks,
126+
_ => throw new ArgumentException($"Invalid {nameof(tryType)} {tryType}"),
127+
};
128+
126129
if (handledBlocks.Contains(currentBlock))
127130
return currentBlock.branchType;
128131
handledBlocks.Add(currentBlock);
@@ -135,19 +138,19 @@ public BranchType CoverSingleTry(BasicBlock currentBlock,
135138
foreach (int callaTarget in contractInBasicBlocks.coverage.pushaTargets.Keys)
136139
{
137140
BasicBlock callABlock = contractInBasicBlocks.basicBlocksByStartAddr[callaTarget];
138-
CoverSingleTry(callABlock, tryStack);
141+
ContinueToBlock(callABlock, tryStack);
139142
// TODO: if a PUSHA cannot be covered, do not add it as a CALLA target
140143
}
141144
else
142145
{
143146
int callTarget = ComputeJumpTarget(currentBlock.lastAddr, instruction);
144147
BasicBlock calledBlock = contractInBasicBlocks.basicBlocksByStartAddr[callTarget];
145-
CoverSingleTry(calledBlock, tryStack);
148+
ContinueToBlock(calledBlock, tryStack);
146149
}
147150
if (currentBlock.branchType == BranchType.OK)
148151
if (currentBlock.nextBlock != null)
149152
// nextBlock can still be null, when we call a method in try that ABORTs
150-
return CoverSingleTry(currentBlock.nextBlock!, tryStack);
153+
return ContinueToBlock(currentBlock.nextBlock!, tryStack);
151154
return currentBlock.branchType;
152155
}
153156
if (instruction.OpCode == OpCode.RET)
@@ -156,6 +159,8 @@ public BranchType CoverSingleTry(BasicBlock currentBlock,
156159
{
157160
if (instruction.OpCode == OpCode.TRY || instruction.OpCode == OpCode.TRY_L)
158161
{
162+
if (EnableDiagnostics)
163+
Console.Error.WriteLine($"[TryCoverage] TRY in block {currentBlock.startAddr}, state={tryType}, stack={FormatStack(tryStack)}");
159164
HashSet<TryCatchFinallySingleCoverage> handledCoverage = tryType switch
160165
{
161166
TryType.TRY => allTry[tryBlock].nestedTrysInTry,
@@ -165,7 +170,9 @@ public BranchType CoverSingleTry(BasicBlock currentBlock,
165170
};
166171
handledCoverage.Add(allTry[currentBlock.nextBlock!]);
167172
tryStack.Push((currentBlock.nextBlock!, null, TryType.TRY, true));
168-
CoverSingleTry(currentBlock.nextBlock!, tryStack);
173+
if (EnableDiagnostics)
174+
Console.Error.WriteLine($"[TryCoverage] push TRY block {currentBlock.nextBlock!.startAddr}, stack={FormatStack(tryStack)}");
175+
ContinueToBlock(currentBlock.nextBlock!, tryStack);
169176
}
170177
if (instruction.OpCode == OpCode.THROW)
171178
{
@@ -179,59 +186,81 @@ public BranchType CoverSingleTry(BasicBlock currentBlock,
179186
if (allTry[prevTryBlock].catchBlock != null)
180187
{
181188
tryStack.Push((prevTryBlock, null, TryType.CATCH, true));
182-
return CoverSingleTry(allTry[prevTryBlock].catchBlock!, tryStack);
189+
return ContinueToBlock(allTry[prevTryBlock].catchBlock!, tryStack);
183190
}
184191
else if (allTry[prevTryBlock].finallyBlock != null)
185192
{
186193
tryStack.Push((prevTryBlock, null, TryType.FINALLY, false));
187-
if (CoverSingleTry(allTry[prevTryBlock].finallyBlock!, tryStack) == BranchType.ABORT)
194+
if (ContinueToBlock(allTry[prevTryBlock].finallyBlock!, tryStack) == BranchType.ABORT)
188195
return BranchType.ABORT;
189196
return BranchType.THROW;
190197
}
191198
}
192199
if (tryType == TryType.CATCH && allTry[tryBlock].finallyBlock != null)
193200
{
194201
tryStack.Push((prevTryBlock, null, TryType.FINALLY, false));
195-
if (CoverSingleTry(allTry[prevTryBlock].finallyBlock!, tryStack) == BranchType.ABORT)
202+
if (ContinueToBlock(allTry[prevTryBlock].finallyBlock!, tryStack) == BranchType.ABORT)
196203
return BranchType.ABORT;
197204
}
198205
return BranchType.THROW;
199206
}
200207
if (instruction.OpCode == OpCode.ENDTRY || instruction.OpCode == OpCode.ENDTRY_L)
201208
{
209+
if (EnableDiagnostics)
210+
Console.Error.WriteLine($"[TryCoverage] ENDTRY in block {currentBlock.startAddr}, state={tryType}, stack={FormatStack(tryStack)}");
202211
if (tryType != TryType.TRY && tryType != TryType.CATCH)
203-
throw new BadScriptException("No try stack on ENDTRY");
212+
{
213+
if (EnableDiagnostics)
214+
{
215+
Console.Error.WriteLine($"[TryCoverage] Unexpected ENDTRY without TRY/CATCH context. state={tryType}, block={currentBlock.startAddr}, tryBlock={tryBlock.startAddr}");
216+
foreach (var inst in currentBlock.instructions)
217+
Console.Error.WriteLine($" {inst.OpCode}");
218+
}
219+
throw new BadScriptException($"No try stack on ENDTRY (state={tryType}, block={currentBlock.startAddr}, tryBlock={tryBlock.startAddr})");
220+
}
204221
tryStack.Pop(); // pop the ending TRY or CATCH
222+
if (EnableDiagnostics)
223+
Console.Error.WriteLine($"[TryCoverage] pop after ENDTRY, stack={FormatStack(tryStack)}");
205224
if (tryType == TryType.TRY && allTry[tryBlock].catchBlock != null)
206225
{
207226
tryStack.Push((tryBlock, null, TryType.CATCH, true));
208-
CoverSingleTry(allTry[tryBlock].catchBlock!, tryStack);
227+
if (EnableDiagnostics)
228+
Console.Error.WriteLine($"[TryCoverage] push CATCH block {allTry[tryBlock].catchBlock!.startAddr}, stack={FormatStack(tryStack)}");
229+
ContinueToBlock(allTry[tryBlock].catchBlock!, tryStack);
209230
tryStack.Pop(); // Pop the CATCH
231+
if (EnableDiagnostics)
232+
Console.Error.WriteLine($"[TryCoverage] pop CATCH, stack={FormatStack(tryStack)}");
210233
}
211234
int endPointer = ComputeJumpTarget(currentBlock.lastAddr, instruction);
212235
endFinallyBlock = contractInBasicBlocks.basicBlocksByStartAddr[endPointer];
213236
BasicBlock nextBlock;
214237
if (allTry[tryBlock].finallyBlock != null)
215238
{
216239
tryStack.Push(new(tryBlock, endFinallyBlock, TryType.FINALLY, true));
240+
if (EnableDiagnostics)
241+
Console.Error.WriteLine($"[TryCoverage] push FINALLY block {allTry[tryBlock].finallyBlock!.startAddr}, end={endFinallyBlock.startAddr}, stack={FormatStack(tryStack)}");
217242
nextBlock = allTry[tryBlock].finallyBlock!;
218243
}
219244
else
220245
{
221246
allTry[tryBlock].endingBlocks.Add(endFinallyBlock);
222247
nextBlock = endFinallyBlock;
223248
}
224-
return CoverSingleTry(nextBlock, tryStack);
249+
return ContinueToBlock(nextBlock, tryStack);
225250
}
226251
if (instruction.OpCode == OpCode.ENDFINALLY)
227252
{
253+
if (EnableDiagnostics)
254+
Console.Error.WriteLine($"[TryCoverage] ENDFINALLY in block {currentBlock.startAddr}, state={tryType}, stack={FormatStack(tryStack)}");
228255
if (tryType != TryType.FINALLY)
229256
throw new BadScriptException("No finally stack on ENDFINALLY");
230257
tryStack.Pop(); // pop the ending FINALLY
258+
if (EnableDiagnostics)
259+
Console.Error.WriteLine($"[TryCoverage] pop FINALLY, stack={FormatStack(tryStack)}");
231260
if (continueAfterFinally)
232261
{
233262
allTry[tryBlock].endingBlocks.Add(endFinallyBlock!);
234-
return CoverSingleTry(endFinallyBlock!, tryStack);
263+
return ContinueToBlock(endFinallyBlock!, tryStack);
235264
}
236265
// For this basic block in finally, the branch type is OK
237266
// The throw is caused by previous codes
@@ -242,14 +271,14 @@ public BranchType CoverSingleTry(BasicBlock currentBlock,
242271
{
243272
int target = ComputeJumpTarget(currentBlock.lastAddr, instruction);
244273
BasicBlock targetBlock = contractInBasicBlocks.basicBlocksByStartAddr[target];
245-
return CoverSingleTry(targetBlock, tryStack);
274+
return ContinueToBlock(targetBlock, tryStack);
246275
}
247276
if (conditionalJump.Contains(instruction.OpCode) || conditionalJump_L.Contains(instruction.OpCode))
248277
{
249-
BranchType noJump = CoverSingleTry(currentBlock.nextBlock!, tryStack);
278+
BranchType noJump = ContinueToBlock(currentBlock.nextBlock!, tryStack);
250279
int target = ComputeJumpTarget(currentBlock.lastAddr, instruction);
251280
BasicBlock targetBlock = contractInBasicBlocks.basicBlocksByStartAddr[target];
252-
BranchType jump = CoverSingleTry(targetBlock, tryStack);
281+
BranchType jump = ContinueToBlock(targetBlock, tryStack);
253282
if (noJump == BranchType.OK || jump == BranchType.OK)
254283
return BranchType.OK;
255284
if (noJump == BranchType.ABORT && jump == BranchType.ABORT)
@@ -260,5 +289,32 @@ public BranchType CoverSingleTry(BasicBlock currentBlock,
260289
currentBlock = currentBlock.nextBlock!;
261290
}
262291
}
292+
293+
private static string FormatStack(Stack<(BasicBlock tryBlock, BasicBlock? endFinallyBlock, TryType tryType, bool continueAfterFinally)> stack)
294+
=> string.Join(" | ", stack.Select(s => $"[{s.tryBlock.startAddr}:{s.tryType}:{(s.endFinallyBlock?.startAddr.ToString() ?? "-")}:cont={s.continueAfterFinally}]"));
295+
296+
private BranchType ContinueToBlock(
297+
BasicBlock targetBlock,
298+
Stack<(BasicBlock tryBlock, BasicBlock? endFinallyBlock, TryType tryType, bool continueAfterFinally)> tryStack)
299+
{
300+
if (EnableDiagnostics)
301+
Console.Error.WriteLine($"[TryCoverage] Continue to block {targetBlock.startAddr}, stack={FormatStack(tryStack)}");
302+
if (tryStack.Count > 0)
303+
{
304+
(BasicBlock tryBlock, BasicBlock? endFinallyBlock, TryType tryType, bool continueAfterFinally) = tryStack.Peek();
305+
if (tryType == TryType.FINALLY && endFinallyBlock != null && ReferenceEquals(targetBlock, endFinallyBlock))
306+
{
307+
Stack<(BasicBlock tryBlock, BasicBlock? endFinallyBlock, TryType tryType, bool continueAfterFinally)> stackAfterFinally = InstructionCoverage.CopyStack(tryStack);
308+
stackAfterFinally.Pop();
309+
if (EnableDiagnostics)
310+
Console.Error.WriteLine($"[TryCoverage] transition after finally end, stack={FormatStack(stackAfterFinally)}");
311+
if (!continueAfterFinally)
312+
return BranchType.OK;
313+
allTry[tryBlock].endingBlocks.Add(endFinallyBlock);
314+
return CoverSingleTry(targetBlock, stackAfterFinally);
315+
}
316+
}
317+
return CoverSingleTry(targetBlock, tryStack);
318+
}
263319
}
264320
}

‎src/Neo.Compiler.CSharp/SecurityAnalyzer/WriteInTryAnalyzer.cs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,14 @@ namespace Neo.Compiler.SecurityAnalyzer
2323
{
2424
public static class WriteInTryAnalzyer
2525
{
26+
private static readonly HashSet<uint> StorageWriteHashes = new()
27+
{
28+
ApplicationEngine.System_Storage_Put.Hash,
29+
ApplicationEngine.System_Storage_Delete.Hash,
30+
ApplicationEngine.System_Storage_Local_Put.Hash,
31+
ApplicationEngine.System_Storage_Local_Delete.Hash
32+
};
33+
2634
public class WriteInTryVulnerability
2735
{
2836
// key block writes storage; value blocks in try
@@ -63,8 +71,7 @@ public string GetWarningInfo(bool print = false)
6371
foreach (VM.Instruction i in b.instructions)
6472
{
6573
if (i.OpCode == VM.OpCode.SYSCALL
66-
&& (i.TokenU32 == ApplicationEngine.System_Storage_Put.Hash
67-
|| i.TokenU32 == ApplicationEngine.System_Storage_Delete.Hash))
74+
&& StorageWriteHashes.Contains(i.TokenU32))
6875
writeAddrs.Add(a);
6976
a += i.Size;
7077
}
@@ -133,8 +140,7 @@ public static WriteInTryVulnerability AnalyzeWriteInTry
133140
foreach (BasicBlock block in contractInBasicBlocks.sortedBasicBlocks)
134141
foreach (VM.Instruction i in block.instructions)
135142
if (i.OpCode == VM.OpCode.SYSCALL
136-
&& (i.TokenU32 == ApplicationEngine.System_Storage_Put.Hash
137-
|| i.TokenU32 == ApplicationEngine.System_Storage_Delete.Hash))
143+
&& StorageWriteHashes.Contains(i.TokenU32))
138144
allBasicBlocksWritingStorage.Add(block);
139145
foreach (TryCatchFinallySingleCoverage c in tryCatchFinallyCoverage.allTry.Values)
140146
{

‎tests/Neo.Compiler.CSharp.UnitTests/SecurityAnalyzer/UnitTest_WriteInTry.cs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,13 @@
1010
// modifications are permitted.
1111

1212
using Microsoft.VisualStudio.TestTools.UnitTesting;
13+
using System.Linq;
1314
using Neo.Compiler.SecurityAnalyzer;
1415
using Neo.Json;
1516
using Neo.Optimizer;
17+
using Neo.SmartContract;
1618
using Neo.SmartContract.Testing;
19+
using Neo.VM;
1720

1821
namespace Neo.Compiler.CSharp.UnitTests.SecurityAnalyzer
1922
{
@@ -72,5 +75,24 @@ public void Test_WriteInTryWithEnhancedDiagnostics()
7275
// Message should be more detailed than just addresses
7376
Assert.IsTrue(warningInfo.Length > 150, "Enhanced diagnostic message should be more detailed than simple address listing");
7477
}
78+
79+
[TestMethod]
80+
public void Test_WriteInTryDetectsLocalStorageSyscalls()
81+
{
82+
WriteInTryAnalzyer.WriteInTryVulnerability v =
83+
WriteInTryAnalzyer.AnalyzeWriteInTry(NefFile, Manifest);
84+
85+
Assert.AreEqual(2, v.vulnerabilities.Count);
86+
87+
foreach (var block in v.vulnerabilities.Keys)
88+
{
89+
bool hasLocalStorageWrite = block.instructions.Any(instruction =>
90+
instruction.OpCode == OpCode.SYSCALL &&
91+
(instruction.TokenU32 == ApplicationEngine.System_Storage_Local_Put.Hash
92+
|| instruction.TokenU32 == ApplicationEngine.System_Storage_Local_Delete.Hash));
93+
94+
Assert.IsTrue(hasLocalStorageWrite, "Expected vulnerability block to contain a local storage write syscall");
95+
}
96+
}
7597
}
7698
}

0 commit comments

Comments
 (0)