diff --git a/src/Backend/AutopilotMonitor.Functions.Tests/EndpointNormalizerTests.cs b/src/Backend/AutopilotMonitor.Functions.Tests/EndpointNormalizerTests.cs new file mode 100644 index 00000000..f590175f --- /dev/null +++ b/src/Backend/AutopilotMonitor.Functions.Tests/EndpointNormalizerTests.cs @@ -0,0 +1,57 @@ +using AutopilotMonitor.Functions.Helpers; + +namespace AutopilotMonitor.Functions.Tests; + +public class EndpointNormalizerTests +{ + [Fact] + public void Normalize_StripsApiPrefixReplacesGuidsAndKeyChars() + { + var normalized = EndpointNormalizer.Normalize("/api/Sessions/0f8fad5b-d9cb-469f-a165-70867728950e/events?x=1"); + Assert.Equal("sessions_{id}_events_x=1", normalized); + } + + [Fact] + public void ToolNameKey_RealToolName_PassesThroughUnchanged() + { + // Existing UserUsageLog rows are keyed by ":"; a real tool name must keep its key. + Assert.Equal("get_session_events", EndpointNormalizer.ToolNameKey("get_session_events")); + } + + [Fact] + public void ToolNameKey_NullOrEmpty_ReturnsEmpty() + { + Assert.Equal(string.Empty, EndpointNormalizer.ToolNameKey(null)); + Assert.Equal(string.Empty, EndpointNormalizer.ToolNameKey(string.Empty)); + } + + [Fact] + public void ToolNameKey_WithCrlfAndControlChars_StripsThem() + { + // The log-forging vector: a header that tries to inject a fake log line through the row key. + var forged = $"get_metrics\r\n[McpQuota] FAKE{(char)0}{(char)27}{(char)0x2028}x"; + var key = EndpointNormalizer.ToolNameKey(forged); + Assert.Equal("get_metrics[McpQuota] FAKEx", key); + } + + [Fact] + public void ToolNameKey_WithTableKeyIllegalChars_ReplacesThem() + { + // '/', '\', '#', '?' are rejected by Table Storage in keys — a header carrying them would make + // every usage increment fail with 400, so they are replaced the same way the path is. + Assert.Equal("a_b_c_d_e", EndpointNormalizer.ToolNameKey("a/b\\c#d?e")); + } + + [Fact] + public void ToolNameKey_LongerThanCap_IsTruncated() + { + var key = EndpointNormalizer.ToolNameKey(new string('t', EndpointNormalizer.MaxToolNameLength + 10)); + Assert.Equal(EndpointNormalizer.MaxToolNameLength, key.Length); + } + + [Fact] + public void ToolNameKey_OnlyControlChars_ReturnsEmpty() + { + Assert.Equal(string.Empty, EndpointNormalizer.ToolNameKey("\r\n ")); + } +} diff --git a/src/Backend/AutopilotMonitor.Functions/Helpers/EndpointNormalizer.cs b/src/Backend/AutopilotMonitor.Functions/Helpers/EndpointNormalizer.cs index 8de78a8c..cd3c555c 100644 --- a/src/Backend/AutopilotMonitor.Functions/Helpers/EndpointNormalizer.cs +++ b/src/Backend/AutopilotMonitor.Functions/Helpers/EndpointNormalizer.cs @@ -12,6 +12,16 @@ public static class EndpointNormalizer @"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}", RegexOptions.Compiled); + // Same classes LogSanitizer strips: C0/C1 control characters (Table Storage rejects them in + // keys) plus the Unicode line/paragraph separators a log viewer breaks a line on. + private static readonly Regex ControlChars = new(@"[\p{Cc}\p{Zl}\p{Zp}]", RegexOptions.Compiled); + + /// + /// Cap for the caller-supplied tool-name prefix of a usage row key. Real MCP tool names are + /// well under this; the cap only bounds what an arbitrary header can push into a key. + /// + public const int MaxToolNameLength = 64; + public static string Normalize(string path) { // Strip /api/ prefix @@ -22,9 +32,32 @@ public static class EndpointNormalizer // Replace GUIDs with {id} normalized = GuidPattern.Replace(normalized, "{id}"); - // Replace characters that are illegal in Azure Table Storage keys (/, \, #, ?) - normalized = normalized.Replace('/', '_').Replace('\\', '_').Replace('#', '_').Replace('?', '_'); - - return normalized.ToLowerInvariant(); + return ReplaceIllegalKeyChars(normalized).ToLowerInvariant(); } + + /// + /// Makes the X-MCP-Tool-Name header value safe to use as the prefix of a usage row key. + /// The header is caller-supplied (the MCP server sets it to the real tool name, but any + /// authenticated MCP-source caller can send it), and the key ends up in the CAS-conflict warning + /// of the usage counter: control characters and line separators are removed, the characters Table + /// Storage rejects in keys are replaced, and the length is capped. A real tool name + /// (get_session_events) passes through unchanged, so existing usage rows keep their keys. + /// Returns an empty string when nothing key-safe remains; callers then skip the prefix. + /// + public static string ToolNameKey(string? toolName) + { + if (string.IsNullOrEmpty(toolName)) + return string.Empty; + + // Explicit CR/LF removal first — the log-injection barrier the CodeQL taint model + // recognizes; the regex pass then drops any remaining control character or separator. + var cleaned = toolName.Replace("\r", string.Empty).Replace("\n", string.Empty); + cleaned = ControlChars.Replace(cleaned, string.Empty); + cleaned = ReplaceIllegalKeyChars(cleaned).Trim(); + return cleaned.Length <= MaxToolNameLength ? cleaned : cleaned[..MaxToolNameLength]; + } + + /// Replaces the characters Azure Table Storage rejects in keys (/, \, #, ?). + private static string ReplaceIllegalKeyChars(string value) + => value.Replace('/', '_').Replace('\\', '_').Replace('#', '_').Replace('?', '_'); } diff --git a/src/Backend/AutopilotMonitor.Functions/Middleware/McpQuotaEnforcementMiddleware.cs b/src/Backend/AutopilotMonitor.Functions/Middleware/McpQuotaEnforcementMiddleware.cs index a4acb804..96171150 100644 --- a/src/Backend/AutopilotMonitor.Functions/Middleware/McpQuotaEnforcementMiddleware.cs +++ b/src/Backend/AutopilotMonitor.Functions/Middleware/McpQuotaEnforcementMiddleware.cs @@ -354,8 +354,9 @@ public class McpQuotaEnforcementMiddleware : IFunctionsWorkerMiddleware private void TrackUsage(HttpContext httpContext, string oid, string? upn, string homeTenantId, IReadOnlyCollection chargeTenantIds) { var normalizedEndpoint = EndpointNormalizer.Normalize(httpContext.Request.Path.Value ?? string.Empty); - var mcpToolName = httpContext.Request.Headers["X-MCP-Tool-Name"].FirstOrDefault(); - if (!string.IsNullOrEmpty(mcpToolName)) + // Key-safe copy of the header: the value becomes part of the UserUsageLog row key. + var mcpToolName = EndpointNormalizer.ToolNameKey(httpContext.Request.Headers["X-MCP-Tool-Name"].FirstOrDefault()); + if (mcpToolName.Length > 0) normalizedEndpoint = $"{mcpToolName}:{normalizedEndpoint}"; var repo = _userUsageRepo;