Implement AES-256-GCM encryption, secure initialization, and improve SkyFi tool robustness

lib/actions/skyfi.ts

@@ -7,8 +7,20 @@ import { eq } from 'drizzle-orm';
 import { getCurrentUserIdOnServer } from '@/lib/auth/get-current-user';
 import { SkyfiOAuthProvider } from '@/lib/skyfi/provider';
 import crypto from 'crypto';
+import { redirect } from 'next/navigation';
+import { headers } from 'next/headers';
 
 export async function getRedirectUri(): Promise<string> {
+  try {
+    const headersList = await headers();
+    const host = headersList.get('host');
+    if (host) {
+      const protocol = host.startsWith('localhost') || host.startsWith('127.0.0.1') ? 'http' : 'https';
+      return `${protocol}://${host}/api/skyfi/callback`;
+    }
+  } catch (e) {
+    console.warn('[Skyfi] Failed to get host from headers, falling back to ENV:', e);
+  }
   const baseUrl = process.env.NEXT_PUBLIC_APP_URL || 'http://localhost:3000';
   return `${baseUrl}/api/skyfi/callback`;
 }
@@ -17,13 +29,15 @@ export async function getRedirectUri(): Promise<string> {
  * Ensures the client is registered dynamically with the SkyFi MCP server.
  * Uses an AbortController with a 10s timeout to prevent hanging.
  */
-async function ensureClientRegistered(provider: SkyfiOAuthProvider): Promise<string> {
-  const currentInfo = await provider.clientInformation();
-  if (currentInfo?.client_id) {
-    return currentInfo.client_id;
+async function ensureClientRegistered(provider: SkyfiOAuthProvider, forceRegister: boolean = false): Promise<string> {
+  if (!forceRegister) {
+    const currentInfo = await provider.clientInformation();
+    if (currentInfo?.client_id) {
+      return currentInfo.client_id;
+    }
   }
 
-  console.log('[SkyFiAction] Client not registered. Registering dynamically...');
+  console.log('[SkyFiAction] Client not registered or force-register active. Registering dynamically...');
 
   const controller = new AbortController();
   const timeoutId = setTimeout(() => controller.abort(), 10000);
@@ -68,16 +82,16 @@ async function ensureClientRegistered(provider: SkyfiOAuthProvider): Promise<str
  * Performs DCR, generates PKCE, and returns the authorization URL.
  */
 export async function startSkyfiConnection(): Promise<{ url?: string; error?: string }> {
-  try {
-    const userId = await getCurrentUserIdOnServer();
-    if (!userId) {
-      return { error: 'Unauthorized: User not authenticated.' };
-    }
+  const userId = await getCurrentUserIdOnServer();
+  if (!userId) {
+    redirect('/sign-in');
+  }
 
+  try {
     const redirectUri = await getRedirectUri();
     const provider = new SkyfiOAuthProvider(userId, redirectUri);
 
-    const clientId = await ensureClientRegistered(provider);
+    const clientId = await ensureClientRegistered(provider, true);
 
     const verifier = generateCodeVerifier();
     const challenge = generateCodeChallenge(verifier);

lib/agents/tools/skyfi.tsx

@@ -11,10 +11,21 @@ import { skyfiQuerySchema } from '@/lib/schema/skyfi';
 import { DrawnFeature } from '@/lib/agents/resolution-search';
 import { z } from 'zod';
 import crypto from 'crypto';
+import { headers } from 'next/headers';
 
 export type McpClient = MCPClientClass;
 
-function getRedirectUri(): string {
+async function getRedirectUri(): Promise<string> {
+  try {
+    const headersList = await headers();
+    const host = headersList.get('host');
+    if (host) {
+      const protocol = host.startsWith('localhost') || host.startsWith('127.0.0.1') ? 'http' : 'https';
+      return `${protocol}://${host}/api/skyfi/callback`;
+    }
+  } catch (e) {
+    console.warn('[SkyfiTool] Failed to get host from headers, falling back to ENV:', e);
+  }
   const baseUrl = process.env.NEXT_PUBLIC_APP_URL || 'http://localhost:3000';
   return `${baseUrl}/api/skyfi/callback`;
 }
@@ -29,6 +40,14 @@ function geoJsonToWkt(geometry: any): string | null {
   return null;
 }
 
+function deterministicSerialize(val: any): string {
+  if (val === null || val === undefined) return '';
+  if (typeof val !== 'object') return String(val);
+  const keys = Object.keys(val).sort();
+  const parts = keys.map(k => `${k}:${deterministicSerialize(val[k])}`);
+  return `{${parts.join(',')}}`;
+}
+
 /**
  * Establish connection to the SkyFi MCP server with the user's stored OAuth token and support AbortSignal.
  */
@@ -118,8 +137,8 @@ export const skyfiTool = ({
   Always ask the user for permission before placing a billable order! Ensure you call 'validate_order' first to display the cost.`,
   parameters: skyfiQuerySchema,
   execute: async (params: z.infer<typeof skyfiQuerySchema>) => {
-    const { queryType, location, aoi, params: extraParams, maxResults } = params;
-    console.log('[SkyfiTool] Execute called with:', params);
+    const { queryType, location, aoi, orderId, params: extraParams, maxResults } = params;
+    console.log('[SkyfiTool] Execute called for queryType:', queryType, { maxResults, confirm: params.confirm });
 
     const uiFeedbackStream = createStreamableValue<string>();
     uiStream.append(<BotMessage content={uiFeedbackStream.value} />);
@@ -135,7 +154,7 @@ export const skyfiTool = ({
     }
 
     // Resolve OAuth tokens
-    const redirectUri = getRedirectUri();
+    const redirectUri = await getRedirectUri();
     const provider = new SkyfiOAuthProvider(userId, redirectUri);
     const tokens = await provider.tokens();
 
@@ -190,6 +209,7 @@ export const skyfiTool = ({
           clearTimeout(timeoutId);
           const invalidAoiMessage = "Your drawn area is not a valid closed Polygon. SkyFi requires a closed Polygon to define your Area of Interest. Please use the map drawing tools to draw a closed Polygon.";
           uiFeedbackStream.update(invalidAoiMessage);
+          await closeClient(mcpClient);
           uiFeedbackStream.done();
           uiStream.update(<BotMessage content={uiFeedbackStream.value} />);
           return { type: 'SKYFI_QUERY', success: false, error: 'Invalid AOI geometry', message: invalidAoiMessage };
@@ -201,8 +221,10 @@ export const skyfiTool = ({
       }
 
       // Stable idempotency key derivation to prevent duplicate order placements
+      const serializedExtraParams = deterministicSerialize(extraParams);
+      const hashInput = `${userId}_${resolvedAoiWkt || ''}_${queryType}_${orderId || ''}_${serializedExtraParams}`;
       stableIdempotencyKey = crypto.createHash('sha256')
-        .update(`${userId}_${resolvedAoiWkt || ''}_${queryType}`)
+        .update(hashInput)
         .digest('hex');
 
       switch (queryType) {
@@ -226,6 +248,7 @@ export const skyfiTool = ({
           mcpToolName = 'skyfi_validate_archive_order';
           mcpArgs = {
             ...extraParams,
+            ...(orderId ? { orderId } : {}),
             aoi: resolvedAoiWkt || extraParams?.aoi
           };
           break;
@@ -238,15 +261,17 @@ export const skyfiTool = ({
             mcpToolName = 'skyfi_place_archive_order';
           }
           mcpArgs = {
-            idempotencyKey: stableIdempotencyKey,
             ...extraParams,
+            ...(orderId ? { orderId } : {}),
+            idempotencyKey: stableIdempotencyKey,
             aoi: resolvedAoiWkt || extraParams?.aoi
           };
           break;
         case 'list_orders':
           mcpToolName = 'skyfi_list_orders';
           mcpArgs = {
-            ...extraParams
+            ...extraParams,
+            ...(orderId ? { orderId } : {})
           };
           break;
         default:
@@ -261,7 +286,7 @@ export const skyfiTool = ({
         const geocodeResult = await mcpClient.callTool({
           name: 'skyfi_geocode_aoi',
           arguments: { place: location }
-        });
+        }, undefined, { signal: controller.signal });
 
         const geocodeText = (geocodeResult as any)?.content?.[0]?.text || '';
         // Extract polygon WKT
@@ -274,9 +299,12 @@ export const skyfiTool = ({
         }
       }
 
-      console.log('[SkyfiTool] Calling tool:', mcpToolName, 'with args:', mcpArgs);
+      console.log('[SkyfiTool] Calling tool:', mcpToolName);
 
-      const mcpResult = await mcpClient.callTool({ name: mcpToolName, arguments: mcpArgs });
+      const mcpResult = await mcpClient.callTool({
+        name: mcpToolName,
+        arguments: mcpArgs
+      }, undefined, { signal: controller.signal });
 
       clearTimeout(timeoutId);
 

lib/utils/encryption.ts

@@ -1,30 +1,54 @@
 import crypto from 'crypto';
 
-const ALGORITHM = 'aes-256-cbc';
-const ENCRYPTION_KEY = process.env.ENCRYPTION_KEY || 'default_32_byte_secret_key_placeholder_for_local_dev'; // Must be 32 bytes
+const ALGORITHM_GCM = 'aes-256-gcm';
+const ALGORITHM_CBC = 'aes-256-cbc';
+
+if (!process.env.ENCRYPTION_KEY) {
+  throw new Error('ENCRYPTION_KEY environment variable is not set');
+}
+const ENCRYPTION_KEY = process.env.ENCRYPTION_KEY;
 
 export function encrypt(text: string | null | undefined): string | null {
   if (!text) return null;
-  const iv = crypto.randomBytes(16);
+  const iv = crypto.randomBytes(12); // Standard IV size for GCM is 12 bytes
   // Ensure the key is exactly 32 bytes
   const key = crypto.createHash('sha256').update(ENCRYPTION_KEY).digest();
-  const cipher = crypto.createCipheriv(ALGORITHM, key, iv);
+  const cipher = crypto.createCipheriv(ALGORITHM_GCM, key, iv);
   let encrypted = cipher.update(text, 'utf8', 'hex');
   encrypted += cipher.final('hex');
-  return `${iv.toString('hex')}:${encrypted}`;
+  const tag = cipher.getAuthTag().toString('hex');
+  return `${iv.toString('hex')}:${encrypted}:${tag}`;
 }
 
 export function decrypt(encryptedText: string | null | undefined): string | null {
   if (!encryptedText) return null;
   try {
-    const [ivHex, encrypted] = encryptedText.split(':');
-    if (!ivHex || !encrypted) return null;
-    const iv = Buffer.from(ivHex, 'hex');
+    const parts = encryptedText.split(':');
     const key = crypto.createHash('sha256').update(ENCRYPTION_KEY).digest();
-    const decipher = crypto.createDecipheriv(ALGORITHM, key, iv);
-    let decrypted = decipher.update(encrypted, 'hex', 'utf8');
-    decrypted += decipher.final('utf8');
-    return decrypted;
+
+    if (parts.length === 3) {
+      // New AES-256-GCM format
+      const [ivHex, encrypted, tagHex] = parts;
+      if (!ivHex || !encrypted || !tagHex) return null;
+      const iv = Buffer.from(ivHex, 'hex');
+      const tag = Buffer.from(tagHex, 'hex');
+      const decipher = crypto.createDecipheriv(ALGORITHM_GCM, key, iv);
+      decipher.setAuthTag(tag);
+      let decrypted = decipher.update(encrypted, 'hex', 'utf8');
+      decrypted += decipher.final('utf8');
+      return decrypted;
+    } else if (parts.length === 2) {
+      // Legacy AES-256-CBC format
+      const [ivHex, encrypted] = parts;
+      if (!ivHex || !encrypted) return null;
+      const iv = Buffer.from(ivHex, 'hex');
+      const decipher = crypto.createDecipheriv(ALGORITHM_CBC, key, iv);
+      let decrypted = decipher.update(encrypted, 'hex', 'utf8');
+      decrypted += decipher.final('utf8');
+      return decrypted;
+    }
+
+    return null;
   } catch (error) {
     console.error('[Encryption] Decryption failed:', error);
     return null;

tests-unit/encryption.test.ts

@@ -0,0 +1,74 @@
+import crypto from 'crypto';
+import { encrypt, decrypt } from '../lib/utils/encryption';
+
+console.log("Starting encryption tests...");
+
+const ENCRYPTION_KEY = process.env.ENCRYPTION_KEY || 'my-super-secret-key-32-chars-long-or-more';
+
+// Test basic GCM encrypt/decrypt
+const originalText = "Hello, GCM authenticated encryption!";
+const encrypted = encrypt(originalText);
+
+if (!encrypted) {
+  throw new Error("Encryption failed: returned null");
+}
+
+console.log("Encrypted GCM Text:", encrypted);
+
+const parts = encrypted.split(':');
+if (parts.length !== 3) {
+  throw new Error(`Encrypted text does not have 3 parts separated by colons. Got: ${parts.length}`);
+}
+
+const decrypted = decrypt(encrypted);
+if (decrypted !== originalText) {
+  throw new Error(`GCM Decryption failed! Expected: "${originalText}", got: "${decrypted}"`);
+}
+console.log("GCM Decryption verified successfully!");
+
+// Test legacy CBC decryption backwards compatibility
+const legacyText = "Hello, legacy CBC encryption!";
+const ivCbc = crypto.randomBytes(16);
+const key = crypto.createHash('sha256').update(ENCRYPTION_KEY).digest();
+const cipherCbc = crypto.createCipheriv('aes-256-cbc', key, ivCbc);
+let encryptedCbc = cipherCbc.update(legacyText, 'utf8', 'hex');
+encryptedCbc += cipherCbc.final('hex');
+const legacyPayload = `${ivCbc.toString('hex')}:${encryptedCbc}`;
+
+console.log("Legacy CBC payload:", legacyPayload);
+const decryptedLegacy = decrypt(legacyPayload);
+if (decryptedLegacy !== legacyText) {
+  throw new Error(`Legacy CBC Decryption failed! Expected: "${legacyText}", got: "${decryptedLegacy}"`);
+}
+console.log("Legacy CBC Decryption backwards compatibility verified successfully!");
+
+// Test integrity protection (GCM authentication tag verification)
+// Modify one byte of the ciphertext
+const [iv, ciphertext, tag] = parts;
+const modifiedCiphertext = ciphertext.substring(0, ciphertext.length - 2) + (ciphertext.endsWith('0') ? '1' : '0');
+const modifiedEncrypted = `${iv}:${modifiedCiphertext}:${tag}`;
+
+// Decrypting tampered or invalid hex string might log an error and return null.
+// We disable console.error temporarily to keep the test logs clean.
+const originalConsoleError = console.error;
+console.error = () => {};
+try {
+  const decryptedModified = decrypt(modifiedEncrypted);
+  if (decryptedModified !== null) {
+    throw new Error("Decryption of modified ciphertext did not return null! GCM authentication tag failed to catch tamper.");
+  }
+} finally {
+  console.error = originalConsoleError;
+}
+console.log("GCM Authentication tag verification passed!");
+
+// Test null/undefined handling
+if (encrypt(null) !== null || encrypt(undefined) !== null) {
+  throw new Error("encrypt(null/undefined) did not return null");
+}
+if (decrypt(null) !== null || decrypt(undefined) !== null) {
+  throw new Error("decrypt(null/undefined) did not return null");
+}
+console.log("Null/undefined handling passed!");
+
+console.log("All encryption tests passed!");