sal-modular/crypto_vault_extension/background.js

let vault = null;
let isInitialized = false;
let currentSession = null;
let keepAliveInterval = null;
let sessionTimeoutDuration = 15; // Default 15 seconds
let sessionTimeoutId = null; // Background timer
let popupPort = null; // Track popup connection

// Utility function to convert Uint8Array to hex
function toHex(uint8Array) {
  return Array.from(uint8Array)
    .map(b => b.toString(16).padStart(2, '0'))
    .join('');
}
// Background session timeout management
async function loadTimeoutSetting() {
  const result = await chrome.storage.local.get(['sessionTimeout']);
  sessionTimeoutDuration = result.sessionTimeout || 15;
}

function startSessionTimeout() {
  clearSessionTimeout();

  if (currentSession && sessionTimeoutDuration > 0) {
    sessionTimeoutId = setTimeout(async () => {
      if (vault && currentSession) {
        // Lock the session
        vault.lock_session();
        await sessionManager.clear();
        // Notify popup if it's open
        if (popupPort) {
          popupPort.postMessage({
            type: 'sessionTimeout',
            message: 'Session timed out due to inactivity'
          });
        }
      }
    }, sessionTimeoutDuration * 1000);
  }
}

function clearSessionTimeout() {
  if (sessionTimeoutId) {
    clearTimeout(sessionTimeoutId);
    sessionTimeoutId = null;
  }
}

function resetSessionTimeout() {
  if (currentSession) {
    startSessionTimeout();
  }
}

// Session persistence functions
async function saveSession(keyspace) {
  currentSession = { keyspace, timestamp: Date.now() };

  // Save to both session and local storage for better persistence
  await chrome.storage.session.set({ cryptoVaultSession: currentSession });
  await chrome.storage.local.set({ cryptoVaultSessionBackup: currentSession });
}

async function loadSession() {
  // Try session storage first
  let result = await chrome.storage.session.get(['cryptoVaultSession']);
  if (result.cryptoVaultSession) {
    currentSession = result.cryptoVaultSession;
    return currentSession;
  }

  // Fallback to local storage
  result = await chrome.storage.local.get(['cryptoVaultSessionBackup']);
  if (result.cryptoVaultSessionBackup) {
    currentSession = result.cryptoVaultSessionBackup;
    // Restore to session storage
    await chrome.storage.session.set({ cryptoVaultSession: currentSession });
    return currentSession;
  }
  return null;
}

async function clearSession() {
  currentSession = null;
  await chrome.storage.session.remove(['cryptoVaultSession']);
  await chrome.storage.local.remove(['cryptoVaultSessionBackup']);
}

// Keep service worker alive
function startKeepAlive() {
  if (keepAliveInterval) {
    clearInterval(keepAliveInterval);
  }

  keepAliveInterval = setInterval(() => {
    chrome.storage.session.get(['keepAlive']).catch(() => {});
  }, 20000);
}

function stopKeepAlive() {
  if (keepAliveInterval) {
    clearInterval(keepAliveInterval);
    keepAliveInterval = null;
  }
}

// Consolidated session management
const sessionManager = {
  async save(keyspace) {
    await saveSession(keyspace);
    startKeepAlive();
    await loadTimeoutSetting();
    startSessionTimeout();
  },
  async clear() {
    await clearSession();
    stopKeepAlive();
    clearSessionTimeout();
  }
};

async function restoreSession() {
  const session = await loadSession();
  if (session && vault) {
    // Check if the session is still valid by testing if vault is unlocked
    const isUnlocked = vault.is_unlocked();
    if (isUnlocked) {
      // Restart keep-alive for restored session
      startKeepAlive();
      return session;
    } else {
      await sessionManager.clear();
    }
  }
  return null;
}

// Import WASM module functions
import init, * as wasmFunctions from './wasm/wasm_app.js';

// Initialize WASM module
async function initVault() {
  try {
    if (vault && isInitialized) return vault;

    // Initialize with the WASM file
    const wasmUrl = chrome.runtime.getURL('wasm/wasm_app_bg.wasm');
    await init(wasmUrl);

    // Use imported functions directly
    vault = wasmFunctions;
    isInitialized = true;

    // Try to restore previous session
    await restoreSession();

    return vault;
  } catch (error) {
    console.error('Failed to initialize CryptoVault:', error);
    throw error;
  }
}


// Consolidated message handlers
const messageHandlers = {
  createKeyspace: async (request) => {
    await vault.create_keyspace(request.keyspace, request.password);
    return { success: true };
  },

  initSession: async (request) => {
    await vault.init_session(request.keyspace, request.password);
    await sessionManager.save(request.keyspace);
    return { success: true };
  },

  isUnlocked: () => ({ success: true, unlocked: vault.is_unlocked() }),

  addKeypair: async (request) => {
    const result = await vault.add_keypair(request.keyType, request.metadata);
    return { success: true, result };
  },

  listKeypairs: async () => {
    if (!vault.is_unlocked()) {
      return { success: false, error: 'Session is not unlocked' };
    }
    const keypairsRaw = await vault.list_keypairs();
    const keypairs = typeof keypairsRaw === 'string' ? JSON.parse(keypairsRaw) : keypairsRaw;
    return { success: true, keypairs };
  },

  selectKeypair: (request) => {
    vault.select_keypair(request.keyId);
    return { success: true };
  },

  getCurrentKeypairMetadata: () => ({ success: true, metadata: vault.current_keypair_metadata() }),

  getCurrentKeypairPublicKey: () => ({ success: true, publicKey: toHex(vault.current_keypair_public_key()) }),

  sign: async (request) => {
    const signature = await vault.sign(new Uint8Array(request.message));
    return { success: true, signature };
  },

  encrypt: async (request) => {
    if (!vault.is_unlocked()) {
      return { success: false, error: 'Session is not unlocked' };
    }
    const messageBytes = new TextEncoder().encode(request.message);
    const encryptedData = await vault.encrypt_data(messageBytes);
    const encryptedMessage = btoa(String.fromCharCode(...new Uint8Array(encryptedData)));
    return { success: true, encryptedMessage };
  },

  decrypt: async (request) => {
    if (!vault.is_unlocked()) {
      return { success: false, error: 'Session is not unlocked' };
    }
    const encryptedBytes = new Uint8Array(atob(request.encryptedMessage).split('').map(c => c.charCodeAt(0)));
    const decryptedData = await vault.decrypt_data(encryptedBytes);
    const decryptedMessage = new TextDecoder().decode(new Uint8Array(decryptedData));
    return { success: true, decryptedMessage };
  },

  verify: async (request) => {
    const metadata = vault.current_keypair_metadata();
    if (!metadata) {
      return { success: false, error: 'No keypair selected' };
    }
    const isValid = await vault.verify(new Uint8Array(request.message), request.signature);
    return { success: true, isValid };
  },

  lockSession: async () => {
    vault.lock_session();
    await sessionManager.clear();
    return { success: true };
  },

  getStatus: async () => {
    const status = vault ? vault.is_unlocked() : false;
    const session = await loadSession();
    return {
      success: true,
      status,
      session: session ? { keyspace: session.keyspace } : null
    };
  },

  // Timeout management handlers
  resetTimeout: async () => {
    resetSessionTimeout();
    return { success: true };
  },

  updateTimeout: async (request) => {
    sessionTimeoutDuration = request.timeout;
    await chrome.storage.local.set({ sessionTimeout: request.timeout });
    resetSessionTimeout(); // Restart with new duration
    return { success: true };
  }
};

// Handle messages from popup and content scripts
chrome.runtime.onMessage.addListener((request, _sender, sendResponse) => {
  const handleRequest = async () => {
    try {
      if (!vault) {
        await initVault();
      }

      const handler = messageHandlers[request.action];
      if (handler) {
        return await handler(request);
      } else {
        throw new Error('Unknown action: ' + request.action);
      }
    } catch (error) {
      return { success: false, error: error.message };
    }
  };

  handleRequest().then(sendResponse);
  return true; // Keep the message channel open for async response
});

// Initialize vault when extension starts
chrome.runtime.onStartup.addListener(() => {
  initVault();
});

chrome.runtime.onInstalled.addListener(() => {
  initVault();
});

// Handle popup connection for keep-alive and timeout notifications
chrome.runtime.onConnect.addListener((port) => {
  if (port.name === 'popup') {
    // Track popup connection
    popupPort = port;

    // If we have an active session, ensure keep-alive is running
    if (currentSession) {
      startKeepAlive();
    }

    port.onDisconnect.addListener(() => {
      // Popup closed, clear reference and stop keep-alive
      popupPort = null;
      stopKeepAlive();
    });
  }
});