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

import type { StorageProvider, IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState, RetiredIdentity } 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);
  }
}