gandalf-reborn/src/dotnet/Suspectus.Gandalf.Mithrandir.Api/Controllers/AppProxyController.cs

using System.Net;
using System.Net.Http.Headers;
using Microsoft.AspNetCore.Mvc;
using Microsoft.Extensions.Caching.Distributed;
using Suspectus.Gandalf.Core.Abstractions.DTOs.Internal.Auth;
using Suspectus.Gandalf.Core.Abstractions.Extensions;
using Suspectus.Gandalf.Palantir.Client;

namespace Suspectus.Gandalf.Mithrandir.Api.Controllers;

[ApiController]
[Route("api/[controller]")]
public class AppProxyController : ControllerBase
{
    private readonly IDistributedCache _tokenCache;
    private readonly IPalantirClient _client;
    private readonly HttpClient _httpClient;

    public AppProxyController(IDistributedCache tokenCache, IPalantirClient client, IHttpClientFactory httpClientFactory)
    {
        _tokenCache = tokenCache;
        _client = client;
        _httpClient = httpClientFactory.CreateClient("test");
    }

    [HttpGet("{appId}/{*path}")]
    public async Task<IActionResult> Get(string appId, string path, CancellationToken cancellationToken)
    {
        var sessionCookie = Request.Cookies["MithrandirSession"];

        if (sessionCookie is null)
        {
            return Unauthorized("Session expired.");
        }

        var appInfoResult = await _client.Internal.App.GetInfo(appId);
        
        if (!appInfoResult.IsSuccess)
        {
            return BadRequest(appInfoResult.Match<string?>(_ => null, e => e.Message));
        }

        var appInfo = appInfoResult.GetValue();
        
        if (appInfo.IsMaster && path.StartsWith("api/internal"))
        {
            return NotFound();
        }

        var token = await GetToken(sessionCookie, appId, cancellationToken);

        if (token is null)
        {
            return Unauthorized("Session expired.");
        }
        
        _httpClient.BaseAddress = new Uri(appInfo.BaseAddress);
        _httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", token);
        
        try
        {
            var clientResponse = await _httpClient.GetAsync(path, cancellationToken);

            if (!clientResponse.IsSuccessStatusCode)
            {
                Response.StatusCode = (int)clientResponse.StatusCode;
            }
            
            foreach (var header in clientResponse.Content.Headers)
            {
                Response.Headers[header.Key] = header.Value.ToArray();
            }
            
            await Response.WriteAsync(await clientResponse.Content.ReadAsStringAsync(cancellationToken), cancellationToken: cancellationToken);
        }
        catch (HttpRequestException e)
        {
            Response.StatusCode = e.StatusCode is null ? (int)HttpStatusCode.InternalServerError : (int)e.StatusCode;
            await Response.WriteAsync(e.Message, cancellationToken: cancellationToken);
        }

        return new EmptyResult();
    }

    private async Task<string?> GetToken(string subjectId, string appId, CancellationToken cancellationToken)
    {
        var accessTokenCacheKey = GetCacheKey(subjectId, appId, "access_token");
        var accessToken = await _tokenCache.GetStringAsync(accessTokenCacheKey, cancellationToken);

        if (accessToken is not null)
        {
            return accessToken;
        }
        
        var refreshTokenCacheKey = GetCacheKey(subjectId, appId, "refresh_token");
        var refreshToken = await _tokenCache.GetStringAsync(refreshTokenCacheKey, cancellationToken);

        if (refreshToken is null)
        {
            return null;
        }
        
        var tokenRequestCommandResponse = await _client.Internal.Auth.Token(new TokenRequestCommand
        {
            GrandType = GrandType.RefreshToken,
            RefreshToken = refreshToken,
            ClientId = appId
        });

        if (tokenRequestCommandResponse.IsFaulted)
        {
            return null;
        }
        
        var tokenRequestResponse = tokenRequestCommandResponse.GetValue();
        
        if (tokenRequestResponse is null)
        {
            return null;
        }
        
        await _tokenCache.SetStringAsync(
            accessTokenCacheKey,
            tokenRequestResponse.AccessToken,
            new DistributedCacheEntryOptions
            {
                AbsoluteExpiration = tokenRequestResponse.AccessTokenExpiresAt.AddSeconds(-10),
            },
            cancellationToken
        );

        var refreshTokenCacheExpiration = tokenRequestResponse.RefreshTokenExpiresAt.AddSeconds(-10);
        
        await _tokenCache.SetStringAsync(
            refreshTokenCacheKey,
            tokenRequestResponse.RefreshToken,
            new DistributedCacheEntryOptions
            {
                AbsoluteExpiration = refreshTokenCacheExpiration,
            },
            cancellationToken
        );
        
        UpdateSessionCookie(subjectId, refreshTokenCacheExpiration);
        return tokenRequestResponse.AccessToken;
    }

    private string GetCacheKey(string subjectId, string appId, string tokenType)
    {
        return $"{subjectId}:{appId}:{tokenType}";
    }
    
    private void UpdateSessionCookie(string subjectId, DateTimeOffset refreshTokenCacheExpiration)
    {
        Response.Cookies.Append("MithrandirSession", subjectId, new CookieOptions
        {
            Secure = true,
            HttpOnly = true,
            SameSite = SameSiteMode.Lax,
            Expires = refreshTokenCacheExpiration
        });
    }
}