xo-stack/packages/stack/src/storage/store-remote-http.ts

import { BaseStorage, BaseStore, StorageOpts } from './base-store.js';
import { SSESession } from '../utils/index.js';
import {
  Bytes,
  Encrypter,
  PrivateKey,
  PublicKey,
  decodeExtendedJson,
} from '@xocash/primitives';
import { pack, unpack } from 'msgpackr';
import { z } from 'zod';

export type EncrptedStoreOpts = StorageOpts & {
  privateKey: PrivateKey;
  encrypter: Encrypter;
};

/**
 * A storage adapter for a remote storage server.
 *
 * All data sent to the server is first packed using MessagePack and then encrypted using the provided encrypter and then packed again using MessagePack.
 * The requests are signed using the provided private key and the timestamp of the request to mitigate replay attacks and to authenticate the request.
 */
export class StorageRemoteHTTP extends BaseStorage {
  /**
   * Creates a new StorageRemoteHTTP instance.
   *
   * @param baseURL The base URL of the remote storage server.
   * @returns A new StorageRemoteHTTP instance.
   */
  static async from(baseURL: string): Promise<StorageRemoteHTTP> {
    const storage = new this(baseURL);
    return storage;
  }

  private stores: Map<string, StoreRemoteHTTP> = new Map();

  constructor(protected baseURL: string) {
    super();
  }

  /**
   * Gets a store from the storage.
   *
   * @param storeName The name of the store to get.
   * @returns The store or null if it does not exist.
   */
  async getStore(storeName: string): Promise<StoreRemoteHTTP | null> {
    return this.stores.get(storeName) || null;
  }

  /**
   * Creates a new store in the remote storage using the `StoreRemoteHTTP.from` method.
   *
   * @param storeName The name of the store to create.
   * @param opts The options for the store. This includes the private key and encrypter.
   * @returns The created store.
   */
  async createStore(storeName: string, opts: EncrptedStoreOpts): Promise<StoreRemoteHTTP> {
    const store = await StoreRemoteHTTP.from(
      this.baseURL,
      opts.privateKey,
      opts.encrypter,
    );
    this.stores.set(storeName, store);
    return store;
  }

  /**
   * Creates a new store in the storage if it does not exist, otherwise returns the existing store.
   *
   * @param storeName The name of the store to create or get.
   * @param opts The options for the store. This includes the private key and encrypter.
   * @returns The created or existing store.
   */
  async createOrGetStore(storeName: string, opts: EncrptedStoreOpts): Promise<StoreRemoteHTTP> {
    const store = await this.getStore(storeName);
    if (store) {
      return store;
    }
    return this.createStore(storeName, opts);
  }

  /**
   * Deletes a store from the storage.
   *
   * @param storeName The name of the store to delete.
   * @returns void
   */
  async deleteStore(storeName: string): Promise<void> {
    this.stores.delete(storeName);
  }

  /**
   * Lists all stores in the storage.
   *
   * @returns An array of store names.
   */
  async listStores(): Promise<string[]> {
    return Array.from(this.stores.keys());
  }

  /**
   * Deletes the entire all stores from the remote storage server.
   *
   * @returns void
   */
  async deleteDatabase(): Promise<void> {
    this.stores.clear();
  }
}

export class StoreRemoteHTTP extends BaseStore {
  static async from(
    baseURL: string,
    privateKey: PrivateKey,
    encrypter: Encrypter,
  ) {
    const store = new this(baseURL, privateKey, encrypter);

    await store.subscribeEvents();

    return store;
  }

  // Optimizations/Caches.
  protected publicKey: PublicKey;
  protected sseSession?: SSESession;

  constructor(
    public readonly baseURL: string,
    protected privateKey: PrivateKey,
    protected encrypter: Encrypter,
  ) {
    super();

    // Derive the public key and cache it to optimize performance.
    this.publicKey = privateKey.derivePublicKey();
  }

  async keys() {
    const response = await this.makeRequest('keys');
    return response;
  }

  async get<T>(key: string) {
    const response = await this.makeRequest('get', { key });

    if (!response.value) {
      return undefined;
    }

    // Value is passed as an object: { index: number }. We need to convert this to an array so we can pass it to uint8array.
    const value: number[] = Object.values(response.value);

    // Decrypt the value using the encrypter
    const decryptedValue = await this.encrypter.decrypt(new Uint8Array(value));

    // Parse the decrypted MessagePack value
    return unpack(decryptedValue) as T;
  }

  async set<T>(key: string, value: T): Promise<void> {
    // Pack the encoded value with MessagePack and encrypt it
    const packedValue = pack(value);
    const encryptedValue = await this.encrypter.encrypt(packedValue);

    await this.makeRequest('set', {
      key,
      value: encryptedValue, // Convert Uint8Array to regular array for MessagePack
    });
  }

  async delete(key: string): Promise<void> {
    await this.makeRequest('delete', { key });
  }

  /**
   * Subscribes to events from the remote storage server.
   *
   * @remarks This method is called automatically when the store is created using the `StoreRemoteHTTP.from` method.
   * @returns void
   */
  async subscribeEvents(): Promise<void> {
    const createSSESessionOptions = async () => {
      // Compute our signature (to authenticate with the SSE endpoint).
      const timestamp = Math.floor(Date.now() / 1000);
      const signature = await StoreRemoteHTTP.signRequest(
        this.privateKey.toHex(),
        timestamp,
        'events',
        {},
      );

      // Define our headers.
      const headers = {
        'Content-Type': 'application/json',
        Accept: 'application/json',
        'X-Public-Key': this.publicKey.toHex(),
        'X-Signature': Bytes.fromBytes(signature).toHex(),
        'X-Timestamp': timestamp.toString(),
      };

      // Return the options.
      return {
        method: 'POST',
        headers,
        body: JSON.stringify({}),
      } as const;
    };

    const options = await createSSESessionOptions();

    // Create our SSE Client.
    // NOTE: SSE only supports text-encoding, so we use Extended JSON instead of msgpackr (which is binary).
    // TODO: It would be really nice to somehow be able to support msgpackr for consistency though.
    //       But, from a bandwidth point of view, this is not crucial: These payloads are "live" payloads (thus, there's not much data).
    this.sseSession = await SSESession.from(`${this.baseURL}events`, {
      onMessage: async (e) => {
        // Validate the payload using Zod.
        // TODO: Use msgpackr
        const payload = StoreRemoteHTTPEventSchema.parse(
          decodeExtendedJson(e.data),
        );

        // Decode, Decrypt, and Unpack the value.
        const encryptedValue = Bytes.fromBase64(payload.value);
        const decryptedValue = await this.encrypter.decrypt(encryptedValue);
        const unpackedValue = unpack(decryptedValue);

        // Emit the event (update) to any listeners.
        this.emit(payload.action, {
          key: payload.key,
          value: unpackedValue,
        });
      },
      onReconnect: createSSESessionOptions,
      ...options,
    });
  }

  async close() {
    await this.sseSession?.close();
    this.sseSession = undefined;
  }

  /**
   * Make a request to the remote storage server.
   *
   * @remarks The payload is encoded with msgpackr. This means that if the request contains an encrypted payload, you will need to 1. unpack the payload, 2. decrypt the value, 3. unpack the value. We pack the value before encrypting it so we can have a single interface that we are encrypting
   *
   * @param endpoint
   * @param payload
   * @returns
   */
  private async makeRequest(endpoint: string, payload: object = {}) {
    const timestamp = Math.floor(Date.now() / 1000);
    const signature = await StoreRemoteHTTP.signRequest(
      this.privateKey.toHex(),
      timestamp,
      endpoint,
      payload,
    );

    const headers = {
      'Content-Type': 'application/msgpackr',
      Accept: 'application/msgpackr',
      'X-Public-Key': this.publicKey.toHex(),
      'X-Signature': Bytes.fromBytes(signature).toHex(),
      'X-Timestamp': timestamp.toString(),
    };

    const requestOptions = {
      method: 'POST',
      headers,
      body: payload ? pack(payload) : pack({}),
    };

    // Use the URL class to compile the full path.
    // NOTE: We do this to cater for those situations where the base URL might contain a trailing slash.
    //       E.g. https://wallet-sync.xo.cash/store/ vs https://wallet-sync.xo.cash/store
    // TODO: THis is broken and currently requires a trailing slash in the BaseURL!
    const requestURL = new URL(endpoint, this.baseURL).toString();

    const response = await fetch(requestURL, requestOptions);

    if (!response.ok) {
      throw new Error(`HTTP error! status: ${response.status}`);
    }

    const responseBuffer = await response.arrayBuffer();
    return decodeExtendedJson(unpack(new Uint8Array(responseBuffer)));
  }

  static async signRequest(
    privateKeyHex: string,
    timestamp: number,
    endpoint: string,
    payload?: object,
  ) {
    const messageBin = pack([endpoint, timestamp, payload]);

    return PrivateKey.fromHex(privateKeyHex).signMessageSchnorr(messageBin);
  }

  static async verifySignature(
    publicKeyHex: string,
    signatureHex: string,
    timestamp: number,
    endpoint: string,
    payload?: object,
  ) {
    const now = Math.floor(Date.now() / 1000);
    if (Math.abs(now - timestamp) > 60) {
      throw new Error(`Request has expired`);
    }

    const messageBin = pack([endpoint, timestamp, payload]);

    const isValidSignature = PublicKey.fromHex(
      publicKeyHex,
    ).verifyMessageSchnorr(messageBin, signatureHex);

    if (!isValidSignature) {
      throw new Error(`The provided signature is invalid`);
    }
  }
}

export const StoreRemoteHTTPEventSchema = z.object({
  publicKey: z.string(),
  action: z.enum(['set', 'delete', 'clear']),
  key: z.string(),
  value: z.string(),
  timestamp: z.number(),
});

export type StoreRemoteHTTPEvent = z.infer<typeof StoreRemoteHTTPEventSchema>;