Shade/packages/shade-crypto-web/src/memory-storage.ts

import type { StorageProvider, IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState } from '@shade/core';

/**
 * In-memory StorageProvider for testing and embedded use.
 * All data is lost when the instance is garbage collected.
 */
export class MemoryStorage implements StorageProvider {
  private identityKeyPair: IdentityKeyPair | null = null;
  private registrationId: number = 0;
  private signedPreKeys = new Map<number, SignedPreKey>();
  private oneTimePreKeys = new Map<number, OneTimePreKey>();
  private sessions = new Map<string, SessionState>();
  private trustedIdentities = new Map<string, Uint8Array>();

  // ─── Identity ──────────────────────────────────────────────

  async getIdentityKeyPair(): Promise<IdentityKeyPair | null> {
    return this.identityKeyPair;
  }

  async saveIdentityKeyPair(keyPair: IdentityKeyPair): Promise<void> {
    this.identityKeyPair = keyPair;
  }

  async getLocalRegistrationId(): Promise<number> {
    return this.registrationId;
  }

  async saveLocalRegistrationId(id: number): Promise<void> {
    this.registrationId = id;
  }

  // ─── Signed Pre-Keys ──────────────────────────────────────

  async getSignedPreKey(keyId: number): Promise<SignedPreKey | null> {
    return this.signedPreKeys.get(keyId) ?? null;
  }

  async saveSignedPreKey(key: SignedPreKey): Promise<void> {
    this.signedPreKeys.set(key.keyId, key);
  }

  async removeSignedPreKey(keyId: number): Promise<void> {
    this.signedPreKeys.delete(keyId);
  }

  // ─── One-Time Pre-Keys ────────────────────────────────────

  async getOneTimePreKey(keyId: number): Promise<OneTimePreKey | null> {
    return this.oneTimePreKeys.get(keyId) ?? null;
  }

  async saveOneTimePreKey(key: OneTimePreKey): Promise<void> {
    this.oneTimePreKeys.set(key.keyId, key);
  }

  async removeOneTimePreKey(keyId: number): Promise<void> {
    this.oneTimePreKeys.delete(keyId);
  }

  async getOneTimePreKeyCount(): Promise<number> {
    return this.oneTimePreKeys.size;
  }

  // ─── Sessions ─────────────────────────────────────────────

  async getSession(address: string): Promise<SessionState | null> {
    return this.sessions.get(address) ?? null;
  }

  async saveSession(address: string, state: SessionState): Promise<void> {
    this.sessions.set(address, state);
  }

  async removeSession(address: string): Promise<void> {
    this.sessions.delete(address);
  }

  // ─── Trust ────────────────────────────────────────────────

  async isTrustedIdentity(address: string, identityKey: Uint8Array): Promise<boolean> {
    const stored = this.trustedIdentities.get(address);
    if (!stored) return true; // TOFU: trust on first use
    return arraysEqual(stored, identityKey);
  }

  async saveTrustedIdentity(address: string, identityKey: Uint8Array): Promise<void> {
    this.trustedIdentities.set(address, identityKey);
  }
}

function arraysEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) {
    if (a[i] !== b[i]) return false;
  }
  return true;
}