Move auth into middleware (#3133)

Closes #2098.

This cleans up the authentication a bit; after this change we have two stages in the middleware pipeline:

- `AuthenticationMiddleware` reads the JWT token (it does not validate it, this is done by the Azure Functions service) and stores it in `FunctionContext.Items["ONEFUZZ_USER_INFO"]`
- `AuthorizationMiddleware` checks the user info against the `[Authorize]` attribute to see if the user has the required permissions
- Functions can read the user info from the `FunctionContext` if needed

The authorize attribute can be `[Authorize(Allow.User)]` or `Allow.Agent` or `Allow.Admin`. The `Admin` case is new and allows this to be declaratively specified rather than being checked in code. We have several functions which could be changed to use this (e.g. Pool POST/DELETE/PATCH, Scaleset POST/DELETE/PATCH), but I have only changed one so far (JinjaToScriban).

One of the benefits here is that this simplifies the test code a lot: we can set the desired user info directly onto our `(Test)FunctionContext` rather than having to supply a fake that pretends to parse the token from the HTTP request. This will also have benefits when running the service locally for testing purposes (refer to internal issue).

The other benefit is the ability to programmatically read the required authentication for each function, which may help with Swagger generation.

src/ApiService/ApiService/Auth/AuthenticationItems.cs

@@ -0,0 +1,16 @@
+using Microsoft.Azure.Functions.Worker;
+
+namespace Microsoft.OneFuzz.Service.Auth;
+
+public static class AuthenticationItems {
+    private const string Key = "ONEFUZZ_USER_INFO";
+
+    public static void SetUserAuthInfo(this FunctionContext context, UserAuthInfo info)
+        => context.Items[Key] = info;
+
+    public static UserAuthInfo GetUserAuthInfo(this FunctionContext context)
+        => (UserAuthInfo)context.Items[Key];
+
+    public static UserAuthInfo? TryGetUserAuthInfo(this FunctionContext context)
+        => context.Items.TryGetValue(Key, out var result) ? (UserAuthInfo)result : null;
+}

src/ApiService/ApiService/Auth/AuthenticationMiddleware.cs

@@ -0,0 +1,111 @@
+using System.IdentityModel.Tokens.Jwt;
+using System.Net;
+using System.Net.Http.Headers;
+using Microsoft.Azure.Functions.Worker;
+using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.Azure.Functions.Worker.Middleware;
+
+namespace Microsoft.OneFuzz.Service.Auth;
+
+public sealed class AuthenticationMiddleware : IFunctionsWorkerMiddleware {
+    private readonly IConfigOperations _config;
+    private readonly ILogTracer _log;
+
+    public AuthenticationMiddleware(IConfigOperations config, ILogTracer log) {
+        _config = config;
+        _log = log;
+    }
+
+    public async Async.Task Invoke(FunctionContext context, FunctionExecutionDelegate next) {
+        var requestData = await context.GetHttpRequestDataAsync();
+        if (requestData is not null) {
+            var authToken = GetAuthToken(requestData);
+            if (authToken is not null) {
+                // note that no validation of the token is performed here
+                // this is done globally by Azure Functions; see the configuration in
+                // 'function.bicep'
+                var token = new JwtSecurityToken(authToken);
+                var allowedTenants = await AllowedTenants();
+                if (!allowedTenants.Contains(token.Issuer)) {
+                    await BadIssuer(requestData, context, token, allowedTenants);
+                    return;
+                }
+
+                context.SetUserAuthInfo(UserInfoFromAuthToken(token));
+            }
+        }
+
+        await next(context);
+    }
+
+    private static UserAuthInfo UserInfoFromAuthToken(JwtSecurityToken token)
+        => token.Payload.Claims.Aggregate(
+            seed: new UserAuthInfo(new UserInfo(null, null, null), new List<string>()),
+            (acc, claim) => {
+                switch (claim.Type) {
+                    case "oid":
+                        return acc with { UserInfo = acc.UserInfo with { ObjectId = Guid.Parse(claim.Value) } };
+                    case "appid":
+                        return acc with { UserInfo = acc.UserInfo with { ApplicationId = Guid.Parse(claim.Value) } };
+                    case "upn":
+                        return acc with { UserInfo = acc.UserInfo with { Upn = claim.Value } };
+                    case "roles":
+                        acc.Roles.Add(claim.Value);
+                        return acc;
+                    default:
+                        return acc;
+                }
+            });
+
+    private async Async.ValueTask BadIssuer(
+        HttpRequestData request,
+        FunctionContext context,
+        JwtSecurityToken token,
+        IEnumerable<string> allowedTenants) {
+
+        var tenantsStr = string.Join("; ", allowedTenants);
+        _log.Error($"issuer not from allowed tenant. issuer: {token.Issuer:Tag:Issuer} - tenants: {tenantsStr:Tag:Tenants}");
+
+        var response = HttpResponseData.CreateResponse(request);
+        var status = HttpStatusCode.BadRequest;
+        await response.WriteAsJsonAsync(
+            new ProblemDetails(
+                status,
+                new Error(
+                    ErrorCode.INVALID_REQUEST,
+                    new List<string> {
+                        "unauthorized AAD issuer. If multi-tenant auth is failing, make sure to include all tenant_ids in the `allowed_aad_tenants` list in the instance_config. To see the current instance_config, run `onefuzz instance_config get`. "
+                    }
+                )),
+            "application/problem+json",
+            status);
+
+        context.GetInvocationResult().Value = response;
+    }
+
+    private async Async.Task<IEnumerable<string>> AllowedTenants() {
+        var config = await _config.Fetch();
+        return config.AllowedAadTenants.Select(t => $"https://sts.windows.net/{t}/");
+    }
+
+    private static string? GetAuthToken(HttpRequestData requestData)
+        => GetBearerToken(requestData) ?? GetAadIdToken(requestData);
+
+    private static string? GetAadIdToken(HttpRequestData requestData) {
+        if (!requestData.Headers.TryGetValues("x-ms-token-aad-id-token", out var values)) {
+            return null;
+        }
+
+        return values.First();
+    }
+
+    private static string? GetBearerToken(HttpRequestData requestData) {
+        if (!requestData.Headers.TryGetValues("Authorization", out var values)
+            || !AuthenticationHeaderValue.TryParse(values.First(), out var headerValue)
+            || !string.Equals(headerValue.Scheme, "Bearer", StringComparison.OrdinalIgnoreCase)) {
+            return null;
+        }
+
+        return headerValue.Parameter;
+    }
+}

src/ApiService/ApiService/Auth/AuthorizationMiddleware.cs

@@ -0,0 +1,103 @@
+using System.Collections.Immutable;
+using System.Diagnostics;
+using System.Net;
+using System.Reflection;
+using Microsoft.Azure.Functions.Worker;
+using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.Azure.Functions.Worker.Middleware;
+
+namespace Microsoft.OneFuzz.Service.Auth;
+
+public sealed class AuthorizationMiddleware : IFunctionsWorkerMiddleware {
+    private readonly IEndpointAuthorization _auth;
+    private readonly ILogTracer _log;
+
+    public AuthorizationMiddleware(IEndpointAuthorization auth, ILogTracer log) {
+        _auth = auth;
+        _log = log;
+    }
+
+    public async Async.Task Invoke(FunctionContext context, FunctionExecutionDelegate next) {
+        var attribute = GetAuthorizeAttribute(context);
+        if (attribute is not null) {
+            var req = await context.GetHttpRequestDataAsync() ?? throw new NotSupportedException("no HTTP request data found");
+            var user = context.TryGetUserAuthInfo();
+            if (user is null) {
+                await Reject(req, context, "no authentication");
+                return;
+            }
+
+            var (isAgent, _) = await _auth.IsAgent(user);
+            if (isAgent) {
+                if (attribute.Allow != Allow.Agent) {
+                    await Reject(req, context, "endpoint not allowed for agents");
+                    return;
+                }
+            } else {
+                if (attribute.Allow == Allow.Agent) {
+                    await Reject(req, context, "endpoint not allowed for users");
+                    return;
+                }
+
+                Debug.Assert(attribute.Allow is Allow.User or Allow.Admin);
+
+                // check access control first
+                var access = await _auth.CheckAccess(req);
+                if (!access.IsOk) {
+                    await Reject(req, context, "access control rejected request");
+                    return;
+                }
+
+                // check admin next
+                if (attribute.Allow == Allow.Admin) {
+                    var adminAccess = await _auth.CheckRequireAdmins(user);
+                    if (!adminAccess.IsOk) {
+                        await Reject(req, context, "must be admin to use this endpoint");
+                        return;
+                    }
+                }
+            }
+        }
+
+        await next(context);
+    }
+
+    private static async Async.ValueTask Reject(HttpRequestData request, FunctionContext context, string reason) {
+        var response = HttpResponseData.CreateResponse(request);
+        var status = HttpStatusCode.Unauthorized;
+        await response.WriteAsJsonAsync(
+            new ProblemDetails(
+                status,
+                Error.Create(ErrorCode.UNAUTHORIZED, reason)),
+            "application/problem+json",
+            status);
+
+        context.GetInvocationResult().Value = response;
+    }
+
+    // use ImmutableDictionary to prevent needing to lock and without the overhead
+    // of ConcurrentDictionary
+    private static ImmutableDictionary<string, AuthorizeAttribute?> _authorizeCache =
+        ImmutableDictionary.Create<string, AuthorizeAttribute?>();
+
+    private static AuthorizeAttribute? GetAuthorizeAttribute(FunctionContext context) {
+        // fully-qualified name of the method
+        var entryPoint = context.FunctionDefinition.EntryPoint;
+        if (_authorizeCache.TryGetValue(entryPoint, out var cached)) {
+            return cached;
+        }
+
+        var lastDot = entryPoint.LastIndexOf('.');
+        var (typeName, methodName) = (entryPoint[..lastDot], entryPoint[(lastDot + 1)..]);
+        var assemblyPath = context.FunctionDefinition.PathToAssembly;
+        var assembly = Assembly.LoadFrom(assemblyPath); // should already be loaded
+        var type = assembly.GetType(typeName)!;
+        var method = type.GetMethod(methodName)!;
+        var result =
+            method.GetCustomAttribute<AuthorizeAttribute>()
+            ?? type.GetCustomAttribute<AuthorizeAttribute>();
+
+        _authorizeCache = _authorizeCache.SetItem(entryPoint, result);
+        return result;
+    }
+}

src/ApiService/ApiService/Auth/AuthorizeAttribute.cs

@@ -0,0 +1,17 @@
+namespace Microsoft.OneFuzz.Service.Auth;
+
+[AttributeUsage(AttributeTargets.Class | AttributeTargets.Method)]
+public sealed class AuthorizeAttribute : Attribute {
+    public AuthorizeAttribute(Allow allow) {
+        Allow = allow;
+    }
+
+    public Allow Allow { get; set; }
+}
+
+public enum Allow {
+    Agent,
+    User,
+    Admin,
+
+}

src/ApiService/ApiService/Functions/AgentCanSchedule.cs

@@ -1,27 +1,23 @@
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.OneFuzz.Service.Auth;
 
 namespace Microsoft.OneFuzz.Service.Functions;
 
 public class AgentCanSchedule {
     private readonly ILogTracer _log;
-    private readonly IEndpointAuthorization _auth;
     private readonly IOnefuzzContext _context;
 
-
-    public AgentCanSchedule(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context) {
+    public AgentCanSchedule(ILogTracer log, IOnefuzzContext context) {
         _log = log;
-        _auth = auth;
         _context = context;
     }
 
     [Function("AgentCanSchedule")]
-    public Async.Task<HttpResponseData> Run(
+    [Authorize(Allow.Agent)]
+    public async Async.Task<HttpResponseData> Run(
         [HttpTrigger(AuthorizationLevel.Anonymous, "POST", Route="agents/can_schedule")]
-        HttpRequestData req)
-        => _auth.CallIfAgent(req, Post);
-
-    private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
+        HttpRequestData req) {
         var request = await RequestHandling.ParseRequest<CanScheduleRequest>(req);
         if (!request.IsOk) {
             _log.Warning($"Cannot schedule due to {request.ErrorV}");

src/ApiService/ApiService/Functions/AgentCommands.cs

@@ -1,28 +1,28 @@
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.OneFuzz.Service.Auth;
 
 namespace Microsoft.OneFuzz.Service.Functions;
 
 public class AgentCommands {
     private readonly ILogTracer _log;
-    private readonly IEndpointAuthorization _auth;
     private readonly IOnefuzzContext _context;
 
-    public AgentCommands(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context) {
+    public AgentCommands(ILogTracer log, IOnefuzzContext context) {
         _log = log;
-        _auth = auth;
         _context = context;
     }
 
     [Function("AgentCommands")]
+    [Authorize(Allow.Agent)]
     public Async.Task<HttpResponseData> Run(
         [HttpTrigger(AuthorizationLevel.Anonymous, "GET", "DELETE", Route="agents/commands")]
         HttpRequestData req)
-        => _auth.CallIfAgent(req, r => r.Method switch {
+        => req.Method switch {
             "GET" => Get(req),
             "DELETE" => Delete(req),
             _ => throw new NotSupportedException($"HTTP Method {req.Method} is not supported for this method")
-        });
+        };
 
     private async Async.Task<HttpResponseData> Get(HttpRequestData req) {
         var request = await RequestHandling.ParseRequest<NodeCommandGet>(req);

src/ApiService/ApiService/Functions/AgentEvents.cs

@@ -1,28 +1,25 @@
 using System.Threading.Tasks;
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.OneFuzz.Service.Auth;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 
 namespace Microsoft.OneFuzz.Service.Functions;
 
 public class AgentEvents {
     private readonly ILogTracer _log;
-    private readonly IEndpointAuthorization _auth;
     private readonly IOnefuzzContext _context;
 
-    public AgentEvents(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context) {
+    public AgentEvents(ILogTracer log, IOnefuzzContext context) {
         _log = log;
-        _auth = auth;
         _context = context;
     }
 
     [Function("AgentEvents")]
-    public Async.Task<HttpResponseData> Run(
+    [Authorize(Allow.Agent)]
+    public async Async.Task<HttpResponseData> Run(
         [HttpTrigger(AuthorizationLevel.Anonymous, "POST", Route="agents/events")]
-        HttpRequestData req)
-        => _auth.CallIfAgent(req, Post);
-
-    private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
+        HttpRequestData req) {
         var request = await RequestHandling.ParseRequest<NodeStateEnvelope>(req);
         if (!request.IsOk) {
             return await _context.RequestHandling.NotOk(req, request.ErrorV, context: "node event");

src/ApiService/ApiService/Functions/AgentRegistration.cs

@@ -1,33 +1,31 @@
 using Azure.Storage.Sas;
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.OneFuzz.Service.Auth;
 
 namespace Microsoft.OneFuzz.Service.Functions;
 
 public class AgentRegistration {
     private readonly ILogTracer _log;
-    private readonly IEndpointAuthorization _auth;
     private readonly IOnefuzzContext _context;
 
-    public AgentRegistration(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context) {
+    public AgentRegistration(ILogTracer log, IOnefuzzContext context) {
         _log = log;
-        _auth = auth;
         _context = context;
     }
 
     [Function("AgentRegistration")]
+    [Authorize(Allow.Agent)]
     public Async.Task<HttpResponseData> Run(
         [HttpTrigger(
             AuthorizationLevel.Anonymous,
             "GET", "POST",
             Route="agents/registration")] HttpRequestData req)
-        => _auth.CallIfAgent(
-            req,
-            r => r.Method switch {
-                "GET" => Get(r),
-                "POST" => Post(r),
-                var m => throw new InvalidOperationException($"method {m} not supported"),
-            });
+        => req.Method switch {
+            "GET" => Get(req),
+            "POST" => Post(req),
+            var m => throw new InvalidOperationException($"method {m} not supported"),
+        };
 
     private async Async.Task<HttpResponseData> Get(HttpRequestData req) {
         var request = await RequestHandling.ParseUri<AgentRegistrationGet>(req);