ClawdDotNet/src/ClawdDotNet.Core/Mcp/McpGrants.cs

using System.Collections.Concurrent;
using System.Security.Cryptography;
using System.Text;
using ClawdDotNet.Core.Storage;

namespace ClawdDotNet.Core.Mcp;

/// <summary>
/// #713: What a token to the MCP server stands for — one agent, and whose run its calls
/// belong to. A call through the server acts as this agent: with its tools, its rights,
/// and nothing more.
/// </summary>
public sealed record McpGrant
{
    public string AgentId { get; init; } = "";

    /// <summary>The run the calls are audited under.</summary>
    public string RunId { get; init; } = "";

    /// <summary>The task the run works on, if any.</summary>
    public string? TaskId { get; init; }

    public DateTime ExpiresAtUtc { get; init; }

    /// <summary>A token a human issued for their own client — not a run's.</summary>
    public bool IsPersonal { get; init; }

    /// <summary>The personal token's id; empty for a run's token.</summary>
    public string TokenId { get; init; } = "";
}

/// <summary>#713: A personal token as it is listed — never the token itself.</summary>
public sealed record McpPersonalToken(
    string Id, string AgentId, string Label, DateTime CreatedAtUtc, DateTime ExpiresAtUtc, DateTime? RevokedAtUtc)
{
    public bool IsActive(DateTime nowUtc) => RevokedAtUtc is null && nowUtc < ExpiresAtUtc;
}

/// <summary>
/// #713: The tokens of the MCP server.
///
/// <para><b>A run's token lives as long as the run.</b> <see cref="OpenRun"/> hands out a
/// token and a handle; disposing the handle — at the end of the run — makes the token
/// worthless, and it has an expiry of its own in case nobody disposes it. Run tokens exist
/// only in memory: after a restart there is no run they could belong to.</para>
///
/// <para><b>A personal token</b> is for a human's own client (their Claude Code on the
/// instance's taskboard). It acts as one agent, has an expiry, can be revoked, and is kept
/// only as a hash: shown once, when it is issued.</para>
///
/// <para>Tokens are compared by their SHA-256 — looked up, not compared character by
/// character, so the time a check takes says nothing about how close a guess was.</para>
/// </summary>
public sealed class McpGrants
{
    /// <summary>A run that forgets to close its token loses it after this long anyway.</summary>
    public static readonly TimeSpan DefaultRunLifetime = TimeSpan.FromHours(4);

    /// <summary>The longest a personal token may live.</summary>
    public const int MaxPersonalDays = 90;

    private readonly ConcurrentDictionary<string, McpGrant> _runs = new(StringComparer.Ordinal);
    private readonly SqliteStorage? _storage;
    private readonly TimeProvider _clock;

    public McpGrants(SqliteStorage? storage = null, TimeProvider? clock = null)
    {
        _storage = storage;
        _clock = clock ?? TimeProvider.System;
    }

    /// <summary>
    /// A token for one run of <paramref name="agentId"/>. Dispose the handle when the run ends.
    /// </summary>
    public McpRunToken OpenRun(string agentId, string runId, string? taskId = null, TimeSpan? lifetime = null)
    {
        var token = NewToken();
        var hash = Hash(token);
        _runs[hash] = new McpGrant
        {
            AgentId = agentId,
            RunId = runId,
            TaskId = taskId,
            ExpiresAtUtc = _clock.GetUtcNow().UtcDateTime + (lifetime ?? DefaultRunLifetime)
        };
        return new McpRunToken(token, () => _runs.TryRemove(hash, out _));
    }

    /// <summary>What a token stands for, or <c>null</c> — unknown, expired or revoked.</summary>
    public async Task<McpGrant?> ResolveAsync(string? token, CancellationToken ct)
    {
        if (string.IsNullOrWhiteSpace(token))
            return null;

        var hash = Hash(token.Trim());
        var now = _clock.GetUtcNow().UtcDateTime;

        if (_runs.TryGetValue(hash, out var run))
        {
            if (now < run.ExpiresAtUtc)
                return run;
            _runs.TryRemove(hash, out _);
            return null;
        }

        if (_storage is null)
            return null;

        await using var conn = await _storage.OpenConnectionAsync(ct);
        using var cmd = conn.CreateCommand();
        cmd.CommandText = "SELECT Id, AgentId, ExpiresAt, RevokedAt FROM McpTokens WHERE Hash = @hash";
        cmd.Parameters.AddWithValue("@hash", hash);
        await using var reader = await cmd.ExecuteReaderAsync(ct);
        if (!await reader.ReadAsync(ct) || !reader.IsDBNull(3))
            return null;

        var expires = Parse(reader.GetString(2));
        if (now >= expires)
            return null;

        var id = reader.GetString(0);
        return new McpGrant
        {
            AgentId = reader.GetString(1),
            // Each personal token is its own "run" in the audit: its calls stay together.
            RunId = "mcp-" + id,
            ExpiresAtUtc = expires,
            IsPersonal = true,
            TokenId = id
        };
    }

    /// <summary>Issues a personal token. Returns it in clear — the only time it exists so.</summary>
    public async Task<(McpPersonalToken Info, string Token)> IssuePersonalAsync(
        string agentId, string label, int days, CancellationToken ct)
    {
        var store = _storage ?? throw new InvalidOperationException("Personal tokens need the instance database.");
        var token = NewToken();
        var now = _clock.GetUtcNow().UtcDateTime;
        var info = new McpPersonalToken(
            Guid.NewGuid().ToString("N")[..12], agentId, label.Trim(), now,
            now.AddDays(Math.Clamp(days, 1, MaxPersonalDays)), null);

        await store.WriteAsync(async conn =>
        {
            using var cmd = conn.CreateCommand();
            cmd.CommandText = """
                INSERT INTO McpTokens (Id, Hash, AgentId, Label, CreatedAt, ExpiresAt, RevokedAt)
                VALUES (@id, @hash, @agent, @label, @created, @expires, NULL)
                """;
            cmd.Parameters.AddWithValue("@id", info.Id);
            cmd.Parameters.AddWithValue("@hash", Hash(token));
            cmd.Parameters.AddWithValue("@agent", info.AgentId);
            cmd.Parameters.AddWithValue("@label", info.Label);
            cmd.Parameters.AddWithValue("@created", Format(info.CreatedAtUtc));
            cmd.Parameters.AddWithValue("@expires", Format(info.ExpiresAtUtc));
            await cmd.ExecuteNonQueryAsync(ct);
        }, ct);

        return (info, token);
    }

    public async Task<IReadOnlyList<McpPersonalToken>> ListPersonalAsync(CancellationToken ct)
    {
        if (_storage is null)
            return [];
        await using var conn = await _storage.OpenConnectionAsync(ct);
        using var cmd = conn.CreateCommand();
        cmd.CommandText = "SELECT Id, AgentId, Label, CreatedAt, ExpiresAt, RevokedAt FROM McpTokens ORDER BY CreatedAt DESC";
        var result = new List<McpPersonalToken>();
        await using var reader = await cmd.ExecuteReaderAsync(ct);
        while (await reader.ReadAsync(ct))
            result.Add(new McpPersonalToken(
                reader.GetString(0), reader.GetString(1), reader.GetString(2),
                Parse(reader.GetString(3)), Parse(reader.GetString(4)),
                reader.IsDBNull(5) ? null : Parse(reader.GetString(5))));
        return result;
    }

    public async Task<bool> RevokeAsync(string id, CancellationToken ct)
    {
        if (_storage is null)
            return false;
        return await _storage.WriteAsync(async conn =>
        {
            using var cmd = conn.CreateCommand();
            cmd.CommandText = "UPDATE McpTokens SET RevokedAt = @now WHERE Id = @id AND RevokedAt IS NULL";
            cmd.Parameters.AddWithValue("@now", Format(_clock.GetUtcNow().UtcDateTime));
            cmd.Parameters.AddWithValue("@id", id.Trim());
            return await cmd.ExecuteNonQueryAsync(ct) > 0;
        }, ct);
    }

    /// <summary>32 random bytes, URL-safe — "cdn_" in front, so a leaked one is recognisable.</summary>
    private static string NewToken()
        => "cdn_" + Convert.ToBase64String(RandomNumberGenerator.GetBytes(32))
            .TrimEnd('=').Replace('+', '-').Replace('/', '_');

    private static string Hash(string token)
        => Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(token)));

    private static string Format(DateTime value) => value.ToUniversalTime().ToString("O");

    private static DateTime Parse(string value)
        => DateTime.TryParse(value, null, System.Globalization.DateTimeStyles.RoundtripKind, out var dt)
            ? dt.ToUniversalTime()
            : DateTime.MinValue;
}

/// <summary>#713: A run's token and the handle that ends it.</summary>
public sealed class McpRunToken(string token, Action close) : IDisposable
{
    private int _closed;

    /// <summary>The token — to hand over by file (0600) or environment, never on a command line.</summary>
    public string Token { get; } = token;

    public void Dispose()
    {
        if (Interlocked.Exchange(ref _closed, 1) == 0)
            close();
    }
}