PranitInitial commit439f83c3 months agoHistory
833 lines · 32.4 KB · Zig
RawDownload
const std = @import("std");const builtin = @import("builtin");const debug_trace = @import("../shared/debug_trace.zig");const host = @import("../hosts/host.zig");const io_mod = @import("../shared/io.zig");const oauth = @import("oauth.zig");const oauth_session = @import("oauth_session.zig");const oauth_transport = @import("oauth_transport.zig");const secret = @import("secret.zig");const types = @import("../shared/types.zig"); pub const Source = types.CredentialSource; pub const CatalogPublicOnly = union(enum) {    no_credential,    fx_login_team_required,    fx_login_refresh_required,    credential_refresh_failed: Source,    authenticated_credential_rejected: Source,     fn credentialSource(self: CatalogPublicOnly) ?Source {        return switch (self) {            .no_credential => null,            .fx_login_team_required, .fx_login_refresh_required => .fx_login,            .credential_refresh_failed => |source| source,            .authenticated_credential_rejected => |source| source,        };    }}; pub const CatalogPublicOnlyReason = std.meta.Tag(CatalogPublicOnly); pub const CatalogAuthenticatedSource = enum {    vercel_oidc_token,    ai_gateway_api_key,    fx_login,    stored_key,     fn credentialSource(self: CatalogAuthenticatedSource) Source {        return switch (self) {            .vercel_oidc_token => .vercel_oidc_token,            .ai_gateway_api_key => .ai_gateway_api_key,            .fx_login => .fx_login,            .stored_key => .stored_key,        };    }}; /// A borrowed authorization decision for one model-catalog request. Public-only/// states cannot carry credential or team bytes; authenticated states carry the/// only values the request is allowed to send.pub const CatalogAccess = union(enum) {    public_only: CatalogPublicOnly,    authenticated: struct {        source: CatalogAuthenticatedSource,        credential: []const u8,        team_context: ?[]const u8,    },     pub fn credentialSource(self: CatalogAccess) ?Source {        return switch (self) {            .public_only => |access| access.credentialSource(),            .authenticated => |access| access.source.credentialSource(),        };    }     pub fn publicOnlyReason(self: CatalogAccess) ?CatalogPublicOnlyReason {        const access = self.publicOnly() orelse return null;        return std.meta.activeTag(access);    }     pub fn publicOnly(self: CatalogAccess) ?CatalogPublicOnly {        return switch (self) {            .public_only => |access| access,            .authenticated => null,        };    }     pub fn publicFallbackAfterRejection(self: CatalogAccess) ?CatalogAccess {        return switch (self) {            .public_only => null,            .authenticated => |access| .{                .public_only = .{                    .authenticated_credential_rejected = access.source.credentialSource(),                },            },        };    }     pub fn authorizationCredential(self: CatalogAccess) ?[]const u8 {        return switch (self) {            .public_only => null,            .authenticated => |access| access.credential,        };    }     pub fn teamContext(self: CatalogAccess) ?[]const u8 {        const team = switch (self) {            .public_only => return null,            .authenticated => |access| access.team_context orelse return null,        };        return if (team.len > 0) team else null;    }}; pub fn catalogAccessAt(credential: ?Credential, now_ms: i64) CatalogAccess {    const selected = credential orelse return .{ .public_only = .no_credential };    if (selected.source == .fx_login and selected.needsRefreshAt(now_ms)) {        return .{ .public_only = .fx_login_refresh_required };    }    return catalogAccessForCredential(        selected.source,        selected.token,        selected.gatewayTeam(),    );} pub fn catalogAccessAfterRefreshFailure(source: Source) CatalogAccess {    return .{        .public_only = .{            .credential_refresh_failed = source,        },    };} pub fn catalogAccessForCredential(    source: ?Source,    credential: []const u8,    team_context: ?[]const u8,) CatalogAccess {    const selected_source = source orelse return .{ .public_only = .no_credential };    const authenticated_source: CatalogAuthenticatedSource = switch (selected_source) {        .vercel_oidc_token => .vercel_oidc_token,        .ai_gateway_api_key => .ai_gateway_api_key,        .stored_key => .stored_key,        .fx_login => blk: {            const team = team_context orelse                return .{ .public_only = .fx_login_team_required };            if (!types.validGatewayTeam(team))                return .{ .public_only = .fx_login_team_required };            break :blk .fx_login;        },    };    return .{        .authenticated = .{            .source = authenticated_source,            .credential = credential,            .team_context = team_context,        },    };} /// Current native product copy. Store mechanics and availability come from the/// injected host port; Core retains the stable user-facing source name.pub const stored_key_backend_label = if (builtin.os.tag == .macos) "macOS Keychain" else "profile file"; /// Both modes resolve the same source set; the mode selects only whether an expired/// fx login session is refreshed first.pub const LoadMode = enum { stored, refresh_if_needed }; const FxLoginRefreshMode = enum { if_needed, force }; pub const missing_credential_message = "Fx needs access to Vercel AI Gateway. Run fx login to sign in, fx setup to use an API key, or set AI_GATEWAY_API_KEY.";pub const missing_interactive_credential_message = "Fx needs access to Vercel AI Gateway. Run /login to sign in, /setup to use an API key, or set AI_GATEWAY_API_KEY.";pub const unreadable_store_message = "Fx could not read the stored API key from " ++ stored_key_backend_label ++ ". A key may be saved but unreadable. Set FX_TRACE_LOG for the failing step, or set AI_GATEWAY_API_KEY."; pub const Credential = struct {    token: []u8,    source: Source,    team_id: ?[]u8 = null,    team_slug: ?[]u8 = null,    refresh_after_ms: ?i64 = null,     pub fn deinit(self: *Credential, alloc: std.mem.Allocator) void {        secret.zeroAndFree(alloc, self.token);        if (self.team_id) |team| alloc.free(team);        if (self.team_slug) |team| alloc.free(team);        self.* = undefined;    }     pub fn gatewayTeam(self: Credential) ?[]const u8 {        if (self.team_id) |team| return team;        return self.team_slug;    }     pub fn needsRefreshAt(self: Credential, now_ms: i64) bool {        const refresh_after_ms = self.refresh_after_ms orelse return false;        return refresh_after_ms <= now_ms;    }}; pub const StoredKeyReadStatus = enum {    not_attempted,    not_found,    unavailable,}; pub const Resolution = struct {    credential: ?Credential = null,    stored_key_status: StoredKeyReadStatus = .not_attempted,}; /// The single credential resolution method. Walks source precedence, then falls back to/// the stored key, reporting why that store was silent when it produced nothing.pub fn resolve(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,    secret_store: host.SecretStore,    mode: LoadMode,) !Resolution {    return resolvePreferring(alloc, transport, secret_store, mode, null);} /// `preferred` is the source the user last chose in the hub. It wins over the/// precedence order below, including over the environment, because it is an/// explicit choice rather than a default. A preferred source that no longer/// resolves falls through to precedence instead of failing.pub fn resolvePreferring(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,    secret_store: host.SecretStore,    mode: LoadMode,    preferred: ?Source,) !Resolution {    if (preferred) |source| {        if (source != .stored_key or !secret_store.isDisabled()) {            const chosen = loadPreferredSource(alloc, transport, secret_store, mode, source) catch |err| blk: {                if (err == error.OutOfMemory) return err;                debug_trace.logf("auth", "preferred source load failed source={t} err={s}", .{ source, @errorName(err) });                break :blk null;            };            if (chosen) |credential| return .{ .credential = credential };            debug_trace.logf("auth", "preferred source unavailable source={t}; using precedence", .{source});        }    }     if (try loadSource(alloc, transport, secret_store, .vercel_oidc_token)) |credential| return .{ .credential = credential };    if (try loadSource(alloc, transport, secret_store, .ai_gateway_api_key)) |credential| return .{ .credential = credential };     const fx_login = switch (mode) {        .stored => try loadStoredFxLoginCredential(alloc),        .refresh_if_needed => try loadFxLoginCredential(alloc, transport),    };    if (fx_login) |credential| return .{ .credential = credential };     if (secret_store.isDisabled()) return .{};     var status: StoredKeyReadStatus = .not_found;    const stored = loadSource(alloc, transport, secret_store, .stored_key) catch |err| blk: {        if (err == error.OutOfMemory) return err;        status = .unavailable;        debug_trace.logf("auth", "stored key load failed err={s} status={t}", .{ @errorName(err), status });        break :blk null;    };    if (stored) |credential| return .{ .credential = credential };    return .{ .stored_key_status = status };} /// `loadSource` always refreshes an expired fx login, which `.stored` mode/// forbids: a diagnostic must not rewrite the session file or make an OAuth/// request. Honour the mode for the preferred source too.fn loadPreferredSource(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,    secret_store: host.SecretStore,    mode: LoadMode,    source: Source,) !?Credential {    if (source != .fx_login) return loadSource(alloc, transport, secret_store, source);    return switch (mode) {        .stored => loadStoredFxLoginCredential(alloc),        .refresh_if_needed => loadFxLoginCredential(alloc, transport),    };} pub fn loadSource(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,    secret_store: host.SecretStore,    source: Source,) !?Credential {    return switch (source) {        .vercel_oidc_token => loadEnvCredential(alloc, "VERCEL_OIDC_TOKEN", source),        .ai_gateway_api_key => loadEnvCredential(alloc, "AI_GATEWAY_API_KEY", source),        .fx_login => loadFxLoginCredential(alloc, transport),        .stored_key => loadStoredKeyCredential(alloc, secret_store),    };} pub fn sourceExists(    alloc: std.mem.Allocator,    secret_store: host.SecretStore,    source: Source,) !bool {    return switch (source) {        .vercel_oidc_token => nonEmptyEnvValue("VERCEL_OIDC_TOKEN") != null,        .ai_gateway_api_key => nonEmptyEnvValue("AI_GATEWAY_API_KEY") != null,        .fx_login => blk: {            const loaded = oauth_session.load(alloc) catch |err| switch (err) {                error.OutOfMemory => return err,                else => {                    debug_trace.logf("auth", "source probe failed source=fx_login err={s}", .{@errorName(err)});                    break :blk false;                },            };            var session = loaded orelse break :blk false;            defer session.deinit(alloc);            break :blk true;        },        .stored_key => blk: {            if (secret_store.isDisabled()) break :blk false;            const stored = secret_store.load(alloc) catch |err| switch (err) {                error.OutOfMemory => return err,                else => {                    debug_trace.logf("auth", "source probe failed source=stored_key err={s}", .{@errorName(err)});                    break :blk false;                },            };            const value = stored orelse break :blk false;            secret.zeroAndFree(alloc, value);            break :blk true;        },    };} fn loadEnvCredential(    alloc: std.mem.Allocator,    name: []const u8,    source: Source,) !?Credential {    const value = nonEmptyEnvValue(name) orelse return null;    return .{        .token = try alloc.dupe(u8, value),        .source = source,    };} fn loadStoredKeyCredential(    alloc: std.mem.Allocator,    secret_store: host.SecretStore,) !?Credential {    if (secret_store.isDisabled()) return null;    const value = (try secret_store.load(alloc)) orelse return null;    return .{ .token = value, .source = .stored_key };} fn nonEmptyEnvValue(name: []const u8) ?[]const u8 {    return nonEmptyValue(io_mod.getenv(name));} fn nonEmptyValue(value: ?[]const u8) ?[]const u8 {    const raw = value orelse return null;    if (std.mem.trim(u8, raw, " \t\r\n").len == 0) return null;    return raw;} pub fn loadFxLoginCredential(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,) !?Credential {    var session = (try oauth_session.load(alloc)) orelse return null;    defer session.deinit(alloc);     if (session.expired(io_mod.milliTimestamp())) {        return refreshFxLoginCredentialLocked(alloc, transport, .if_needed);    }     return takeCredentialFromSession(&session, null);} fn loadStoredFxLoginCredential(alloc: std.mem.Allocator) !?Credential {    var session = (try oauth_session.load(alloc)) orelse return null;    defer session.deinit(alloc);    return takeCredentialFromSession(&session, null);} pub fn refreshFxLoginCredential(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,) !?Credential {    return refreshFxLoginCredentialLocked(alloc, transport, .force);} fn refreshFxLoginCredentialLocked(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,    mode: FxLoginRefreshMode,) !?Credential {    var mutation = (try oauth_session.beginExistingMutation()) orelse return null;    defer mutation.deinit();     var session = (try mutation.load(alloc)) orelse return null;    defer session.deinit(alloc);     var refreshed_at_ms: ?i64 = null;    if (mode == .force or session.expired(io_mod.milliTimestamp())) {        try refreshFxSession(alloc, transport, &mutation, &session);        refreshed_at_ms = io_mod.milliTimestamp();    }    return takeCredentialFromSession(&session, refreshed_at_ms);} pub fn refreshFxSession(    alloc: std.mem.Allocator,    transport: oauth_transport.Provider,    mutation: *oauth_session.Mutation,    session: *oauth_session.Session,) !void {    const issuer_url = session.issuer;    var metadata = try oauth.discover(alloc, transport, issuer_url);    defer metadata.deinit(alloc);    try oauth_session.validateE2EEndpoint(issuer_url, metadata.token_endpoint);     var refreshed = try oauth.refreshToken(        alloc,        transport,        metadata,        session.client_id,        session.refresh_token,    );    defer refreshed.deinit(alloc);     secret.zeroAndFree(alloc, session.access_token);    session.access_token = refreshed.access_token;    refreshed.access_token = &.{};     if (refreshed.refresh_token) |value| {        secret.zeroAndFree(alloc, session.refresh_token);        session.refresh_token = value;        refreshed.refresh_token = null;    }     alloc.free(session.scope);    session.scope = refreshed.scope;    refreshed.scope = &.{};     alloc.free(session.token_type);    session.token_type = refreshed.token_type;    refreshed.token_type = &.{};     session.expires_at_ms = try oauth.expiry_timestamp_ms(io_mod.milliTimestamp(), refreshed.expires_in);    try mutation.save(alloc, session.*);} fn takeCredentialFromSession(session: *oauth_session.Session, refreshed_at_ms: ?i64) Credential {    const token = session.access_token;    session.access_token = &.{};    const team_id = session.team_id;    session.team_id = null;    const team_slug = session.team_slug;    session.team_slug = null;    return .{        .token = token,        .source = .fx_login,        .team_id = team_id,        .team_slug = team_slug,        .refresh_after_ms = credentialRefreshAfterMs(session.expires_at_ms, refreshed_at_ms),    };} fn credentialRefreshAfterMs(expires_at_ms: i64, refreshed_at_ms: ?i64) i64 {    const refresh_after_ms = oauth_session.refresh_deadline_ms(expires_at_ms);    const refreshed_at = refreshed_at_ms orelse return refresh_after_ms;    if (refresh_after_ms > refreshed_at) return refresh_after_ms;    return expires_at_ms;} pub fn sourceLabel(source: Source) []const u8 {    return switch (source) {        .vercel_oidc_token => "VERCEL_OIDC_TOKEN",        .ai_gateway_api_key => "AI_GATEWAY_API_KEY",        .fx_login => "fx login",        .stored_key => "stored API key (" ++ stored_key_backend_label ++ ")",    };} pub fn sourceRefreshable(source: Source) bool {    return source == .fx_login;} test "stored key label discloses the backend that answered" {    try std.testing.expect(std.mem.find(u8, sourceLabel(.stored_key), stored_key_backend_label) != null);    try std.testing.expect(std.mem.find(u8, unreadable_store_message, stored_key_backend_label) != null);    for ([_]Source{ .vercel_oidc_token, .ai_gateway_api_key, .fx_login }) |source| {        try std.testing.expect(!std.mem.eql(u8, sourceLabel(source), sourceLabel(.stored_key)));    }} test "missing credential messages use surface commands in preferred order" {    const cli_login = std.mem.find(u8, missing_credential_message, "fx login").?;    const cli_setup = std.mem.find(u8, missing_credential_message, "fx setup").?;    const cli_env = std.mem.find(u8, missing_credential_message, "AI_GATEWAY_API_KEY").?;     try std.testing.expect(cli_login < cli_setup);    try std.testing.expect(cli_setup < cli_env);     const tui_login = std.mem.find(u8, missing_interactive_credential_message, "/login").?;    const tui_setup = std.mem.find(u8, missing_interactive_credential_message, "/setup").?;    const tui_env = std.mem.find(u8, missing_interactive_credential_message, "AI_GATEWAY_API_KEY").?; 

Lines 1–500 of 833.