feat: API key management and fixes, workload tier updates, and admin CLI improvements

.github/workflows/nilcc-api-ci.yml

@@ -27,7 +27,7 @@ jobs:
           package_json_file: nilcc-api/package.json
       - run: pnpm install
       - run: pnpm exec biome ci
-      - run: tsc
+      - run: pnpm exec tsc
 
   test:
     needs: check

.gitignore

@@ -11,3 +11,4 @@ postgres_data/
 localstack_data/
 nilcc-attester/gpu-attester/verifier.log
 .env
+.vscode/

nilcc-admin-cli/src/api.rs

@@ -40,6 +40,16 @@ impl ApiClient {
         Self::handle_response(response)
     }
 
+    pub fn put<T, O>(&self, path: &str, request: &T) -> Result<O, RequestError>
+    where
+        T: Serialize,
+        O: DeserializeOwned,
+    {
+        let url = self.make_url(path);
+        let response = self.client.put(url).json(request).send()?;
+        Self::handle_response(response)
+    }
+
     fn handle_response<O>(response: Response) -> Result<O, RequestError>
     where
         O: DeserializeOwned,

nilcc-admin-cli/src/main.rs

@@ -1,5 +1,5 @@
 use crate::api::{ApiClient, RequestError};
-use clap::{Args, Parser, Subcommand};
+use clap::{Args, Parser, Subcommand, ValueEnum};
 use serde_json::json;
 use uuid::Uuid;
 
@@ -31,6 +31,10 @@ enum Command {
     #[clap(subcommand)]
     Accounts(AccountsCommand),
 
+    /// Manage account API keys.
+    #[clap(subcommand)]
+    ApiKeys(ApiKeysCommand),
+
     /// Manage tiers.
     #[clap(subcommand)]
     Tiers(TiersCommand),
@@ -59,6 +63,44 @@ enum AccountsCommand {
     Rename(RenameAccountArgs),
 }
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
+enum ApiKeyTypeArg {
+    #[value(name = "account-admin")]
+    AccountAdmin,
+    #[value(name = "user")]
+    User,
+}
+
+impl ApiKeyTypeArg {
+    fn as_str(self) -> &'static str {
+        match self {
+            Self::AccountAdmin => "account-admin",
+            Self::User => "user",
+        }
+    }
+}
+
+#[derive(Subcommand)]
+enum ApiKeysCommand {
+    /// Create an API key for an account.
+    Create(CreateApiKeyArgs),
+
+    /// List the API keys for an account.
+    List {
+        /// The account id.
+        account_id: Uuid,
+    },
+
+    /// Update an API key.
+    Update(UpdateApiKeyArgs),
+
+    /// Delete an API key.
+    Delete {
+        /// The identifier of the API key to be deleted.
+        id: Uuid,
+    },
+}
+
 #[derive(Args)]
 struct CreateAccountArgs {
     /// The account name.
@@ -90,6 +132,34 @@ struct RenameAccountArgs {
     name: String,
 }
 
+#[derive(Args)]
+struct CreateApiKeyArgs {
+    /// The account id.
+    account_id: Uuid,
+
+    /// The type of API key to create.
+    #[clap(long, value_enum)]
+    key_type: ApiKeyTypeArg,
+
+    /// Create the key as inactive.
+    #[clap(long, default_value_t = false)]
+    inactive: bool,
+}
+
+#[derive(Args)]
+struct UpdateApiKeyArgs {
+    /// The identifier of the API key to update.
+    id: Uuid,
+
+    /// Set the API key type.
+    #[clap(long, value_enum)]
+    key_type: Option<ApiKeyTypeArg>,
+
+    /// Set whether the API key is active.
+    #[clap(long)]
+    active: Option<bool>,
+}
+
 #[derive(Subcommand)]
 enum TiersCommand {
     /// Create a tier.
@@ -98,6 +168,9 @@ enum TiersCommand {
     /// List tiers.
     List,
 
+    /// Update a tier.
+    Update(UpdateTierArgs),
+
     /// Delete a tier.
     Delete {
         /// The identifier of the tier to be deleted.
@@ -131,6 +204,35 @@ struct CreateTierArgs {
     disk_gb: u64,
 }
 
+#[derive(Args)]
+struct UpdateTierArgs {
+    /// The identifier of the tier to be updated.
+    id: Uuid,
+
+    /// The tier name.
+    name: String,
+
+    /// The tier cost in USD/minute.
+    #[clap(long)]
+    cost: f64,
+
+    /// The number of CPUs that are granted with this tier.
+    #[clap(long)]
+    cpus: u64,
+
+    /// The number of GPUs that are granted with this tier.
+    #[clap(long)]
+    gpus: u64,
+
+    /// The amount of memory in this tier, in MBs.
+    #[clap(long)]
+    memory_mb: u64,
+
+    /// The amount of disk in this tier, in GBs.
+    #[clap(long)]
+    disk_gb: u64,
+}
+
 #[derive(Subcommand)]
 enum ArtifactsCommand {
     /// Enable an artifact version.
@@ -192,6 +294,35 @@ impl Runner {
         self.client.post("/api/v1/accounts/update", &request)
     }
 
+    fn create_api_key(&self, args: CreateApiKeyArgs) -> Result<serde_json::Value, RequestError> {
+        let CreateApiKeyArgs { account_id, key_type, inactive } = args;
+        let request = models::api_keys::CreateApiKeyRequest {
+            account_id,
+            r#type: key_type.as_str().to_string(),
+            active: !inactive,
+        };
+        self.client.post("/api/v1/api-keys/create", &request)
+    }
+
+    fn list_api_keys(&self, account_id: Uuid) -> Result<serde_json::Value, RequestError> {
+        self.client.get(&format!("/api/v1/api-keys/account/{account_id}"))
+    }
+
+    fn update_api_key(&self, args: UpdateApiKeyArgs) -> Result<serde_json::Value, RequestError> {
+        let UpdateApiKeyArgs { id, key_type, active } = args;
+        let request = models::api_keys::UpdateApiKeyRequest {
+            id,
+            r#type: key_type.map(|value| value.as_str().to_string()),
+            active,
+        };
+        self.client.put("/api/v1/api-keys/update", &request)
+    }
+
+    fn delete_api_key(&self, id: Uuid) -> Result<serde_json::Value, RequestError> {
+        let request = models::api_keys::DeleteApiKeyRequest { id };
+        self.client.post("/api/v1/api-keys/delete", &request)
+    }
+
     fn create_tier(&self, args: CreateTierArgs) -> Result<serde_json::Value, RequestError> {
         let CreateTierArgs { name, cost, cpus, gpus, memory_mb, disk_gb } = args;
         let request = models::tiers::CreateTierRequest { name, cost, cpus, gpus, memory_mb, disk_gb };
@@ -202,6 +333,12 @@ impl Runner {
         self.client.get("/api/v1/workload-tiers/list")
     }
 
+    fn update_tier(&self, args: UpdateTierArgs) -> Result<serde_json::Value, RequestError> {
+        let UpdateTierArgs { id, name, cost, cpus, gpus, memory_mb, disk_gb } = args;
+        let request = models::tiers::UpdateTierRequest { tier_id: id, name, cost, cpus, gpus, memory_mb, disk_gb };
+        self.client.put("/api/v1/workload-tiers/update", &request)
+    }
+
     fn delete_tier(&self, tier_id: Uuid) -> Result<serde_json::Value, RequestError> {
         let request = models::tiers::DeleteTierRequest { tier_id };
         self.client.post("/api/v1/workload-tiers/delete", &request)
@@ -239,8 +376,13 @@ fn main() {
         Command::Accounts(AccountsCommand::List) => runner.list_accounts(),
         Command::Accounts(AccountsCommand::AddBalance(args)) => runner.add_balance(args),
         Command::Accounts(AccountsCommand::Rename(args)) => runner.rename(args),
+        Command::ApiKeys(ApiKeysCommand::Create(args)) => runner.create_api_key(args),
+        Command::ApiKeys(ApiKeysCommand::List { account_id }) => runner.list_api_keys(account_id),
+        Command::ApiKeys(ApiKeysCommand::Update(args)) => runner.update_api_key(args),
+        Command::ApiKeys(ApiKeysCommand::Delete { id }) => runner.delete_api_key(id),
         Command::Tiers(TiersCommand::Create(args)) => runner.create_tier(args),
         Command::Tiers(TiersCommand::List) => runner.list_tiers(),
+        Command::Tiers(TiersCommand::Update(args)) => runner.update_tier(args),
         Command::Tiers(TiersCommand::Delete { id }) => runner.delete_tier(id),
         Command::Artifacts(ArtifactsCommand::Enable { version }) => runner.enable_artifact_version(version),
         Command::Artifacts(ArtifactsCommand::List) => runner.list_artifact_versions(),

nilcc-admin-cli/src/models.rs

@@ -27,6 +27,34 @@ pub mod accounts {
     }
 }
 
+pub mod api_keys {
+    use super::*;
+
+    #[derive(Serialize)]
+    #[serde(rename_all = "camelCase")]
+    pub struct CreateApiKeyRequest {
+        pub account_id: Uuid,
+        pub r#type: String,
+        pub active: bool,
+    }
+
+    #[derive(Serialize)]
+    #[serde(rename_all = "camelCase")]
+    pub struct UpdateApiKeyRequest {
+        pub id: Uuid,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub r#type: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        pub active: Option<bool>,
+    }
+
+    #[derive(Serialize)]
+    #[serde(rename_all = "camelCase")]
+    pub struct DeleteApiKeyRequest {
+        pub id: Uuid,
+    }
+}
+
 pub mod tiers {
     use super::*;
 
@@ -41,6 +69,18 @@ pub mod tiers {
         pub disk_gb: u64,
     }
 
+    #[derive(Serialize)]
+    #[serde(rename_all = "camelCase")]
+    pub struct UpdateTierRequest {
+        pub tier_id: Uuid,
+        pub name: String,
+        pub cost: f64,
+        pub cpus: u64,
+        pub gpus: u64,
+        pub memory_mb: u64,
+        pub disk_gb: u64,
+    }
+
     #[derive(Serialize)]
     #[serde(rename_all = "camelCase")]
     pub struct DeleteTierRequest {

nilcc-api/migrations/1772193140110-ApiKeys.ts

@@ -6,7 +6,7 @@ export class ApiKeys1772193140110 implements MigrationInterface {
   public async up(queryRunner: QueryRunner): Promise<void> {
     await queryRunner.query(`
       CREATE TABLE api_keys (
-        id UUID PRIMARY KEY,
+        id VARCHAR PRIMARY KEY,
         account_id VARCHAR NOT NULL,
         type VARCHAR NOT NULL,
         active BOOLEAN NOT NULL DEFAULT true,

nilcc-api/migrations/1773000000000-WalletAuth.ts

@@ -13,15 +13,14 @@ export class WalletAuth1773000000000 implements MigrationInterface {
     await queryRunner.query(`
       INSERT INTO api_keys (id, account_id, type, active, created_at, updated_at)
       SELECT
-        CAST(api_token AS UUID),
+        api_token,
         id,
         'account-admin',
         true,
         NOW(),
         NOW()
       FROM accounts
       WHERE api_token IS NOT NULL
-        AND api_token ~ '^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$'
     `);
 
     await queryRunner.query("ALTER TABLE accounts DROP COLUMN api_token");

nilcc-api/migrations/1775000000000-ApiKeyIdVarchar.ts

@@ -0,0 +1,36 @@
+import type { MigrationInterface, QueryRunner } from "typeorm";
+
+const UUID_PATTERN =
+  "^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$";
+
+export class ApiKeyIdVarchar1775000000000 implements MigrationInterface {
+  name = "ApiKeyIdVarchar1775000000000";
+
+  public async up(queryRunner: QueryRunner): Promise<void> {
+    await queryRunner.query(`
+      ALTER TABLE api_keys
+      ALTER COLUMN id TYPE VARCHAR USING id::VARCHAR
+    `);
+  }
+
+  public async down(queryRunner: QueryRunner): Promise<void> {
+    await queryRunner.query(`
+      DO $$
+      BEGIN
+        IF EXISTS (
+          SELECT 1
+          FROM api_keys
+          WHERE id !~* '${UUID_PATTERN}'
+        ) THEN
+          RAISE EXCEPTION 'cannot narrow api_keys.id back to UUID while non-UUID API keys exist';
+        END IF;
+      END
+      $$;
+    `);
+
+    await queryRunner.query(`
+      ALTER TABLE api_keys
+      ALTER COLUMN id TYPE UUID USING id::UUID
+    `);
+  }
+}

nilcc-api/src/account/account.controller.ts

@@ -3,9 +3,9 @@ import { resolver } from "hono-openapi/zod";
 import { z } from "zod";
 import {
   accountIdentityAdminAuthentication,
+  accountIdentityAuthentication,
   adminAuthentication,
   assertCanManageIdentityAccount,
-  jwtAuthentication,
 } from "#/common/auth";
 import { EntityNotFound } from "#/common/errors";
 import { microdollarsToUsd } from "#/common/nil";
@@ -169,7 +169,7 @@ export function me(options: ControllerOptions) {
         ...OpenApiSpecCommonErrorResponses,
       },
     }),
-    jwtAuthentication(bindings),
+    accountIdentityAuthentication(bindings),
     async (c) => {
       const account = c.get("account");
       const outputAccount = accountMapper.entityToResponse(account);