support api_key as part of the url

Author: shaykeren

src/DigmaSSEServer/Authentication/ApiKeyValidatorMiddleware.cs

@@ -0,0 +1,37 @@
+using DigmaSSEServer.Options;
+using Microsoft.Extensions.Options;
+
+namespace DigmaSSEServer.Authentication;
+
+public class ApiKeyValidatorMiddleware
+{ 
+    private readonly RequestDelegate _next;
+    private readonly string _token;
+    const string INVALID_API_KEY = "unauthorized access invalid api key";
+
+    public ApiKeyValidatorMiddleware(RequestDelegate rd, IOptions<AuthOptions> options)
+    {
+        _next = rd;
+        _token = options.Value.Token;
+    }
+    private async Task UnauthorizedAccess(HttpContext httpContext, string errorMessage)
+    {
+        httpContext.Response.StatusCode = StatusCodes.Status401Unauthorized;
+        await httpContext.Response.WriteAsync(errorMessage);
+    }
+    
+    public async Task Invoke(HttpContext httpContext)
+    {
+        var routeValue = httpContext.GetRouteValue(Constants.API_KEY_NAME);
+        if (routeValue == null) await UnauthorizedAccess(httpContext,INVALID_API_KEY);
+        var apiKey = (string)routeValue!;
+        if (apiKey != _token)
+        {
+            await UnauthorizedAccess(httpContext,INVALID_API_KEY);
+            return;
+        }
+
+        
+        await _next.Invoke(httpContext);
+    }
+}

src/DigmaSSEServer/Authentication/Constants.cs

@@ -0,0 +1,7 @@
+namespace DigmaSSEServer.Authentication;
+
+public class Constants
+{
+    public const string API_KEY_NAME = "api_key";
+    public const string RoutePatternPrefix = $"/mcp/{{{API_KEY_NAME}}}";
+}

src/DigmaSSEServer/DigmaConfigurator.cs

@@ -26,24 +26,10 @@ static HttpClientHandler CreateHandler()
             return handler;
         }
 
-        services.AddHttpClient(TokenClientName, (sp, c) =>
-            {
-                var digma = sp.GetRequiredService<IOptions<DigmaOptions>>().Value;
-                var auth = sp.GetRequiredService<IOptions<AuthOptions>>().Value;
-
-                c.BaseAddress = new Uri(digma.AnalyticsApi);
-                c.DefaultRequestHeaders.Add("Digma-Access-Token", auth.Token);
-            })
+        services.AddHttpClient(TokenClientName, ConfigureHttpClient)
             .ConfigurePrimaryHttpMessageHandler(CreateHandler);
 
-        services.AddHttpClient<DigmaClient>((sp, c) =>
-            {
-                var digma = sp.GetRequiredService<IOptions<DigmaOptions>>().Value;
-                var auth = sp.GetRequiredService<IOptions<AuthOptions>>().Value;
-
-                c.BaseAddress = new Uri(digma.AnalyticsApi);
-                c.DefaultRequestHeaders.Add("Digma-Access-Token", auth.Token);
-            })
+        services.AddHttpClient<DigmaClient>(ConfigureHttpClient)
             .ConfigurePrimaryHttpMessageHandler(sp =>
             {
                 var delegating = sp.GetRequiredService<BearerTokenHandler>();
@@ -61,4 +47,13 @@ static HttpClientHandler CreateHandler()
 
         return services;
     }
+
+    private static void ConfigureHttpClient(IServiceProvider sp, HttpClient c)
+    {
+        var digma = sp.GetRequiredService<IOptions<DigmaOptions>>().Value;
+        var auth = sp.GetRequiredService<IOptions<AuthOptions>>().Value;
+
+        c.BaseAddress = new Uri(digma.AnalyticsApi);
+        c.DefaultRequestHeaders.Add("Digma-Access-Token", $"Token {auth.Token}");
+    }
 }

src/DigmaSSEServer/Program.cs

@@ -1,23 +1,38 @@
-using System.Net;
 using DigmaSSEServer;
+using DigmaSSEServer.Authentication;
+using DigmaSSEServer.Options;
 using DigmaSSEServer.Tools;
+using Microsoft.Extensions.Options;
 using ModelContextProtocol.Protocol.Types;
 
 var builder = WebApplication.CreateBuilder(args);
 
 builder.Services.AddDigma(builder.Configuration);
-
 builder.Services
-    .AddMcpServer(options => options.ServerInfo = new Implementation()
+    .AddMcpServer(options =>
     {
-        Name = "Digma Server",
-        Version = "1.0"
+        options.ServerInfo = new Implementation()
+        {
+            Name = "Digma Server",
+            Version = "1.0"
+            
+        };
     })
     .WithHttpTransport() // OR WithStdioServerTransport()
     .WithTools<CodeObservabilityTool>();
 
 var app = builder.Build();
+var authOptions = app.Services.GetRequiredService<IOptions<AuthOptions>>();
+if (string.IsNullOrWhiteSpace(authOptions.Value.Token))
+{
+    Console.WriteLine("ERROR: Missing required configuration 'Auth__Token must be specified'");
+    Environment.Exit(1); // Exit with failure code
+    return;
+}
+
+app.MapMcp(Constants.RoutePatternPrefix);
+
 
-app.MapMcp();
+app.UseMiddleware<ApiKeyValidatorMiddleware>();
 
 app.Run();

src/DigmaSSEServer/appsettings.json

@@ -5,6 +5,13 @@
       "Microsoft.AspNetCore": "Warning"
     }
   },
+  "Kestrel": {
+    "Endpoints": {
+      "Http": {
+        "Url": "http://0.0.0.0:3000"
+      }
+    }
+  },
   "AllowedHosts": "*",
   "Digma": {
     "AnalyticsApi": ""