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

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

  // ─── Peer verifications (V3.3) ────────────────────────────

  private peerVerifications = new Map<string, PeerVerification>();
  private peerIdentityVersions = new Map<string, number>();

  async savePeerVerification(v: PeerVerification): Promise<void> {
    this.peerVerifications.set(v.peerAddress, { ...v });
  }

  async getPeerVerification(address: string): Promise<PeerVerification | null> {
    const v = this.peerVerifications.get(address);
    return v ? { ...v } : null;
  }

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

  async getPeerIdentityVersion(address: string): Promise<number> {
    return this.peerIdentityVersions.get(address) ?? 1;
  }

  async bumpPeerIdentityVersion(address: string): Promise<number> {
    const next = (this.peerIdentityVersions.get(address) ?? 1) + 1;
    this.peerIdentityVersions.set(address, next);
    return next;
  }

  // ─── 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);
      }
    }
  }

  // ─── Broadcast channels (V4.6) ────────────────────────────

  private broadcastChannels = new Map<string, BroadcastChannelRecord>();
  private broadcastMembers = new Map<string, Map<string, BroadcastMemberRecord>>();

  async saveBroadcastChannel(channel: BroadcastChannelRecord): Promise<void> {
    this.broadcastChannels.set(channel.channelId, cloneChannel(channel));
  }

  async getBroadcastChannel(channelId: string): Promise<BroadcastChannelRecord | null> {
    const v = this.broadcastChannels.get(channelId);
    return v ? cloneChannel(v) : null;
  }

  async listBroadcastChannels(): Promise<BroadcastChannelRecord[]> {
    return [...this.broadcastChannels.values()].map(cloneChannel);
  }

  async removeBroadcastChannel(channelId: string): Promise<void> {
    this.broadcastChannels.delete(channelId);
    this.broadcastMembers.delete(channelId);
  }

  async saveBroadcastMember(member: BroadcastMemberRecord): Promise<void> {
    let inner = this.broadcastMembers.get(member.channelId);
    if (!inner) {
      inner = new Map();
      this.broadcastMembers.set(member.channelId, inner);
    }
    inner.set(member.peerAddress, { ...member });
  }

  async getBroadcastMembers(channelId: string): Promise<BroadcastMemberRecord[]> {
    const inner = this.broadcastMembers.get(channelId);
    return inner ? [...inner.values()].map((m) => ({ ...m })) : [];
  }

  async removeBroadcastMember(channelId: string, peerAddress: string): Promise<void> {
    this.broadcastMembers.get(channelId)?.delete(peerAddress);
  }
}

function cloneChannel(c: BroadcastChannelRecord): BroadcastChannelRecord {
  const out: BroadcastChannelRecord = {
    channelId: c.channelId,
    ownerRole: c.ownerRole,
    ownerAddress: c.ownerAddress,
    generation: c.generation,
    chainKey: new Uint8Array(c.chainKey),
    iteration: c.iteration,
    signingPublicKey: new Uint8Array(c.signingPublicKey),
    createdAt: c.createdAt,
    updatedAt: c.updatedAt,
  };
  if (c.label !== undefined) out.label = c.label;
  if (c.signingPrivateKey !== undefined) out.signingPrivateKey = new Uint8Array(c.signingPrivateKey);
  return out;
}