Orleans Multi-Tenant Credentials Cache - Complete Single File Application This is a complete Orleans application demonstrating multi-tenant credential caching in a single C# file. Features: - Orleans 9.2.1 with localhost clustering - Multi-tenant credential isolation using IGrainWithStringKey - Azure.Identity integration with automatic MSAL toke…

orleans-credentials-cache-single-file.cs

// Orleans Credentials Cache - Single File Application
// This demonstrates a complete Orleans multi-tenant credentials cache application in a single .cs file
// 
// To run: dotnet run orleans-cache.cs
// To publish as single-file executable: dotnet publish -c Release -r win-x64 -p:PublishSingleFile=true
//
// This single file contains:
// - Orleans host configuration
// - Grain interfaces and implementations
// - Azure.Identity integration with built-in caching
// - Complete demonstration program
//
// Features demonstrated:
// - Multi-tenant credential isolation (each tenant gets its own grain instance)
// - Azure.Identity TokenCredential with automatic MSAL caching
// - Manual caching pattern for educational purposes
// - Proper Orleans grain lifecycle management (OnActivateAsync)

using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.DependencyInjection;
using Orleans;
using Orleans.Hosting;
using Orleans.Runtime;
using Azure.Identity;
using Azure.Core;

// Create Orleans host builder
var builder = Host.CreateDefaultBuilder(args)
    .UseOrleans(siloBuilder =>
    {
        // Configure the silo for local development
        siloBuilder.UseLocalhostClustering();
    })
    .UseConsoleLifetime();

// Build and run the host
using var host = builder.Build();

Console.WriteLine("Orleans Silo is starting up...");
await host.StartAsync();

Console.WriteLine("Orleans Silo is running.");

// Demonstrate the credentials cache grain
await DemonstrateCredentialsCache(host.Services);

// Demonstrate the Azure Identity credentials cache grain
await DemonstrateAzureIdentityCredentialsCache(host.Services);

Console.WriteLine("Press any key to stop.");
Console.ReadKey();

Console.WriteLine("Orleans Silo is shutting down...");
await host.StopAsync();

/// <summary>
/// Demonstrates the usage of the credentials cache grain
/// </summary>
static async Task DemonstrateCredentialsCache(IServiceProvider services)
{
    var grainFactory = services.GetRequiredService<IGrainFactory>();
    
    Console.WriteLine("\n=== Credentials Cache Demonstration ===");
    
    // Get grain instances for different tenants
    var tenant1Grain = grainFactory.GetGrain<ICredentialsCacheGrain>("tenant-001");
    var tenant2Grain = grainFactory.GetGrain<ICredentialsCacheGrain>("tenant-002");
    
    try
    {
        // Test multiple scopes for tenant 1
        Console.WriteLine("\n--- Testing Tenant 1 ---");
        var token1a = await tenant1Grain.GetAccessToken("https://graph.microsoft.com/.default");
        var token1b = await tenant1Grain.GetAccessToken("https://vault.azure.net/.default");
        
        // Request the same token again to demonstrate caching
        Console.WriteLine("\n--- Testing Cache for Tenant 1 ---");
        var token1a_cached = await tenant1Grain.GetAccessToken("https://graph.microsoft.com/.default");
        var token1b_cached = await tenant1Grain.GetAccessToken("https://vault.azure.net/.default");
        
        Console.WriteLine($"Graph token matches cached: {token1a == token1a_cached}");
        Console.WriteLine($"Vault token matches cached: {token1b == token1b_cached}");
        
        // Test tenant 2 to show isolation
        Console.WriteLine("\n--- Testing Tenant 2 (Different Tenant) ---");
        var token2a = await tenant2Grain.GetAccessToken("https://graph.microsoft.com/.default");
        
        Console.WriteLine($"Tenant 1 and 2 have different tokens: {token1a != token2a}");
        
        // Test parallel access to the same grain
        Console.WriteLine("\n--- Testing Parallel Access ---");
        var parallelTasks = Enumerable.Range(1, 5)
            .Select(i => tenant1Grain.GetAccessToken($"scope-{i}"))
            .ToList();
        
        var parallelResults = await Task.WhenAll(parallelTasks);
        Console.WriteLine($"Generated {parallelResults.Length} tokens in parallel");
        
        Console.WriteLine("\n=== Demonstration Complete ===");
    }
    catch (Exception ex)
    {
        Console.WriteLine($"Error during demonstration: {ex.Message}");
    }
}

/// <summary>
/// Demonstrates the usage of the Azure Identity credentials cache grain
/// </summary>
static async Task DemonstrateAzureIdentityCredentialsCache(IServiceProvider services)
{
    var grainFactory = services.GetRequiredService<IGrainFactory>();
    
    Console.WriteLine("\n=== Azure Identity Credentials Cache Demonstration ===");
    
    // Get grain instances for different tenants
    var tenant1Grain = grainFactory.GetGrain<IAzureIdentityCredentialsCacheGrain>("tenant-001");
    var tenant2Grain = grainFactory.GetGrain<IAzureIdentityCredentialsCacheGrain>("tenant-002");
    
    try
    {
        // Test multiple scopes for tenant 1 using Azure.Identity
        Console.WriteLine("\n--- Testing Azure Identity Tenant 1 ---");
        var token1a = await tenant1Grain.GetAccessToken("https://graph.microsoft.com/.default");
        var token1b = await tenant1Grain.GetAccessToken("https://vault.azure.net/.default");
        
        // Request the same token again to demonstrate caching
        Console.WriteLine("\n--- Testing Azure Identity Cache for Tenant 1 ---");
        var token1a_cached = await tenant1Grain.GetAccessToken("https://graph.microsoft.com/.default");
        var token1b_cached = await tenant1Grain.GetAccessToken("https://vault.azure.net/.default");
        
        Console.WriteLine($"Graph token matches cached: {token1a == token1a_cached}");
        Console.WriteLine($"Vault token matches cached: {token1b == token1b_cached}");
        
        // Test tenant 2 to show isolation
        Console.WriteLine("\n--- Testing Azure Identity Tenant 2 (Different Tenant) ---");
        var token2a = await tenant2Grain.GetAccessToken("https://graph.microsoft.com/.default");
        
        Console.WriteLine($"Tenant 1 and 2 have different tokens: {token1a != token2a}");
        
        // Test parallel access to the same grain
        Console.WriteLine("\n--- Testing Azure Identity Parallel Access ---");
        var parallelTasks = Enumerable.Range(1, 3)
            .Select(i => tenant1Grain.GetAccessToken($"https://custom-api-{i}.example.com/.default"))
            .ToList();
        
        var parallelResults = await Task.WhenAll(parallelTasks);
        Console.WriteLine($"Generated {parallelResults.Length} Azure Identity tokens in parallel");
        
        Console.WriteLine("\n=== Azure Identity Demonstration Complete ===");
    }
    catch (Exception ex)
    {
        Console.WriteLine($"Error during Azure Identity demonstration: {ex.Message}");
    }
}

/// <summary>
/// Grain interface for caching tenant-specific credentials
/// </summary>
public interface ICredentialsCacheGrain : IGrainWithStringKey
{
    /// <summary>
    /// Get access token for the specified scope. Handles caching and refresh internally.
    /// The grain key represents the tenant ID.
    /// </summary>
    /// <param name="scope">The OAuth scope for the access token</param>
    /// <returns>Valid access token for the tenant and scope</returns>
    Task<string> GetAccessToken(string scope);
}

/// <summary>
/// Grain interface for caching tenant-specific credentials using Azure.Identity
/// </summary>
public interface IAzureIdentityCredentialsCacheGrain : IGrainWithStringKey
{
    /// <summary>
    /// Get access token for the specified scope using Azure.Identity TokenCredential.
    /// The grain key represents the tenant ID.
    /// </summary>
    /// <param name="scope">The OAuth scope for the access token</param>
    /// <returns>Valid access token for the tenant and scope</returns>
    Task<string> GetAccessToken(string scope);
}

public readonly record struct CachedToken(string AccessToken, DateTime ExpiresAt, string Scope)
{
    public readonly bool IsExpired => DateTime.UtcNow.AddMinutes(5) >= ExpiresAt;
}

/// <summary>
/// Orleans grain implementation for caching tenant-specific access tokens
/// </summary>
public class CredentialsCacheGrain : Grain, ICredentialsCacheGrain
{
    // In-memory cache of tokens by scope
    private readonly Dictionary<string, CachedToken> _tokenCache = new();

    public async Task<string> GetAccessToken(string scope)
    {
        var tenantId = this.GetPrimaryKeyString();
        
        // Check if we have a valid cached token for this scope
        if (_tokenCache.TryGetValue(scope, out var cachedToken) && !cachedToken.IsExpired)
        {
            Console.WriteLine($"Returning cached token for tenant '{tenantId}' and scope '{scope}'");
            return cachedToken.AccessToken;
        }

        // Token is missing or expired, fetch a new one
        Console.WriteLine($"Fetching new token for tenant '{tenantId}' and scope '{scope}'");
        
        var newToken = await FetchNewAccessToken(tenantId, scope);
        
        // Cache the new token
        _tokenCache[scope] = newToken;
        
        return newToken.AccessToken;
    }

    /// <summary>
    /// Simulates fetching a new access token from an identity provider
    /// In a real implementation, this would call your OAuth2/identity provider
    /// </summary>
    private async Task<CachedToken> FetchNewAccessToken(string tenantId, string scope)
    {
        // Simulate async call to identity provider
        await Task.Delay(100);
        
        // In a real implementation, you would:
        // 1. Load tenant-specific client credentials from secure storage
        // 2. Call the OAuth2 token endpoint with client_credentials grant
        // 3. Parse the response and extract the access token and expiry
        
        // For demonstration, generate a mock token
        var accessToken = $"mock_token_{tenantId}_{scope}_{Guid.NewGuid():N}";
        var expiresAt = DateTime.UtcNow.AddHours(1); // Tokens typically expire in 1 hour
        
        Console.WriteLine($"Generated new access token for tenant '{tenantId}' and scope '{scope}', expires at {expiresAt:yyyy-MM-dd HH:mm:ss} UTC");
        
        return new CachedToken
        {
            AccessToken = accessToken,
            ExpiresAt = expiresAt,
            Scope = scope
        };
    }
}

/// <summary>
/// Configuration for tenant-specific Azure credentials
/// </summary>
public record TenantCredentialsConfig
{
    public required string TenantId { get; init; }
    public required string ClientId { get; init; }
    public required string ClientSecret { get; init; }
}

/// <summary>
/// Orleans grain implementation for caching tenant-specific access tokens using Azure.Identity
/// </summary>
public class AzureIdentityCredentialsCacheGrain : Grain, IAzureIdentityCredentialsCacheGrain
{
    // TokenCredential per tenant - Azure.Identity handles caching internally
    private TokenCredential? _tokenCredential;

    public override async Task OnActivateAsync(CancellationToken cancellationToken)
    {
        var tenantId = this.GetPrimaryKeyString();
        
        // Initialize TokenCredential during grain activation
        var config = GetTenantCredentialsConfig(tenantId);
        
        // Create ClientSecretCredential using Azure.Identity
        // The underlying MSAL library handles all token caching automatically
        _tokenCredential = new ClientSecretCredential(
            tenantId: config.TenantId,
            clientId: config.ClientId,
            clientSecret: config.ClientSecret
        );
        
        Console.WriteLine($"[Azure.Identity] Initialized TokenCredential for tenant '{tenantId}'");
        
        await base.OnActivateAsync(cancellationToken);
    }

    public async Task<string> GetAccessToken(string scope)
    {
        var tenantId = this.GetPrimaryKeyString();
        
        Console.WriteLine($"[Azure.Identity] Requesting token for tenant '{tenantId}' and scope '{scope}'");
        
        // TokenCredential is already initialized during OnActivateAsync
        if (_tokenCredential == null)
        {
            throw new InvalidOperationException("TokenCredential not initialized. Grain activation may have failed.");
        }
        
        try
        {
            // Create TokenRequestContext for the scope
            var tokenRequestContext = new TokenRequestContext(new[] { scope });
            
            // Get access token - Azure.Identity handles caching automatically
            var accessToken = await _tokenCredential.GetTokenAsync(tokenRequestContext, CancellationToken.None);
            
            Console.WriteLine($"[Azure.Identity] Token acquired for tenant '{tenantId}' and scope '{scope}', expires at {accessToken.ExpiresOn:yyyy-MM-dd HH:mm:ss} UTC");
            
            return accessToken.Token;
        }
        catch (Exception ex)
        {
            // In case of authentication failure, fall back to mock token for demonstration
            Console.WriteLine($"[Azure.Identity] Authentication failed for tenant '{tenantId}': {ex.Message}");
            Console.WriteLine($"[Azure.Identity] Falling back to mock token for demonstration purposes");
            
            var mockAccessToken = $"mock_azure_token_{tenantId}_{scope}_{Guid.NewGuid():N}";
            return mockAccessToken;
        }
    }

    /// <summary>
    /// Gets tenant credentials configuration (mock implementation for demonstration)
    /// In a real implementation, this would securely load from Key Vault or configuration
    /// </summary>
    private TenantCredentialsConfig GetTenantCredentialsConfig(string tenantId)
    {
        // Mock configuration - in a real implementation, you would:
        // 1. Load from Azure Key Vault
        // 2. Load from secure configuration store
        // 3. Use managed identity or certificate-based authentication
        
        return new TenantCredentialsConfig
        {
            TenantId = tenantId,
            ClientId = Environment.GetEnvironmentVariable("AZURE_CLIENT_ID") ?? "mock-client-id",
            ClientSecret = Environment.GetEnvironmentVariable("AZURE_CLIENT_SECRET") ?? "mock-client-secret"
        };
    }
}

to use for azure-services, we could probably add something like this and register the OrleansTokenCredential as a singleton

namespace TestOrleans
{
    public class OrleansTokenCredential(IClusterClient clusterClient) : TokenCredential
    {
        public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken) => GetTokenAsync(requestContext, cancellationToken).GetAwaiter().GetResult();
        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
        {
            var grain = clusterClient.GetGrain<IAzureIdentityCredentialsCacheGrain>(requestContext.TenantId);
            return await grain.GetAccessToken(requestContext.Scopes.Single(), cancellationToken);
        }
    }

    public interface IAzureIdentityCredentialsCacheGrain : IGrainWithStringKey
    {
        Task<AccessToken> GetAccessToken(string scope, CancellationToken cancellationToken);
    }

    public class AzureIdentityCredentialsCacheGrain : Grain, IAzureIdentityCredentialsCacheGrain
    {
        private TokenCredential _credential;

        public override Task OnActivateAsync(CancellationToken cancellationToken)
        {
            var tenantId = this.GetPrimaryKeyString();
            _credential = new ClientSecretCredential(tenantId, "<clientId>", "<clientSecret>");
            return base.OnActivateAsync(cancellationToken);
        }

        public async Task<AccessToken> GetAccessToken(string scope, CancellationToken cancellationToken)
            => await _credential.GetTokenAsync(new([scope]), cancellationToken);
    }
}