feat: update gcp key handler for advanced tee architecture

.gitignore

@@ -131,3 +131,5 @@ dist
 **/.yarn/build-state.yml
 **/.yarn/install-state.gz
 .pnp.*
+
+.vscode

src/key/handlers/AzureSKRSidecarKeyHandler/index.ts → src/key/handlers/AzureSKRKeyHandler/index.ts

@@ -1,30 +1,69 @@
+import { exec } from "child_process";
 import { privateDecrypt, constants } from "crypto";
+import fs from "fs/promises";
+import { promisify } from "util";
 
 import axios from "axios";
 import jwkToPem from "jwk-to-pem";
 
-import { getConfigFromEnv } from "./utils";
+import { generateKeys, getConfigFromEnv } from "./utils";
 
 import { getLogger } from "../../../logging";
 
 import type { KeyHandler } from "@fr0ntier-x/polaris-sdk";
 
-export class AzureSKRSidecarKeyHandler implements KeyHandler {
+export class AzureSKRKeyHandler implements KeyHandler {
   private privateKey: string | undefined;
   private publicKey: string | undefined;
 
   constructor() {}
 
   private async init() {
-    const { maxSKRRequestRetries, skrRetryInterval, keyReleaseEndpoint, ...skrConfig } = getConfigFromEnv();
+    const { maxSKRRequestRetries, skrRetryInterval, keyReleaseEndpoint, advancedTEECheckEnabled, ...skrConfig } =
+      getConfigFromEnv();
 
     let attempt = 0;
     while (attempt < maxSKRRequestRetries) {
       try {
         attempt++;
-        const { data } = await axios.post(`${keyReleaseEndpoint}/key/release`, skrConfig);
+        let requestConfig = { ...skrConfig };
+        if (advancedTEECheckEnabled) {
+          try {
+            const attestationCommand = `AzureAttestSKR -a ${skrConfig.maa_endpoint} -g`.replace(/\r?\n|\r/g, " ");
+
+            const execAsync = promisify(exec);
+            const { stdout: attestationToken } = await execAsync(attestationCommand);
+
+            const pcrOutput = await fs.readFile("/tmp/PCR_output_file.pcrs", "base64");
+            const messageOutput = await fs.readFile("/tmp/message_output_file.msg", "base64");
+            const signatureOutput = await fs.readFile("/tmp/signature_output_file.sig", "base64");
+
+            requestConfig = {
+              ...requestConfig,
+              attestation_token: attestationToken,
+              pcr_file_content: pcrOutput,
+              message_file_content: messageOutput,
+              signature_file_content: signatureOutput,
+            };
+          } catch (error: any) {
+            throw new Error(`Failed to read attestation files: ${error.message}`);
+          }
+        }
+
+        const { data } = await axios.post(`${keyReleaseEndpoint}/key/release`, requestConfig);
+
+        if (advancedTEECheckEnabled) {
+          if (!data.cipherText) {
+            throw new Error("Expected cipherText in response for advanced TEE check");
+          }
+
+          const decryptCommand = `AzureAttestSKR -d ${data.cipherText}`;
+
+          const { privateKey, publicKey } = await generateKeys(decryptCommand);
 
-        if (data.key) {
+          this.privateKey = privateKey;
+          this.publicKey = publicKey;
+        } else if (data.key) {
           const key = JSON.parse(data.key);
           this.privateKey = jwkToPem(key as jwkToPem.JWK, { private: true });
           this.publicKey = jwkToPem(key as jwkToPem.JWK);

src/key/handlers/AzureSKRSidecarKeyHandler/types/index.ts → src/key/handlers/AzureSKRKeyHandler/types/index.ts

@@ -6,4 +6,9 @@ export interface AzureSKRSidecarKeyConfig {
   access_token?: string;
   maxSKRRequestRetries: number;
   skrRetryInterval: number;
+  advancedTEECheckEnabled?: boolean;
+  attestation_token?: string;
+  pcr_file_content?: string;
+  message_file_content?: string;
+  signature_file_content?: string;
 }

src/key/handlers/AzureSKRKeyHandler/utils/index.ts

@@ -0,0 +1,68 @@
+import { exec } from "child_process";
+import { createPrivateKey, createPublicKey } from "crypto";
+import { promisify } from "util";
+
+import type { AzureSKRSidecarKeyConfig } from "../types";
+import type { KeyObject } from "crypto";
+
+export const getConfigFromEnv = (): AzureSKRSidecarKeyConfig => {
+  const keyReleaseEndpoint = process.env.POLARIS_CONTAINER_AZURE_SKR_KEY_RELEASE_ENDPOINT || "http://localhost:8080";
+  const maaEndpoint = process.env.POLARIS_CONTAINER_AZURE_SKR_MAA_ENDPOINT;
+  const akvEndpoint = process.env.POLARIS_CONTAINER_AZURE_SKR_AKV_ENDPOINT;
+  const kid = process.env.POLARIS_CONTAINER_AZURE_SKR_KID;
+  const accessToken = process.env.POLARIS_CONTAINER_AZURE_SKR_ACCESS_TOKEN;
+  const maxSKRRequestRetries = Number(process.env.POLARIS_CONTAINER_AZURE_SKR_MAX_REQUEST_RETRIES ?? 5);
+  const skrRetryInterval = Number(process.env.POLARIS_CONTAINER_AZURE_SKR_RETRY_INTERVAL ?? 60000);
+  const advancedTEECheckEnabled = process.env.POLARIS_CONTAINER_ADVANCED_TEE_CHECK_ENABLED === "true";
+
+  if (!maaEndpoint) throw new Error("MAA_ENDPOINT environment variable not defined");
+  if (!akvEndpoint) throw new Error("AKV_ENDPOINT environment variable not defined");
+  if (!kid) throw new Error("KID environment variable not defined");
+  if (advancedTEECheckEnabled && !keyReleaseEndpoint) {
+    throw new Error("KEY_RELEASE_ENDPOINT environment variable not defined for advanced TEE check");
+  }
+
+  return {
+    maa_endpoint: maaEndpoint,
+    akv_endpoint: akvEndpoint,
+    kid,
+    access_token: accessToken,
+    maxSKRRequestRetries,
+    skrRetryInterval,
+    keyReleaseEndpoint,
+    advancedTEECheckEnabled,
+  };
+};
+
+export const convertToBase64EncodedString = (base64UrlEncodedData: string) => {
+  let base64EncodedString = base64UrlEncodedData.replace(/-/g, "+").replace(/_/g, "/");
+
+  while (base64EncodedString.length % 4 !== 0) {
+    base64EncodedString += "=";
+  }
+  return base64EncodedString;
+};
+
+export const convertToPemFormat = (base64UrlEncodedData: string) => {
+  const base64EncodedString = convertToBase64EncodedString(base64UrlEncodedData);
+
+  const formattedKey = base64EncodedString.match(/.{1,64}/g)?.join("\n");
+  return `-----BEGIN PRIVATE KEY-----\n${formattedKey}\n-----END PRIVATE KEY-----`;
+};
+
+export const generateKeys = async (command: string): Promise<{ privateKey: string; publicKey: any }> => {
+  const execAsync = promisify(exec);
+
+  const { stdout } = await execAsync(command);
+
+  const cleanOutput = stdout.replace(/[\s\u0000-\u001F\u007F-\u009F\uFEFF]+$/g, "");
+
+  const privateKey = convertToPemFormat(cleanOutput);
+
+  const privateKeyObject: KeyObject = createPrivateKey(privateKey);
+  const publicKeyObject: KeyObject = createPublicKey(privateKeyObject);
+
+  const publicKey = publicKeyObject.export({ type: "spki", format: "pem" });
+
+  return { privateKey, publicKey };
+};

src/key/handlers/GoogleKeyHandler/index.ts

@@ -0,0 +1,106 @@
+import { promises as fs } from "fs";
+
+import { KeyManagementServiceClient } from "@google-cloud/kms";
+import axios from "axios";
+
+import { getConfigFromEnv } from "./utils";
+
+import type { GoogleKeyConfig } from "./types";
+import type { KeyHandler } from "@fr0ntier-x/polaris-sdk";
+
+export class GoogleKeyHandler implements KeyHandler {
+  private keyWrappingKeyName!: string;
+
+  private keyWrappingKeyKMSClient!: KeyManagementServiceClient;
+
+  private config: GoogleKeyConfig;
+
+  constructor() {
+    this.config = getConfigFromEnv();
+  }
+
+  private async initializeKMSClient() {
+    let credentials: any;
+
+    if (!this.config.advancedTEECheckEnabled) {
+      credentials = {
+        type: "external_account",
+        audience: this.config.federatedCredentialsAudience,
+        service_account_impersonation_url: this.config.federatedCredentialsServiceAccount,
+        subject_token_type: "urn:ietf:params:oauth:token-type:jwt",
+        token_url: "https://sts.googleapis.com/v1/token",
+        credential_source: {
+          file: "/run/container_launcher/attestation_verifier_claims_token",
+        },
+      };
+    } else {
+      await this.verifyAttestation();
+      credentials = this.config.isDevMode
+        ? {
+            client_email: this.config.polarisServiceAccountClientEmail,
+            private_key: this.config.polarisServiceAccountPrivateKey!.replace(/\\n/g, "\n"),
+          }
+        : undefined;
+    }
+
+    this.keyWrappingKeyKMSClient = new KeyManagementServiceClient({ credentials });
+    this.keyWrappingKeyName = this.keyWrappingKeyKMSClient.cryptoKeyVersionPath(
+      this.config.projectId,
+      this.config.location,
+      this.config.keyRingId,
+      this.config.keyId,
+      "1"
+    );
+  }
+
+  private async verifyAttestation() {
+    try {
+      const attestation = await fs.readFile("/tmp/attestation_report.textproto", "utf-8");
+
+      await axios.post(
+        `${this.config.polarisPolicyManagerUrl}/verify`,
+        {
+          attestation_report: attestation,
+          nonce: this.config.nonce,
+          service_account_email: this.config.polarisServiceAccountClientEmail,
+        },
+        {
+          headers: {
+            "Content-Type": "application/json",
+          },
+        }
+      );
+    } catch (error: any) {
+      console.log(error, "Error during attestation verification");
+    }
+  }
+
+  async unwrapKey(wrappedKey: Buffer): Promise<Buffer> {
+    if (!this.keyWrappingKeyKMSClient) {
+      await this.initializeKMSClient();
+    }
+
+    const [decryptResponse] = await this.keyWrappingKeyKMSClient.asymmetricDecrypt({
+      name: this.keyWrappingKeyName,
+      ciphertext: wrappedKey,
+    });
+
+    const decryptedKey = decryptResponse?.plaintext;
+    if (!decryptedKey) throw new Error("Decryption error");
+
+    return Buffer.from(decryptedKey);
+  }
+
+  async getPublicKey(): Promise<string> {
+    if (!this.keyWrappingKeyKMSClient) {
+      await this.initializeKMSClient();
+    }
+    const [publicKey] = await this.keyWrappingKeyKMSClient.getPublicKey({
+      name: this.keyWrappingKeyName,
+    });
+
+    if (!publicKey?.pem) throw new Error("Public key not found");
+
+    return publicKey.pem;
+  }
+}