(feat): better connection management, utility functions & tests

src/index.ts

@@ -4,7 +4,7 @@ import {
   READY_STATE,
 } from "@tensamin/ttp-core";
 import { createSchema } from "./schema";
-import type { ZodObject } from "zod";
+import type { z, ZodObject } from "zod";
 import { x448 } from "@noble/curves/ed448.js";
 import Bun from "bun";
 import { get } from "./user";
@@ -19,6 +19,7 @@ export class TAuthClient {
   identifier: string;
   privateKey: string;
   publicKey: string;
+  saveSession: (userId: number, sessionId: number) => Promise<void> | void;
   redirectUrl: URL;
   schema: ReturnType<typeof createSchema>;
   httpServer: {
@@ -26,12 +27,14 @@ export class TAuthClient {
     hostname: string;
   };
   clientMap = new Map<string, TransportClient<typeof this.schema>>();
+  private appData: ZodObject;
 
   constructor({
     frontendUrl,
     identifier,
     privateKey,
     publicKey,
+    saveSession,
     redirectUrl,
     appData,
     httpServer,
@@ -41,6 +44,7 @@ export class TAuthClient {
     identifier: string;
     privateKey: string;
     publicKey: string;
+    saveSession: (userId: number, sessionId: number) => Promise<void> | void;
     redirectUrl: string;
     appData: ZodObject;
     httpServer?: {
@@ -53,8 +57,13 @@ export class TAuthClient {
     this.identifier = identifier;
     this.privateKey = privateKey;
     this.publicKey = publicKey;
+    if (typeof saveSession !== "function") {
+      throw new Error("[tauth] saveSession must be provided as a function");
+    }
+    this.saveSession = saveSession;
     this.redirectUrl = new URL(redirectUrl);
     this.schema = createSchema(appData);
+    this.appData = appData;
     this.httpServer = httpServer ?? {
       port: 7878,
       hostname: "localhost",
@@ -102,11 +111,7 @@ export class TAuthClient {
               (await this.solveChallenge(userId, originalChallenge)) ===
               potentialChallenge
             ) {
-              const { ip_address } = await fetch(
-                "https://omega.tensamin.net/api/get/omikron/" + userId,
-              ).then((res) => res.json());
-
-              this.createTTP(userId, potentialSessionId, ip_address);
+              await this.saveSession(userId, potentialSessionId);
 
               return new Response(
                 htmlSuccessPage || (await readFile("./index.html")),
@@ -210,6 +215,76 @@ export class TAuthClient {
     this.clientMap.get(`${userId}:${sessionId}`)?.connect();
   }
 
+  async loadData(userId: number, sessionId: number) {
+    const connection = await this.createTemporaryTTPConnection(
+      userId,
+      sessionId,
+    );
+
+    const rawAppData = await connection.send("load_app_data", {});
+    const appData = JSON.parse(rawAppData.data.app_data);
+    const safeAppData = this.appData.safeParse(appData);
+
+    await connection.close();
+
+    if (!safeAppData.success) {
+      throw new Error(
+        `Invalid app data: ${JSON.stringify(safeAppData.error.issues)}`,
+      );
+    }
+
+    return safeAppData.data;
+  }
+
+  async saveData(userId: number, sessionId: number, data: any) {
+    const connection = await this.createTemporaryTTPConnection(
+      userId,
+      sessionId,
+    );
+
+    const safeData = this.appData.safeParse(data);
+
+    if (!safeData.success) {
+      throw new Error(
+        `Invalid app data: ${JSON.stringify(safeData.error.issues)}`,
+      );
+    }
+
+    await connection.send("save_app_data", {
+      app_data: JSON.stringify(safeData.data),
+    });
+
+    await connection.close();
+  }
+
+  private async createTemporaryTTPConnection(
+    userId: number,
+    sessionId: number,
+  ): Promise<TransportClient<typeof this.schema>> {
+    const { ip_address } = await fetch(
+      "https://omega.tensamin.net/api/get/omikron/" + userId,
+    ).then((res) => res.json());
+
+    let temporaryClient: TransportClient<typeof this.schema> | undefined;
+    temporaryClient = createTransportClient(this.schema, {
+      url: ip_address,
+      onReadyStateChange: (stateIndex) => {
+        const state = getFriendlyReadyState(stateIndex);
+        if (state === "OPEN") {
+          temporaryClient?.send("identification", {
+            app_identifier: this.identifier,
+            app_session_id: sessionId,
+            app_public_key: this.publicKey,
+            user_id: userId,
+          });
+        }
+      },
+    });
+
+    temporaryClient.connect();
+    return temporaryClient;
+  }
+
   generateLink(challenge?: string): string {
     return new URL(
       `?identifier=${this.identifier}&redirect=${this.redirectUrl}${challenge ? `&challenge=${challenge}` : ""}`,
@@ -251,4 +326,4 @@ export function generateKeyPair(): { private: string; public: string } {
 
 export * from "./crypto";
 export * from "./schema";
-export * from "./user";
+export * from "./user";

src/schema.ts

@@ -18,9 +18,20 @@ export function createSchema(appData: ZodObject) {
       request: z.object({
         challenge: z.base64(),
       }),
+      response: z.object({}),
+    },
+
+    load_app_data: {
+      request: z.object({}),
       response: z.object({
-        app_data: appData,
+        app_data: z.string(),
       }),
     },
+    save_app_data: {
+      request: z.object({
+        app_data: z.string(),
+      }),
+      response: z.object({}),
+    },
   };
 }