fix: Harden scope mapping

2026-06-16 21:04:09 +08:00 · 2026-05-07 11:51:24 -04:00
parent f346f2a618
commit cebd3512cd
2 changed files with 117 additions and 15 deletions
@@ -6,6 +6,7 @@ using System.Security.Claims;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
+using Microsoft.Shared.Diagnostics;

 namespace Microsoft.Agents.AI.Hosting;

@@ -36,7 +37,7 @@ public class UserIdentityScopedSessionStore : DelegatingAgentSessionStore
    /// </param>
    /// <param name="strict">
    /// If <see langword="true"/>, an exception is thrown when the specified claim is not found.
-    /// If <see langword="false"/>, operations proceed without scoping when the claim is absent.
+    /// If <see langword="false"/>, the conversation ID is passed through unmodified when the claim is absent.
    /// </param>
    public UserIdentityScopedSessionStore(AgentSessionStore innerStore,
                                          IHttpContextAccessor? contextAccessor,
@@ -44,7 +45,7 @@ public class UserIdentityScopedSessionStore : DelegatingAgentSessionStore
                                          bool strict = true) : base(innerStore)
    {
        this._httpContextAccessor = contextAccessor;
-        this._claimType = claimType;
+        this._claimType = Throw.IfNullOrWhitespace(claimType);
        this._strict = strict;
    }

@@ -64,7 +65,18 @@ public class UserIdentityScopedSessionStore : DelegatingAgentSessionStore

    private string? ScopeId => this.GetScopeFromIdentity();

-    private string GetScopedConversationId(string bareConversationId) => $"{this.ScopeId}:{bareConversationId}";
+    private static string EscapeScopeId(string scopeId) => scopeId.Replace("\\", "\\\\").Replace(":", "\\:");
+
+    private string GetScopedConversationId(string bareConversationId)
+    {
+        string? scopeId = this.ScopeId;
+        if (scopeId == null)
+        {
+            return bareConversationId;
+        }
+
+        return $"{EscapeScopeId(scopeId)}::{bareConversationId}";
+    }

    /// <inheritdoc />
    public override ValueTask<AgentSession> GetSessionAsync(AIAgent agent, string conversationId, CancellationToken cancellationToken = default)
@@ -65,6 +65,27 @@ public class UserIdentityScopedSessionStoreTests
        Assert.NotNull(store);
    }

+    /// <summary>
+    /// Verify that constructor throws ArgumentException when claimType is null.
+    /// </summary>
+    [Fact]
+    public void RequiresClaimType_NotNull() =>
+        Assert.Throws<ArgumentNullException>("claimType", () => new UserIdentityScopedSessionStore(this._innerStoreMock.Object, this._httpContextAccessorMock.Object, claimType: null!));
+
+    /// <summary>
+    /// Verify that constructor throws ArgumentException when claimType is empty.
+    /// </summary>
+    [Fact]
+    public void RequiresClaimType_NotEmpty() =>
+        Assert.Throws<ArgumentException>("claimType", () => new UserIdentityScopedSessionStore(this._innerStoreMock.Object, this._httpContextAccessorMock.Object, claimType: string.Empty));
+
+    /// <summary>
+    /// Verify that constructor throws ArgumentException when claimType is whitespace.
+    /// </summary>
+    [Fact]
+    public void RequiresClaimType_NotWhitespace() =>
+        Assert.Throws<ArgumentException>("claimType", () => new UserIdentityScopedSessionStore(this._innerStoreMock.Object, this._httpContextAccessorMock.Object, claimType: "   "));
+
    #endregion

    #region GetSessionAsync Tests
@@ -86,7 +107,7 @@ public class UserIdentityScopedSessionStoreTests
        this._innerStoreMock.Verify(
            x => x.GetSessionAsync(
                this._agentMock.Object,
-                $"{TestUserId}:{TestConversationId}",
+                $"{TestUserId}::{TestConversationId}",
                It.IsAny<CancellationToken>()),
            Times.Once);
    }
@@ -111,7 +132,7 @@ public class UserIdentityScopedSessionStoreTests
        this._innerStoreMock.Verify(
            x => x.GetSessionAsync(
                this._agentMock.Object,
-                $"{CustomClaimValue}:{TestConversationId}",
+                $"{CustomClaimValue}::{TestConversationId}",
                It.IsAny<CancellationToken>()),
            Times.Once);
    }
@@ -152,11 +173,11 @@ public class UserIdentityScopedSessionStoreTests
        // Act - should not throw
        await store.GetSessionAsync(this._agentMock.Object, TestConversationId);

-        // Assert - conversation ID should use null scope
+        // Assert - conversation ID should be passed through unmodified
        this._innerStoreMock.Verify(
            x => x.GetSessionAsync(
                this._agentMock.Object,
-                $":{TestConversationId}",
+                TestConversationId,
                It.IsAny<CancellationToken>()),
            Times.Once);
    }
@@ -200,7 +221,7 @@ public class UserIdentityScopedSessionStoreTests
        this._innerStoreMock.Verify(
            x => x.SaveSessionAsync(
                this._agentMock.Object,
-                $"{TestUserId}:{TestConversationId}",
+                $"{TestUserId}::{TestConversationId}",
                sessionToSave,
                It.IsAny<CancellationToken>()),
            Times.Once);
@@ -227,7 +248,7 @@ public class UserIdentityScopedSessionStoreTests
        this._innerStoreMock.Verify(
            x => x.SaveSessionAsync(
                this._agentMock.Object,
-                $"{CustomClaimValue}:{TestConversationId}",
+                $"{CustomClaimValue}::{TestConversationId}",
                sessionToSave,
                It.IsAny<CancellationToken>()),
            Times.Once);
@@ -271,11 +292,11 @@ public class UserIdentityScopedSessionStoreTests
        // Act - should not throw
        await store.SaveSessionAsync(this._agentMock.Object, TestConversationId, sessionToSave);

-        // Assert - conversation ID should use null scope
+        // Assert - conversation ID should be passed through unmodified
        this._innerStoreMock.Verify(
            x => x.SaveSessionAsync(
                this._agentMock.Object,
-                $":{TestConversationId}",
+                TestConversationId,
                sessionToSave,
                It.IsAny<CancellationToken>()),
            Times.Once);
@@ -319,11 +340,11 @@ public class UserIdentityScopedSessionStoreTests
        // Act - should not throw
        await store.GetSessionAsync(this._agentMock.Object, TestConversationId);

-        // Assert
+        // Assert - conversation ID should be passed through unmodified
        this._innerStoreMock.Verify(
            x => x.GetSessionAsync(
                this._agentMock.Object,
-                $":{TestConversationId}",
+                TestConversationId,
                It.IsAny<CancellationToken>()),
            Times.Once);
    }
@@ -364,11 +385,80 @@ public class UserIdentityScopedSessionStoreTests
        await store2.GetSessionAsync(this._agentMock.Object, TestConversationId);

        // Assert
-        Assert.Equal($"{User1}:{TestConversationId}", capturedConversationId1);
-        Assert.Equal($"{User2}:{TestConversationId}", capturedConversationId2);
+        Assert.Equal($"{User1}::{TestConversationId}", capturedConversationId1);
+        Assert.Equal($"{User2}::{TestConversationId}", capturedConversationId2);
        Assert.NotEqual(capturedConversationId1, capturedConversationId2);
    }

+    /// <summary>
+    /// Verify that colons in user claim values are escaped.
+    /// </summary>
+    [Fact]
+    public async Task EscapesColonsInUserClaimValueAsync()
+    {
+        // Arrange
+        const string UserIdWithColon = "user:with:colons";
+        this.SetupHttpContextWithClaim(ClaimsIdentity.DefaultNameClaimType, UserIdWithColon);
+        var store = new UserIdentityScopedSessionStore(this._innerStoreMock.Object, this._httpContextAccessorMock.Object);
+
+        // Act
+        await store.GetSessionAsync(this._agentMock.Object, TestConversationId);
+
+        // Assert - colons should be escaped as \:
+        this._innerStoreMock.Verify(
+            x => x.GetSessionAsync(
+                this._agentMock.Object,
+                $"user\\:with\\:colons::{TestConversationId}",
+                It.IsAny<CancellationToken>()),
+            Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that backslashes in user claim values are escaped.
+    /// </summary>
+    [Fact]
+    public async Task EscapesBackslashesInUserClaimValueAsync()
+    {
+        // Arrange
+        const string UserIdWithBackslash = @"domain\user";
+        this.SetupHttpContextWithClaim(ClaimsIdentity.DefaultNameClaimType, UserIdWithBackslash);
+        var store = new UserIdentityScopedSessionStore(this._innerStoreMock.Object, this._httpContextAccessorMock.Object);
+
+        // Act
+        await store.GetSessionAsync(this._agentMock.Object, TestConversationId);
+
+        // Assert - backslashes should be escaped as \\
+        this._innerStoreMock.Verify(
+            x => x.GetSessionAsync(
+                this._agentMock.Object,
+                $"domain\\\\user::{TestConversationId}",
+                It.IsAny<CancellationToken>()),
+            Times.Once);
+    }
+
+    /// <summary>
+    /// Verify that both backslashes and colons in user claim values are escaped correctly.
+    /// </summary>
+    [Fact]
+    public async Task EscapesBothBackslashesAndColonsInUserClaimValueAsync()
+    {
+        // Arrange
+        const string UserIdWithBoth = @"domain\user:role";
+        this.SetupHttpContextWithClaim(ClaimsIdentity.DefaultNameClaimType, UserIdWithBoth);
+        var store = new UserIdentityScopedSessionStore(this._innerStoreMock.Object, this._httpContextAccessorMock.Object);
+
+        // Act
+        await store.GetSessionAsync(this._agentMock.Object, TestConversationId);
+
+        // Assert - backslashes escaped first, then colons
+        this._innerStoreMock.Verify(
+            x => x.GetSessionAsync(
+                this._agentMock.Object,
+                $"domain\\\\user\\:role::{TestConversationId}",
+                It.IsAny<CancellationToken>()),
+            Times.Once);
+    }
+
    #endregion

    #region Helper Methods