Resolve "Fix(Storage): Fix composite and RemoteHTTPStorage"

packages/stack/src/storage/store-composite.ts

@@ -1,3 +1,4 @@
+import type { OffCallback } from '../utils/index.js';
 import { BaseStore } from './base-store.js';
 import { StoreInMemory } from './store-inmemory.js';
 
@@ -21,11 +22,26 @@ export class StoreComposite extends BaseStore {
     return compositeStore;
   }
 
+  // Store the event listeners that are relevent to each store.
+  private eventListeners: Map<BaseStore, Array<OffCallback>> = new Map();
+
   constructor(
-    private stores: Array<BaseStore> = [],
+    public stores: Array<BaseStore> = [],
     private primaryStore: BaseStore,
   ) {
     super();
+
+    this.primaryStore.on('set', (payload) => {
+      this.emit('set', payload);
+    });
+
+    this.primaryStore.on('delete', (payload) => {
+      this.emit('delete', payload);
+    });
+
+    this.primaryStore.on('clear', () => {
+      this.emit('clear', undefined);
+    });
   }
 
   async all<T>() {
@@ -42,17 +58,14 @@ export class StoreComposite extends BaseStore {
 
   async set<T>(key: string, value: T) {
     await this.primaryStore.set(key, value);
-    this.emit('set', { key, value });
   }
 
   async delete(key: string) {
     await this.primaryStore.delete(key);
-    this.emit('delete', { key });
   }
 
   async clear() {
     await this.primaryStore.clear();
-    this.emit('clear', undefined);
   }
 
   async has(key: string) {
@@ -72,42 +85,76 @@ export class StoreComposite extends BaseStore {
   async addStore(remoteStore: BaseStore): Promise<void> {
     this.stores.push(remoteStore);
 
+    const offEventListeners = [];
+
     // Propagate existing local events to the new remote
-    this.primaryStore.on('set', async (payload) => {
+    offEventListeners.push(
-      if (await remoteStore.has(payload.key)) return;
+      this.primaryStore.on('set', async (payload) => {
-      await remoteStore.set(payload.key, payload.value);
+        if (await remoteStore.has(payload.key)) return;
-    });
+        await remoteStore.set(payload.key, payload.value);
+      }),
+    );
 
-    this.primaryStore.on('delete', async (payload) => {
+    offEventListeners.push(
-      if (await remoteStore.has(payload.key)) {
+      this.primaryStore.on('delete', async (payload) => {
-        await remoteStore.delete(payload.key);
+        if (await remoteStore.has(payload.key)) {
-      }
+          await remoteStore.delete(payload.key);
-    });
+        }
+      }),
+    );
 
-    this.primaryStore.on('clear', async () => {
+    offEventListeners.push(
-      await remoteStore.clear();
+      this.primaryStore.on('clear', async () => {
-    });
+        await remoteStore.clear();
+      }),
+    );
 
     // Set up remote to local propagation
-    remoteStore.on('set', async (payload) => {
+    offEventListeners.push(
-      if (await this.primaryStore.has(payload.key)) return;
+      remoteStore.on('set', async (payload) => {
-      await this.primaryStore.set(payload.key, payload.value);
+        if (await this.primaryStore.has(payload.key)) return;
-    });
+        await this.primaryStore.set(payload.key, payload.value);
+      }),
+    );
 
-    remoteStore.on('delete', async (payload) => {
+    offEventListeners.push(
-      if (await this.primaryStore.has(payload.key)) {
+      remoteStore.on('delete', async (payload) => {
-        await this.primaryStore.delete(payload.key);
+        if (await this.primaryStore.has(payload.key)) {
-      }
+          await this.primaryStore.delete(payload.key);
-    });
+        }
+      }),
+    );
 
-    remoteStore.on('clear', async () => {
+    offEventListeners.push(
-      await this.primaryStore.clear();
+      remoteStore.on('clear', async () => {
-    });
+        await this.primaryStore.clear();
+      }),
+    );
+
+    // Store our off callbacks for this store
+    this.eventListeners.set(remoteStore, offEventListeners);
 
     // Sync with the new remote
     await this.syncStore(remoteStore);
   }
 
+  /**
+   * Removes a store from the composite store
+   * @param store The store to remove
+   */
+  async removeStore(storeToRemove: BaseStore): Promise<void> {
+    this.stores = this.stores.filter((store) => store !== storeToRemove);
+
+    // Clean up event listeners
+    this.eventListeners
+      .get(storeToRemove)
+      ?.forEach((offCallback) => offCallback());
+    this.eventListeners.delete(storeToRemove);
+
+    // Close the store
+    await storeToRemove.close();
+  }
+
   async sync(): Promise<void> {
     // Sync all remote stores in parallel
     await Promise.all(

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

@@ -1,4 +1,4 @@
-import { BaseStore } from './base-store.js';
+import { BaseStorage, BaseStore, StorageOpts } from './base-store.js';
 import { SSESession } from '../utils/index.js';
 import {
   Bytes,
@@ -10,6 +10,106 @@ import {
 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,
@@ -25,9 +125,10 @@ export class StoreRemoteHTTP extends BaseStore {
 
   // Optimizations/Caches.
   protected publicKey: PublicKey;
+  protected sseSession?: SSESession;
 
   constructor(
-    protected baseURL: string,
+    public readonly baseURL: string,
     protected privateKey: PrivateKey,
     protected encrypter: Encrypter,
   ) {
@@ -49,10 +150,11 @@ export class StoreRemoteHTTP extends BaseStore {
       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(
+    const decryptedValue = await this.encrypter.decrypt(new Uint8Array(value));
-      new Uint8Array(response.value),
-    );
 
     // Parse the decrypted MessagePack value
     return unpack(decryptedValue) as T;
@@ -73,49 +175,84 @@ export class StoreRemoteHTTP extends BaseStore {
     await this.makeRequest('delete', { key });
   }
 
-  async subscribeEvents() {
+  /**
-    // Compute our signature (to authenticate with the SSE endpoint).
+   * Subscribes to events from the remote storage server.
-    const timestamp = Math.floor(Date.now() / 1000);
+   *
-    const signature = await StoreRemoteHTTP.signRequest(
+   * @remarks This method is called automatically when the store is created using the `StoreRemoteHTTP.from` method.
-      this.privateKey.toHex(),
+   * @returns void
-      timestamp,
+   */
-      'events',
+  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.
+      // Define our headers.
-    const headers = {
+      const headers = {
-      'Content-Type': 'application/json',
+        'Content-Type': 'application/json',
-      Accept: 'application/json',
+        Accept: 'application/json',
-      'X-Public-Key': this.publicKey.toHex(),
+        'X-Public-Key': this.publicKey.toHex(),
-      'X-Signature': Bytes.fromBytes(signature).toHex(),
+        'X-Signature': Bytes.fromBytes(signature).toHex(),
-      'X-Timestamp': timestamp.toString(),
+        '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).
-    await SSESession.from(`${this.baseURL}events`, {
+    this.sseSession = await SSESession.from(`${this.baseURL}events`, {
-      method: 'POST',
+      onMessage: async (e) => {
-      headers,
-      onMessage: (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: payload.value,
+          value: unpackedValue,
         });
       },
-      body: JSON.stringify({}),
+      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(
@@ -152,7 +289,7 @@ export class StoreRemoteHTTP extends BaseStore {
     }
 
     const responseBuffer = await response.arrayBuffer();
-    return unpack(new Uint8Array(responseBuffer));
+    return decodeExtendedJson(unpack(new Uint8Array(responseBuffer)));
   }
 
   static async signRequest(

packages/stack/src/utils/sse-session.ts

@@ -53,6 +53,7 @@ export class SSESession {
       onMessage: () => {},
       onError: (error) => console.error('SSESession error:', error),
       onDisconnected: () => {},
+      onReconnect: (options) => Promise.resolve(options),
 
       // Reconnection options
       attemptReconnect: true,
@@ -117,7 +118,17 @@ export class SSESession {
 
     // Establish the connection and get the reader using the exponential backoff
     const reader = await exponentialBackoff.run(async () => {
-      const res = await this.options.fetch(this.url, fetchOptions);
+      const reconnectOptions = await this.handleCallback(
+        this.options.onReconnect,
+        fetchOptions,
+      );
+
+      const updatedFetchOptions = {
+        ...fetchOptions,
+        ...reconnectOptions,
+      };
+
+      const res = await this.options.fetch(this.url, updatedFetchOptions);
       if (!res.ok) {
         throw new Error(`HTTP error! Status: ${res.status}`);
       }
@@ -130,7 +141,7 @@ export class SSESession {
     });
 
     // Call the onConnected callback
-    this.handleCallback(this.options.onConnected, undefined);
+    this.handleCallback(this.options.onConnected);
 
     const readStream = async () => {
       try {
@@ -292,9 +303,12 @@ export class SSESession {
    * @param callback The callback to handle.
    * @param args The arguments to pass to the callback.
    */
-  private handleCallback<T>(callback: (args: T) => void, args: T): void {
+  private handleCallback<T extends (...args: Parameters<T>) => ReturnType<T>>(
+    callback: T,
+    ...args: Parameters<T>
+  ): ReturnType<T> | undefined {
     try {
-      callback(args);
+      return callback(...args);
     } catch (error) {
       try {
         this.options.onError(error);
@@ -352,7 +366,12 @@ export interface SSESessionOptions {
   /**
    * Called when the connection is closed.
    */
-  onDisconnected: () => void;
+  onDisconnected: (error: unknown) => void;
+
+  /*
+   * Called when the connection is going to try to reconnect.
+   */
+  onReconnect: (options: RequestInit) => Promise<RequestInit>;
 
   /**
    * Whether to attempt to reconnect.