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

import type { StorageProvider, IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState, RetiredIdentity, PersistedStreamState } from '@shade/core';
import { constantTimeEqual } 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>();
  private retiredIdentities: RetiredIdentity[] = [];

  // ─── 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 constantTimeEqual(stored, identityKey);
  }

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

  // ─── Identity History ─────────────────────────────────────

  async addRetiredIdentity(identity: RetiredIdentity): Promise<void> {
    this.retiredIdentities.push(identity);
  }

  async getRetiredIdentities(): Promise<RetiredIdentity[]> {
    return [...this.retiredIdentities];
  }

  async pruneRetiredIdentities(olderThan: number): Promise<void> {
    this.retiredIdentities = this.retiredIdentities.filter((r) => r.retiredAt >= olderThan);
  }

  // ─── Stream-transfer resume state (v0.2.0) ────────────────

  private streamStates = new Map<string, PersistedStreamState>();

  async saveStreamState(state: PersistedStreamState): Promise<void> {
    this.streamStates.set(state.streamId, { ...state });
  }

  async getStreamState(streamId: string): Promise<PersistedStreamState | null> {
    const v = this.streamStates.get(streamId);
    return v ? { ...v } : null;
  }

  async removeStreamState(streamId: string): Promise<void> {
    this.streamStates.delete(streamId);
  }

  async listActiveStreamStates(direction?: 'send' | 'receive'): Promise<PersistedStreamState[]> {
    const out: PersistedStreamState[] = [];
    for (const s of this.streamStates.values()) {
      if (s.status !== 'active' && s.status !== 'paused') continue;
      if (direction !== undefined && s.direction !== direction) continue;
      out.push({ ...s });
    }
    out.sort((a, b) => b.updatedAt - a.updatedAt);
    return out;
  }

  async pruneStreamStates(olderThan: number): Promise<void> {
    for (const [id, s] of this.streamStates) {
      if ((s.status === 'finished' || s.status === 'aborted') && s.updatedAt < olderThan) {
        this.streamStates.delete(id);
      }
    }
  }
}