PotentialStorage/test/kv-storage/kv-storage-encrypted.test.ts

import { afterEach, describe, expect, it } from 'vitest'
import { KVStorageComposite } from '../../src/kv-storage/kv-storage-composite.js'
import { KVStorageEncrypted } from '../../src/kv-storage/kv-storage-encrypted.js'
import { KVStorageInMemory } from '../../src/kv-storage/kv-storage-in-memory.js'
import { collect } from '../helpers.js'
import { runKVStorageConformance } from './kv-storage.conformance.js'
import { TestNetwork } from './sync-helpers.js'

const key = await KVStorageEncrypted.generateKey()

runKVStorageConformance('KVStorageEncrypted (in memory)', () => new KVStorageEncrypted(new KVStorageInMemory(), { key }))

let network: TestNetwork | undefined
afterEach(async () => {
  await network?.close()
  network = undefined
})

describe('KVStorageEncrypted', () => {
  it('stores only ciphertext, with readable keys', async () => {
    const inner = new KVStorageInMemory()
    await new KVStorageEncrypted(inner, { key }).child('users').set('1', { name: 'Ada', n: 1n })
    const [[storedKey, stored]] = (await collect(inner.scan({ values: true }))) as [[string, Uint8Array]]
    expect(storedKey).toBe('users/1')
    expect(stored).toBeInstanceOf(Uint8Array)
    expect(new TextDecoder().decode(stored)).not.toContain('Ada')
  })

  it('uses a fresh IV for every write', async () => {
    const inner = new KVStorageInMemory()
    const db = new KVStorageEncrypted(inner, { key })
    await db.set('a', 'same')
    await db.set('b', 'same')
    const [a, b] = (await inner.readMany(['a', 'b'])) as Uint8Array[]
    expect(a).not.toEqual(b)
  })

  it('rejects a wrong key, a value moved to another key, tampering and plaintext', async () => {
    const inner = new KVStorageInMemory()
    await new KVStorageEncrypted(inner, { key }).set('a', 'secret')
    const other = new KVStorageEncrypted(inner, { key: await KVStorageEncrypted.generateKey() })
    await expect(other.get('a')).rejects.toThrow(/Cannot decrypt/)

    const db = new KVStorageEncrypted(inner, { key })
    const [stored] = (await inner.readMany(['a'])) as [Uint8Array]
    await inner.writeMany([{ key: 'b', value: stored }])
    await expect(db.get('b')).rejects.toThrow(/Cannot decrypt/)

    const tampered = stored.slice()
    tampered[tampered.length - 1]! ^= 1
    await inner.writeMany([{ key: 'c', value: tampered }])
    await expect(db.get('c')).rejects.toThrow(/Cannot decrypt/)

    await inner.writeMany([{ key: 'd', value: 'plain' }])
    await expect(db.get('d')).rejects.toThrow(/not encrypted/)
  })

  it('round-trips keys through export/import, and derives the same key from the same password and salt', async () => {
    const inner = new KVStorageInMemory()
    await new KVStorageEncrypted(inner, { key }).set('a', 1)
    const imported = await KVStorageEncrypted.importKey(await KVStorageEncrypted.exportKey(key))
    expect(await new KVStorageEncrypted(inner, { key: imported }).get('a')).toBe(1)

    const salt = crypto.getRandomValues(new Uint8Array(16))
    await new KVStorageEncrypted(inner, { key: await KVStorageEncrypted.deriveKey('hunter2', salt, 1000) }).set('p', 'x')
    expect(await new KVStorageEncrypted(inner, { key: await KVStorageEncrypted.deriveKey('hunter2', salt, 1000) }).get('p')).toBe('x')
    await expect(new KVStorageEncrypted(inner, { key: await KVStorageEncrypted.deriveKey('wrong', salt, 1000) }).get('p')).rejects.toThrow(/Cannot decrypt/)
  })

  it('decrypts pushed changes and passes subscriptions through', async () => {
    network = new TestNetwork()
    const watcher = new KVStorageEncrypted(network.client(), { key })
    await watcher.child('notes').subscribe()
    const pushed = watcher.waitFor('change', (e) => e.origin === 'remote')
    await new KVStorageEncrypted(network.client(), { key }).child('notes').set('1', { text: 'hi' })
    expect((await pushed).changes).toEqual([{ key: 'notes/1', value: { text: 'hi' } }])
  })

  it('works as a composite source: the server stores ciphertext, replicas hold plaintext', async () => {
    network = new TestNetwork()
    const replica = new KVStorageInMemory()
    const a = new KVStorageComposite({ source: new KVStorageEncrypted(network.client(), { key }), replicas: [replica] })
    const b = new KVStorageComposite({ source: new KVStorageEncrypted(network.client(), { key }), replicas: [new KVStorageInMemory()] })
    await Promise.all([a.sync(), b.sync()])
    const seen = b.waitFor('change', (e) => e.origin === 'remote')
    await a.child('notes').set('1', { text: 'hi' })
    expect((await seen).changes).toEqual([{ key: 'notes/1', value: { text: 'hi' } }])
    expect(await replica.readMany(['notes/1'])).toEqual([{ text: 'hi' }])
    expect((await network.store.readMany(['notes/1']))[0]).toBeInstanceOf(Uint8Array)
    await Promise.all([a.close(), b.close()])
  })
})