sal-modular/extension/popup/wasm.js

// WebAssembly API functions for accessing WASM operations directly
// and synchronizing state with background service worker

// Get session state from the background service worker
export function getStatus() {
  return new Promise((resolve) => {
    chrome.runtime.sendMessage({ action: 'get_session' }, (response) => {
      resolve(response);
    });
  });
}

// Debug function to examine vault state using Rhai scripts
export async function debugVaultState(wasmModule) {
  if (!wasmModule) {
    throw new Error('WASM module not loaded');
  }
  
  try {
    console.log('🔍 Debugging vault state...');
    
    // First check if we have a valid session
    const sessionCheck = `
      let has_session = vault::has_active_session();
      let keyspace = "";
      if has_session {
          keyspace = vault::get_current_keyspace();
      }
      
      // Return info about the session
      {
          "has_session": has_session,
          "keyspace": keyspace
      }
    `;
    
    console.log('Checking session status...');
    const sessionStatus = await wasmModule.run_rhai(sessionCheck);
    console.log('Session status:', sessionStatus);
    
    // Only try to get keypairs if we have an active session
    if (sessionStatus && sessionStatus.has_session) {
      // Get information about all keypairs
      const keypairsScript = `
        // Get all keypairs for the current keyspace
        let keypairs = vault::list_keypairs();
        
        // Add more diagnostic information
        let diagnostic = {
          "keypair_count": keypairs.len(),
          "keyspace": vault::get_current_keyspace(),
          "keypairs": keypairs
        };
        
        diagnostic
      `;
      
      console.log('Fetching keypair details...');
      const keypairDiagnostic = await wasmModule.run_rhai(keypairsScript);
      console.log('Keypair diagnostic:', keypairDiagnostic);
      
      return keypairDiagnostic;
    } else {
      console.log('No active session, cannot fetch keypairs');
      return { error: 'No active session' };
    }
  } catch (error) {
    console.error('Error in debug function:', error);
    return { error: error.toString() };
  }
}

// Fetch all keypairs from the WebAssembly vault
export async function getKeypairsFromVault(wasmModule) {
  if (!wasmModule) {
    throw new Error('WASM module not loaded');
  }
  
  try {
    // First run diagnostics for debugging
    await debugVaultState(wasmModule);
    
    console.log('Calling list_keypairs WebAssembly binding...');
    
    // Use our new direct WebAssembly binding instead of Rhai script
    const keypairList = await wasmModule.list_keypairs();
    console.log('Retrieved keypairs from vault:', keypairList);
    
    // Transform the keypairs into the expected format
    // The WebAssembly binding returns an array of objects with id, type, and metadata
    const formattedKeypairs = Array.isArray(keypairList) ? keypairList.map(kp => {
      // Parse metadata if it's a string
      let metadata = {};
      if (kp.metadata) {
        try {
          if (typeof kp.metadata === 'string') {
            metadata = JSON.parse(kp.metadata);
          } else {
            metadata = kp.metadata;
          }
        } catch (e) {
          console.warn('Failed to parse keypair metadata:', e);
        }
      }
      
      return {
        id: kp.id,
        label: metadata.label || `${kp.type}-Key-${kp.id.substring(0, 4)}`
      };
    }) : [];
    
    console.log('Formatted keypairs:', formattedKeypairs);
    
    // Update the keypairs in the background service worker
    return new Promise((resolve) => {
      chrome.runtime.sendMessage({
        action: 'update_session',
        type: 'keypairs_loaded',
        data: formattedKeypairs
      }, (response) => {
        if (response && response.success) {
          console.log('Successfully updated keypairs in background');
          resolve(formattedKeypairs);
        } else {
          console.error('Failed to update keypairs in background:', response?.error);
          resolve([]);
        }
      });
    });
  } catch (error) {
    console.error('Error fetching keypairs from vault:', error);
    return [];
  }
}

// Initialize session with the WASM module
export function initSession(wasmModule, keyspace, password) {
  return new Promise(async (resolve, reject) => {
    if (!wasmModule) {
      reject('WASM module not loaded');
      return;
    }
    
    try {
      // Call the WASM init_session function
      console.log(`Initializing session for keyspace: ${keyspace}`);
      await wasmModule.init_session(keyspace, password);
      
      // Update the session state in the background service worker
      chrome.runtime.sendMessage({ 
        action: 'update_session', 
        type: 'keyspace', 
        data: keyspace 
      }, async (response) => {
        if (response && response.success) {
          try {
            // After successful session initialization, fetch keypairs from the vault
            console.log('Session initialized, fetching keypairs from vault...');
            const keypairs = await getKeypairsFromVault(wasmModule);
            console.log('Keypairs loaded:', keypairs);
            resolve(keypairs);
          } catch (fetchError) {
            console.error('Error fetching keypairs:', fetchError);
            // Even if fetching keypairs fails, the session is initialized
            resolve([]);
          }
        } else {
          reject(response && response.error ? response.error : 'Failed to update session state');
        }
      });
    } catch (error) {
      console.error('Session initialization error:', error);
      reject(error.message || 'Failed to initialize session');
    }
  });
}

// Lock the session using the WASM module
export function lockSession(wasmModule) {
  return new Promise(async (resolve, reject) => {
    if (!wasmModule) {
      reject('WASM module not loaded');
      return;
    }
    
    try {
      // Call the WASM lock_session function
      wasmModule.lock_session();
      
      // Update the session state in the background service worker
      chrome.runtime.sendMessage({ 
        action: 'update_session', 
        type: 'session_locked'
      }, (response) => {
        if (response && response.success) {
          resolve();
        } else {
          reject(response && response.error ? response.error : 'Failed to update session state');
        }
      });
    } catch (error) {
      reject(error.message || 'Failed to lock session');
    }
  });
}

// Add a keypair using the WASM module
export function addKeypair(wasmModule, keyType = 'Secp256k1', label = null) {
  return new Promise(async (resolve, reject) => {
    if (!wasmModule) {
      reject('WASM module not loaded');
      return;
    }
    
    try {
      // Create a default label if none provided
      const keyLabel = label || `${keyType}-Key-${Date.now().toString(16).slice(-4)}`;
      
      // Create metadata JSON for the keypair
      const metadata = JSON.stringify({
        label: keyLabel,
        created: new Date().toISOString(),
        type: keyType
      });
      
      console.log(`Adding new keypair of type ${keyType} with label ${keyLabel}`);
      
      // Call the WASM add_keypair function with metadata
      const keyId = await wasmModule.add_keypair(keyType, metadata);
      console.log(`Keypair created with ID: ${keyId}`);
      
      // Create keypair object with ID and label
      const newKeypair = { 
        id: keyId, 
        label: keyLabel 
      };
      
      // Update the session state in the background service worker
      chrome.runtime.sendMessage({ 
        action: 'update_session', 
        type: 'keypair_added', 
        data: newKeypair 
      }, (response) => {
        if (response && response.success) {
          // After adding a keypair, refresh the whole list from the vault
          getKeypairsFromVault(wasmModule)
            .then(() => {
              console.log('Keypair list refreshed from vault');
              resolve(keyId);
            })
            .catch(refreshError => {
              console.warn('Error refreshing keypair list:', refreshError);
              // Still resolve with the key ID since the key was created
              resolve(keyId);
            });
        } else {
          reject(response && response.error ? response.error : 'Failed to update session state');
        }
      });
    } catch (error) {
      console.error('Error adding keypair:', error);
      reject(error.message || 'Failed to add keypair');
    }
  });
}

// Select a keypair using the WASM module
export function selectKeypair(wasmModule, keyId) {
  return new Promise(async (resolve, reject) => {
    if (!wasmModule) {
      reject('WASM module not loaded');
      return;
    }
    
    try {
      // Call the WASM select_keypair function
      await wasmModule.select_keypair(keyId);
      
      // Update the session state in the background service worker
      chrome.runtime.sendMessage({ 
        action: 'update_session', 
        type: 'keypair_selected', 
        data: keyId 
      }, (response) => {
        if (response && response.success) {
          resolve();
        } else {
          reject(response && response.error ? response.error : 'Failed to update session state');
        }
      });
    } catch (error) {
      reject(error.message || 'Failed to select keypair');
    }
  });
}

// Sign a message using the WASM module
export function sign(wasmModule, message) {
  return new Promise(async (resolve, reject) => {
    if (!wasmModule) {
      reject('WASM module not loaded');
      return;
    }
    
    try {
      // Convert message to Uint8Array for WASM
      const encoder = new TextEncoder();
      const messageBytes = encoder.encode(message);
      
      // Call the WASM sign function
      const signature = await wasmModule.sign(messageBytes);
      resolve(signature);
    } catch (error) {
      reject(error.message || 'Failed to sign message');
    }
  });
}