Skip to content

Commit a9db408

Browse files
lokitothlokitoth
authored
fix: FS Checkpoint storage special character support (microsoft#4730)
The `sessionId`, an optional parameter when starting a new session when running a workflow is an arbitrary string. This allows consumers to support whatever ids are needed by other systems, but can result in errors when an OS special or forbidden character is included. The fix is to escape the paths, in a 1:1 manner. We rely on EncodeDataString to do this. * Also modifies the index file to make it easier to determine what the name of the file on disk is for a given `sessionId`.
1 parent 87962e5 commit a9db408

3 files changed

Lines changed: 185 additions & 43 deletions

File tree

‎dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/FileSystemJsonCheckpointStore.cs‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,14 @@
55
using System.IO;
66
using System.Text;
77
using System.Text.Json;
8+
using System.Text.Json.Serialization.Metadata;
89
using System.Threading;
910
using System.Threading.Tasks;
1011

1112
namespace Microsoft.Agents.AI.Workflows.Checkpointing;
1213

14+
internal record CheckpointFileIndexEntry(CheckpointInfo CheckpointInfo, string FileName);
15+
1316
/// <summary>
1417
/// Provides a file system-based implementation of a JSON checkpoint store that persists checkpoint data and index
1518
/// information to disk using JSON files.
@@ -28,6 +31,8 @@ public sealed class FileSystemJsonCheckpointStore : JsonCheckpointStore, IDispos
2831
internal DirectoryInfo Directory { get; }
2932
internal HashSet<CheckpointInfo> CheckpointIndex { get; }
3033

34+
private static JsonTypeInfo<CheckpointFileIndexEntry> EntryTypeInfo => WorkflowsJsonUtilities.JsonContext.Default.CheckpointFileIndexEntry;
35+
3136
/// <summary>
3237
/// Initializes a new instance of the <see cref="FileSystemJsonCheckpointStore"/> class that uses the specified directory
3338
/// </summary>
@@ -64,9 +69,11 @@ public FileSystemJsonCheckpointStore(DirectoryInfo directory)
6469
using StreamReader reader = new(this._indexFile, encoding: Encoding.UTF8, detectEncodingFromByteOrderMarks: false, BufferSize, leaveOpen: true);
6570
while (reader.ReadLine() is string line)
6671
{
67-
if (JsonSerializer.Deserialize(line, KeyTypeInfo) is { } info)
72+
if (JsonSerializer.Deserialize(line, EntryTypeInfo) is { } entry)
6873
{
69-
this.CheckpointIndex.Add(info);
74+
// We never actually use the file names from the index entries since they can be derived from the CheckpointInfo, but it is useful to
75+
// have the UrlEncoded file names in the index file for human readability
76+
this.CheckpointIndex.Add(entry.CheckpointInfo);
7077
}
7178
}
7279
}
@@ -93,8 +100,14 @@ private void CheckDisposed()
93100
}
94101
}
95102

96-
private string GetFileNameForCheckpoint(string sessionId, CheckpointInfo key)
97-
=> Path.Combine(this.Directory.FullName, $"{sessionId}_{key.CheckpointId}.json");
103+
internal string GetFileNameForCheckpoint(string sessionId, CheckpointInfo key)
104+
{
105+
string protoPath = $"{sessionId}_{key.CheckpointId}.json";
106+
107+
// Escape the protoPath to ensure it is a valid file name, especially if sessionId or CheckpointId contain path separators, etc.
108+
return Uri.EscapeDataString(protoPath) // This takes care of most of the invalid path characters
109+
.Replace(".", "%2E"); // This takes care of escaping the root folder, since EscapeDataString does not escape dots
110+
}
98111

99112
private CheckpointInfo GetUnusedCheckpointInfo(string sessionId)
100113
{
@@ -116,13 +129,16 @@ public override async ValueTask<CheckpointInfo> CreateCheckpointAsync(string ses
116129

117130
CheckpointInfo key = this.GetUnusedCheckpointInfo(sessionId);
118131
string fileName = this.GetFileNameForCheckpoint(sessionId, key);
132+
string filePath = Path.Combine(this.Directory.FullName, fileName);
133+
119134
try
120135
{
121-
using Stream checkpointStream = File.Open(fileName, FileMode.Create, FileAccess.Write, FileShare.None);
136+
using Stream checkpointStream = File.Open(filePath, FileMode.Create, FileAccess.Write, FileShare.None);
122137
using Utf8JsonWriter jsonWriter = new(checkpointStream, new JsonWriterOptions() { Indented = false });
123138
value.WriteTo(jsonWriter);
124139

125-
JsonSerializer.Serialize(this._indexFile!, key, KeyTypeInfo);
140+
CheckpointFileIndexEntry entry = new(key, fileName);
141+
JsonSerializer.Serialize(this._indexFile!, entry, EntryTypeInfo);
126142
byte[] bytes = Encoding.UTF8.GetBytes(Environment.NewLine);
127143
await this._indexFile!.WriteAsync(bytes, 0, bytes.Length, CancellationToken.None).ConfigureAwait(false);
128144
await this._indexFile!.FlushAsync(CancellationToken.None).ConfigureAwait(false);
@@ -136,7 +152,7 @@ public override async ValueTask<CheckpointInfo> CreateCheckpointAsync(string ses
136152
try
137153
{
138154
// try to clean up after ourselves
139-
File.Delete(fileName);
155+
File.Delete(filePath);
140156
}
141157
catch { }
142158

@@ -149,14 +165,15 @@ public override async ValueTask<JsonElement> RetrieveCheckpointAsync(string sess
149165
{
150166
this.CheckDisposed();
151167
string fileName = this.GetFileNameForCheckpoint(sessionId, key);
168+
string filePath = Path.Combine(this.Directory.FullName, fileName);
152169

153170
if (!this.CheckpointIndex.Contains(key) ||
154171
!File.Exists(fileName))
155172
{
156173
throw new KeyNotFoundException($"Checkpoint '{key.CheckpointId}' not found in store at '{this.Directory.FullName}'.");
157174
}
158175

159-
using FileStream checkpointFileStream = File.Open(fileName, FileMode.Open, FileAccess.Read, FileShare.Read);
176+
using FileStream checkpointFileStream = File.Open(filePath, FileMode.Open, FileAccess.Read, FileShare.Read);
160177
using JsonDocument document = await JsonDocument.ParseAsync(checkpointFileStream).ConfigureAwait(false);
161178

162179
return document.RootElement.Clone();

‎dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ private static JsonSerializerOptions CreateDefaultOptions()
7171
[JsonSerializable(typeof(PortableValue))]
7272
[JsonSerializable(typeof(PortableMessageEnvelope))]
7373
[JsonSerializable(typeof(InMemoryCheckpointManager))]
74+
[JsonSerializable(typeof(CheckpointFileIndexEntry))]
7475

7576
// Runtime State Types
7677
[JsonSerializable(typeof(ScopeKey))]

‎dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/FileSystemJsonCheckpointStoreTests.cs‎

Lines changed: 159 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -9,49 +9,173 @@
99

1010
namespace Microsoft.Agents.AI.Workflows.UnitTests;
1111

12-
public sealed class FileSystemJsonCheckpointStoreTests
12+
internal sealed class TempDirectory : IDisposable
1313
{
14-
[Fact]
15-
public async Task CreateCheckpointAsync_ShouldPersistIndexToDiskBeforeDisposeAsync()
14+
public DirectoryInfo DirectoryInfo { get; }
15+
16+
public TempDirectory()
1617
{
17-
// Arrange
18-
DirectoryInfo tempDir = new(Path.Combine(Path.GetTempPath(), Guid.NewGuid().ToString()));
19-
FileSystemJsonCheckpointStore? store = null;
18+
string tempDirPath = Path.Combine(Path.GetTempPath(), Guid.NewGuid().ToString());
19+
this.DirectoryInfo = Directory.CreateDirectory(tempDirPath);
20+
}
21+
22+
public void Dispose()
23+
{
24+
this.DisposeInternal();
25+
GC.SuppressFinalize(this);
26+
}
27+
28+
private void DisposeInternal()
29+
{
30+
if (this.DirectoryInfo.Exists)
31+
{
32+
try
33+
{
34+
// Best efforts
35+
this.DirectoryInfo.Delete(recursive: true);
36+
}
37+
catch { }
38+
}
39+
}
40+
41+
~TempDirectory()
42+
{
43+
// Best efforts
44+
this.DisposeInternal();
45+
}
46+
47+
public static implicit operator DirectoryInfo(TempDirectory tempDirectory) => tempDirectory.DirectoryInfo;
48+
49+
public string FullName => this.DirectoryInfo.FullName;
2050

21-
try
51+
public bool IsParentOf(FileInfo candidate)
52+
{
53+
if (candidate.Directory is null)
2254
{
23-
store = new(tempDir);
24-
string runId = Guid.NewGuid().ToString("N");
25-
JsonElement testData = JsonSerializer.SerializeToElement(new { test = "data" });
26-
27-
// Act
28-
CheckpointInfo checkpoint = await store.CreateCheckpointAsync(runId, testData);
29-
30-
// Assert - Check the file size before disposing to verify data was flushed to disk
31-
// The index.jsonl file is held exclusively by the store, so we check via FileInfo
32-
string indexPath = Path.Combine(tempDir.FullName, "index.jsonl");
33-
FileInfo indexFile = new(indexPath);
34-
indexFile.Refresh();
35-
long fileSizeBeforeDispose = indexFile.Length;
36-
37-
// Data should already be on disk (file size > 0) before we dispose
38-
fileSizeBeforeDispose.Should().BeGreaterThan(0, "index.jsonl should be flushed to disk after CreateCheckpointAsync");
39-
40-
// Dispose to release file lock before final verification
41-
store.Dispose();
42-
store = null;
43-
44-
string[] lines = File.ReadAllLines(indexPath);
45-
lines.Should().HaveCount(1);
46-
lines[0].Should().Contain(checkpoint.CheckpointId);
55+
return false;
4756
}
48-
finally
57+
58+
if (candidate.Directory.FullName == this.DirectoryInfo.FullName)
59+
{
60+
return true;
61+
}
62+
63+
return this.IsParentOf(candidate.Directory);
64+
}
65+
66+
public bool IsParentOf(DirectoryInfo candidate)
67+
{
68+
while (candidate.Parent is not null)
4969
{
50-
store?.Dispose();
51-
if (tempDir.Exists)
70+
if (candidate.Parent.FullName == this.DirectoryInfo.FullName)
5271
{
53-
tempDir.Delete(recursive: true);
72+
return true;
5473
}
74+
75+
candidate = candidate.Parent;
5576
}
77+
78+
return false;
79+
}
80+
}
81+
public sealed class FileSystemJsonCheckpointStoreTests
82+
{
83+
public static JsonElement TestData => JsonSerializer.SerializeToElement(new { test = "data" });
84+
85+
[Fact]
86+
public async Task CreateCheckpointAsync_ShouldPersistIndexToDiskBeforeDisposeAsync()
87+
{
88+
// Arrange
89+
using TempDirectory tempDirectory = new();
90+
using FileSystemJsonCheckpointStore? store = new(tempDirectory);
91+
92+
string runId = Guid.NewGuid().ToString("N");
93+
94+
// Act
95+
CheckpointInfo checkpoint = await store.CreateCheckpointAsync(runId, TestData);
96+
97+
// Assert - Check the file size before disposing to verify data was flushed to disk
98+
// The index.jsonl file is held exclusively by the store, so we check via FileInfo
99+
string indexPath = Path.Combine(tempDirectory.FullName, "index.jsonl");
100+
FileInfo indexFile = new(indexPath);
101+
indexFile.Refresh();
102+
long fileSizeBeforeDispose = indexFile.Length;
103+
104+
// Data should already be on disk (file size > 0) before we dispose
105+
fileSizeBeforeDispose.Should().BeGreaterThan(0, "index.jsonl should be flushed to disk after CreateCheckpointAsync");
106+
107+
// Dispose to release file lock before final verification
108+
store.Dispose();
109+
110+
string[] lines = File.ReadAllLines(indexPath);
111+
lines.Should().HaveCount(1);
112+
lines[0].Should().Contain(checkpoint.CheckpointId);
113+
}
114+
115+
private async ValueTask Run_EscapeRootFolderTestAsync(string escapingPath)
116+
{
117+
// Arrange
118+
using TempDirectory tempDirectory = new();
119+
using FileSystemJsonCheckpointStore store = new(tempDirectory);
120+
121+
string naivePath = Path.Combine(tempDirectory.DirectoryInfo.FullName, escapingPath);
122+
123+
// Check that the naive path is actually outside the temp directory to validate the test is meaningful
124+
FileInfo naiveCheckpointFile = new(naivePath);
125+
tempDirectory.IsParentOf(naiveCheckpointFile).Should().BeFalse("The naive path should be outside the root folder to validate that escaping is necessary.");
126+
127+
// Act
128+
CheckpointInfo checkpointInfo = await store.CreateCheckpointAsync(escapingPath, TestData);
129+
130+
// Assert
131+
string naivePathWithCheckpointId = Path.Combine(tempDirectory.DirectoryInfo.FullName, $"{escapingPath}_{checkpointInfo.CheckpointId}.json");
132+
new FileInfo(naivePathWithCheckpointId).Exists.Should().BeFalse("The naive path should not be used to save a checkpoint file.");
133+
134+
string actualFileName = store.GetFileNameForCheckpoint(escapingPath, checkpointInfo);
135+
string actualFilePath = Path.Combine(tempDirectory.DirectoryInfo.FullName, actualFileName);
136+
FileInfo actualFile = new(actualFilePath);
137+
138+
tempDirectory.IsParentOf(actualFile).Should().BeTrue("The actual checkpoint should be saved inside the root folder.");
139+
actualFile.Exists.Should().BeTrue("The actual path should be used to save a checkpoint file.");
140+
}
141+
142+
[Fact]
143+
public async Task CreateCheckpointAsync_ShouldNotEscapeRootFolderAsync()
144+
{
145+
// The SessionId is used as part of the file name, but if it contains path characters such as /.. it can escape the root folder.
146+
// Testing that such characters are escaped properly to prevent directory traversal attacks, etc.
147+
148+
await this.Run_EscapeRootFolderTestAsync("../valid_suffix");
149+
150+
#if !NETFRAMEWORK
151+
if (OperatingSystem.IsWindows())
152+
{
153+
// Windows allows both \ and / as path separators, so we test both
154+
await this.Run_EscapeRootFolderTestAsync("..\\valid_suffix");
155+
}
156+
#else
157+
// .NET Framework is always on Windows
158+
await this.Run_EscapeRootFolderTestAsync("..\\valid_suffix");
159+
#endif
160+
}
161+
162+
private const string InvalidPathCharsWin32 = "\\/:*?\"<>|";
163+
private const string InvalidPathCharsUnix = "/";
164+
private const string InvalidPathCharsMacOS = "/:";
165+
166+
[Theory]
167+
[InlineData(InvalidPathCharsWin32)]
168+
[InlineData(InvalidPathCharsUnix)]
169+
[InlineData(InvalidPathCharsMacOS)]
170+
public async Task CreateCheckpointAsync_EscapesInvalidCharsAsync(string invalidChars)
171+
{
172+
// Arrange
173+
using TempDirectory tempDirectory = new();
174+
using FileSystemJsonCheckpointStore store = new(tempDirectory);
175+
176+
string runId = $"prefix_{invalidChars}_suffix";
177+
178+
Func<Task> createCheckpointAction = async () => await store.CreateCheckpointAsync(runId, TestData);
179+
await createCheckpointAction.Should().NotThrowAsync();
56180
}
57181
}

0 commit comments

Comments
 (0)