feat(mtp): move useful stuff over to mtp directly

packages/mtp/src/browser.tsx

@@ -0,0 +1,529 @@
+import { type ReactNode, useEffect, useMemo, useRef, useState } from "react";
+import { toast as sonnerToast } from "@methanium/ui";
+import { base64ToBytes, ConnectionState, MTPClient } from "mtp";
+import createAsyncQueue from "@tensamin/shared/asyncQueue";
+import {
+  mtp as mtpSchemas,
+  type Calls,
+  type Communities,
+  type Contacts,
+} from "@tensamin/shared/data";
+import { log } from "@tensamin/shared/log";
+import { useStorage } from "@tensamin/storage/context";
+
+import {
+  type BoundSendFn,
+  MTPContext,
+  type MTPContextType,
+  type ProtocolMessage,
+  removeMissingContacts,
+  useMessageHandlers,
+} from "./mtpContext";
+import {
+  DISCOVERY_TIMEOUT,
+  INITIAL_SYNC_TIMEOUT,
+  RECONNECT_JITTER,
+  RECONNECT_LONG_INTERVAL,
+  RECONNECT_RESET,
+  RECONNECT_TRIES,
+  RETRY_INTERVAL,
+  STATE_ACK_TIMEOUT,
+} from "./values";
+
+type BrowserMtpClient = Awaited<ReturnType<typeof createBrowserClient>>;
+
+function createBrowserClient(
+  options: Omit<Parameters<typeof MTPClient.create>[0], "schemas">,
+) {
+  return MTPClient.create({
+    ...options,
+    schemas: mtpSchemas,
+    throwProtocolErrors: true,
+    onValidationError: (error) => {
+      log(1, "mtp", "red", "Failed to validate push message", error);
+    },
+  });
+}
+
+function abortError(signal: AbortSignal): Error {
+  return signal.reason instanceof Error
+    ? signal.reason
+    : new Error("Initial state synchronization was cancelled");
+}
+
+function withDeadline<T>(
+  promise: Promise<T>,
+  timeoutMs: number,
+  timeoutMessage: string,
+  signal: AbortSignal,
+): Promise<T> {
+  return new Promise<T>((resolve, reject) => {
+    if (signal.aborted) {
+      reject(abortError(signal));
+      return;
+    }
+    let settled = false;
+    const finish = (complete: () => void) => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timeout);
+      signal.removeEventListener("abort", onAbort);
+      complete();
+    };
+    const timeout = setTimeout(
+      () => finish(() => reject(new Error(timeoutMessage))),
+      timeoutMs,
+    );
+    const onAbort = () => finish(() => reject(abortError(signal)));
+    signal.addEventListener("abort", onAbort, { once: true });
+    promise.then(
+      (value) => finish(() => resolve(value)),
+      (error: unknown) => finish(() => reject(error)),
+    );
+  });
+}
+
+async function completeInitialSynchronization(
+  client: BrowserMtpClient,
+  subscribe: MTPContextType["subscribe"],
+  signal: AbortSignal,
+  syncTimeoutMs = INITIAL_SYNC_TIMEOUT,
+  ackTimeoutMs = STATE_ACK_TIMEOUT,
+): Promise<ProtocolMessage<"ClientStateSync">> {
+  const stateSync = new Promise<ProtocolMessage<"ClientStateSync">>(
+    (resolve, reject) => {
+      let unsubscribeStateSync = () => {};
+      let unsubscribeNoIota = () => {};
+      const cleanup = () => {
+        clearTimeout(timeout);
+        unsubscribeStateSync();
+        unsubscribeNoIota();
+        signal.removeEventListener("abort", onAbort);
+      };
+      const onAbort = () => {
+        cleanup();
+        reject(abortError(signal));
+      };
+      const timeout = setTimeout(() => {
+        cleanup();
+        reject(new Error("Initial state synchronization timed out"));
+      }, syncTimeoutMs);
+      signal.addEventListener("abort", onAbort, { once: true });
+      if (signal.aborted) {
+        onAbort();
+        return;
+      }
+      unsubscribeStateSync = subscribe("ClientStateSync", (message) => {
+        cleanup();
+        resolve(message);
+      });
+      unsubscribeNoIota = subscribe("ErrorNoIota", () => {
+        cleanup();
+        reject(new Error("No Iota is currently connected"));
+      });
+    },
+  );
+  const [, state] = await Promise.all([
+    withDeadline(
+      client.auth(),
+      syncTimeoutMs,
+      "MTP authentication timed out",
+      signal,
+    ),
+    stateSync,
+  ]);
+  await withDeadline(
+    client.request("ClientStateAck", {
+      SessionId: state.data.SessionId,
+      VersionNumber: state.data.VersionNumber,
+    }),
+    ackTimeoutMs,
+    "State acknowledgement timed out",
+    signal,
+  );
+  if (signal.aborted) throw abortError(signal);
+  return state;
+}
+
+function protocolErrorDetails(error: unknown) {
+  if (typeof error !== "object" || error === null || !("type" in error)) {
+    return null;
+  }
+  const protocolError = error as {
+    id?: unknown;
+    type?: unknown;
+    frame?: unknown;
+  };
+  return {
+    id: protocolError.id,
+    type: protocolError.type,
+    frame: protocolError.frame,
+  };
+}
+
+export function BrowserProvider(props: {
+  children: ReactNode;
+  blockConnection?: boolean;
+}) {
+  const { load } = useStorage();
+  const [readyState, setReadyState] = useState<number>(
+    ConnectionState.Disconnected,
+  );
+  const [identified, setIdentified] = useState(false);
+  const [identifying, setIdentifying] = useState(false);
+  const [freshCommunities, setFreshCommunities] = useState<Communities>([]);
+  const [freshContacts, setFreshContacts] = useState<Contacts>([]);
+  const [freshCalls, setFreshCalls] = useState<Calls>([]);
+  const clientRef = useRef<BrowserMtpClient | null>(null);
+  const { addInterceptor, attachSubscriptions, interceptorsRef, subscribe } =
+    useMessageHandlers();
+  const connected = readyState === ConnectionState.Connected;
+
+  const [mtpUrl, setMtpUrl] = useState<string | null>(null);
+  useEffect(() => {
+    load("omega_url").then(setMtpUrl);
+  }, [load]);
+
+  const send: BoundSendFn = useMemo(
+    () => async (type, data, options) => {
+      const client = clientRef.current;
+      if (!client) throw new Error("mtp is not connected");
+      const response = await client.request(type, data, options);
+      if (response.type === "GetStates") {
+        setFreshContacts((contacts) =>
+          removeMissingContacts(
+            contacts,
+            response as ProtocolMessage<"GetStates">,
+          ),
+        );
+      }
+      return response;
+    },
+    [],
+  );
+
+  const resolveConnectionRef = useRef(() => {});
+  useEffect(() => {
+    if (!mtpUrl) return;
+    let attempts = 0;
+    let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+    let reconnectResetTimer: ReturnType<typeof setTimeout> | null = null;
+    let reconnectScheduled = false;
+    let disposed = false;
+    let connectionGeneration = 0;
+    let cleanupConnection = () => {};
+
+    const clearReconnectTimer = () => {
+      if (!reconnectTimer) return;
+      clearTimeout(reconnectTimer);
+      reconnectTimer = null;
+      reconnectScheduled = false;
+    };
+    const clearReconnectResetTimer = () => {
+      if (!reconnectResetTimer) return;
+      clearTimeout(reconnectResetTimer);
+      reconnectResetTimer = null;
+    };
+    const scheduleReconnect = (error: unknown) => {
+      if (disposed || reconnectScheduled) return;
+      attempts += 1;
+      const shortRetry = attempts <= RECONNECT_TRIES;
+      if (!shortRetry) {
+        log(0, "mtp", "red", "Reconnection attempts exhausted", error);
+        sonnerToast.error("Connection failed", {
+          id: "mtp-connection-toast",
+          description:
+            error instanceof Error
+              ? `${error.message.split(":")[0]}. Retrying in the background.`
+              : "Connection lost. Retrying in the background.",
+          icon: null,
+          duration: Infinity,
+          closeButton: true,
+          promise: null,
+        } as unknown as Parameters<typeof sonnerToast.error>[1]);
+      } else {
+        sonnerToast.loading(
+          `Reconnecting to server... (attempt ${attempts} of ${RECONNECT_TRIES})`,
+          { id: "mtp-connection-toast" },
+        );
+      }
+      const baseDelay = shortRetry ? RETRY_INTERVAL : RECONNECT_LONG_INTERVAL;
+      const jitter = 1 + (Math.random() * 2 - 1) * RECONNECT_JITTER;
+      reconnectScheduled = true;
+      reconnectTimer = setTimeout(
+        () => {
+          reconnectScheduled = false;
+          reconnectTimer = null;
+          void connect();
+        },
+        Math.round(baseDelay * jitter),
+      );
+    };
+
+    async function connect() {
+      if (disposed || props.blockConnection) return;
+      const generation = ++connectionGeneration;
+      let client: BrowserMtpClient | null = null;
+      let failed = false;
+      let connectionReady = false;
+      let detachSubscriptions = () => {};
+      let unsubscribeNoIota = () => {};
+      const attemptAbort = new AbortController();
+      const cleanup = () => {
+        attemptAbort.abort(
+          new Error("Initial state synchronization was cancelled"),
+        );
+        unsubscribeNoIota();
+        detachSubscriptions();
+        client?.disconnect();
+        if (clientRef.current === client) clientRef.current = null;
+        clearReconnectResetTimer();
+        if (generation === connectionGeneration) {
+          setReadyState(ConnectionState.Disconnected);
+          setIdentified(false);
+          setIdentifying(false);
+        }
+      };
+      cleanupConnection = cleanup;
+      try {
+        setIdentified(false);
+        setIdentifying(false);
+        const [userId, keyring] = await Promise.all([
+          load("user_id"),
+          load("mtp_keyring"),
+        ]);
+        if (!userId || !keyring) throw new Error("Missing login credentials");
+        const forcedOmikronUrl = await load("forced_omikron_url");
+        const forcedOmikronPublicKey = await load("forced_omikron_public_key");
+        let url = null;
+        let omikronPublicKey = null;
+        if (forcedOmikronUrl && forcedOmikronPublicKey) {
+          url = forcedOmikronUrl;
+          omikronPublicKey = forcedOmikronPublicKey;
+        } else {
+          log(2, "mtp", "purple", "Fetching Omikron data.");
+          const data = await fetch(`${mtpUrl}api/get/omikron/${userId}`, {
+            signal: AbortSignal.any([
+              attemptAbort.signal,
+              AbortSignal.timeout(DISCOVERY_TIMEOUT),
+            ]),
+          });
+          if (data.status === 404) {
+            throw new Error("No Omikron assignment is currently available");
+          }
+          if (!data.ok)
+            throw new Error(`Omikron discovery failed: HTTP ${data.status}`);
+          const omikronData = (await data.json()) as {
+            ip_address: string;
+            port: number;
+            public_key: string;
+          };
+          if (
+            !omikronData.ip_address ||
+            !omikronData.port ||
+            !omikronData.public_key
+          ) {
+            throw new Error("Invalid Omikron data");
+          }
+          url = `https://${omikronData.ip_address}:${omikronData.port}`;
+          omikronPublicKey = omikronData.public_key;
+        }
+        if (!url || !omikronPublicKey) {
+          throw new Error("Missing Omikron URL or Public Key");
+        }
+        log(2, "mtp", "green", "Connecting to: " + url);
+        client = await createBrowserClient({
+          url,
+          credentials: { clientId: userId, keyring: base64ToBytes(keyring) },
+          hostPublicKey: { value: omikronPublicKey, encoding: "base64" },
+          descriptor: "client",
+          pings: true,
+          logger: (event) => {
+            if (event.type === "state") {
+              if (generation !== connectionGeneration) return;
+              const state = client?.state ?? ConnectionState.Disconnected;
+              setReadyState(state);
+              if (
+                state === ConnectionState.Disconnected &&
+                clientRef.current === client &&
+                !failed
+              ) {
+                failed = true;
+                const error = new Error("MTP connection lost");
+                attemptAbort.abort(error);
+                if (connectionReady) {
+                  cleanup();
+                  scheduleReconnect(error);
+                }
+              }
+            }
+            if (event.type !== "Pong" && event.type !== "Ping") {
+              log(
+                2,
+                "mtp",
+                event.type === "state"
+                  ? "purple"
+                  : event.direction === "recv"
+                    ? "cyan"
+                    : event.direction === "send"
+                      ? "gray"
+                      : "blue",
+                event.type === "state"
+                  ? event.data
+                  : event.direction === "recv"
+                    ? "< " + event.type
+                    : event.direction === "send"
+                      ? "> " + event.type
+                      : event.type,
+                event,
+              );
+            }
+          },
+        });
+        if (disposed || generation !== connectionGeneration) {
+          client.disconnect();
+          return;
+        }
+        const activeClient = client;
+        clientRef.current = activeClient;
+        detachSubscriptions = attachSubscriptions(activeClient);
+        unsubscribeNoIota = subscribe("ErrorNoIota", () => {
+          if (clientRef.current !== activeClient || failed) return;
+          failed = true;
+          const error = new Error("No Iota is currently connected");
+          attemptAbort.abort(error);
+          cleanup();
+          scheduleReconnect(error);
+        });
+        setReadyState(activeClient.state);
+        setIdentifying(true);
+        const finalResponse = await completeInitialSynchronization(
+          activeClient,
+          subscribe,
+          attemptAbort.signal,
+        );
+        if (disposed || clientRef.current !== activeClient) return;
+        setFreshContacts(finalResponse.data.Contacts);
+        setFreshCommunities(finalResponse.data.Communities);
+        setFreshCalls(finalResponse.data.Calls);
+        connectionReady = true;
+        setIdentifying(false);
+        setIdentified(true);
+        clearReconnectTimer();
+        clearReconnectResetTimer();
+        reconnectResetTimer = setTimeout(() => {
+          attempts = 0;
+          reconnectResetTimer = null;
+        }, RECONNECT_RESET * 1_000);
+        resolveConnectionRef.current?.();
+      } catch (connectError) {
+        if (disposed || generation !== connectionGeneration) {
+          client?.disconnect();
+          return;
+        }
+        failed = true;
+        cleanup();
+        const message =
+          connectError instanceof Error
+            ? connectError.message
+            : String(connectError ?? "Unknown error");
+        log(
+          0,
+          "mtp",
+          "red",
+          `Connection/authentication attempt failed: ${message}`,
+          protocolErrorDetails(connectError) ?? connectError,
+        );
+        scheduleReconnect(connectError);
+      }
+    }
+
+    void connect();
+    return () => {
+      disposed = true;
+      clearReconnectTimer();
+      clearReconnectResetTimer();
+      cleanupConnection();
+      setReadyState(ConnectionState.Disconnected);
+      setIdentified(false);
+      setIdentifying(false);
+      sonnerToast.dismiss("mtp-connection-toast");
+    };
+  }, [attachSubscriptions, load, mtpUrl, props.blockConnection, subscribe]);
+
+  useEffect(() => {
+    return subscribe("ErrorNoIota", () => {
+      setIdentified(false);
+      setIdentifying(false);
+      sonnerToast.error("We couldn't reach your Iota", {
+        description:
+          "Check your network connection and try restarting your Iota",
+        icon: null,
+        duration: Infinity,
+        closeButton: true,
+      });
+      resolveConnectionRef.current?.();
+    });
+  }, [subscribe]);
+
+  useEffect(
+    () =>
+      subscribe("GetStates", (message) => {
+        setFreshContacts((contacts) =>
+          removeMissingContacts(contacts, message),
+        );
+      }),
+    [subscribe],
+  );
+
+  const loadingDescription = useMemo(() => {
+    if (!mtpUrl) return "Loading connection details";
+    if (readyState === ConnectionState.Connecting || !connected) {
+      return "Establishing transport channel";
+    }
+    if (identifying || !identified) return "Waiting for authenticated session";
+    return "Loading...";
+  }, [connected, identified, identifying, readyState, mtpUrl]);
+  const contextReady = connected && identified && mtpUrl !== null;
+  const mtpRef = useMemo(() => createAsyncQueue<{ send: typeof send }>(), []);
+  useEffect(() => {
+    if (connected && identified && mtpUrl) {
+      mtpRef.set({ send });
+    }
+  }, [connected, identified, mtpUrl, send, mtpRef]);
+
+  const sendQueued: BoundSendFn = useMemo(
+    () => async (type, data, options) => {
+      const mtp = await mtpRef.get();
+      const response = await mtp.send(type, data, options);
+      for (const interceptor of interceptorsRef.current) {
+        void Promise.resolve(interceptor({ type, data, response })).catch(
+          (error) => {
+            log(1, "mtp", "yellow", "MTP interceptor failed", error, { type });
+          },
+        );
+      }
+      return response;
+    },
+    [interceptorsRef, mtpRef],
+  );
+
+  return (
+    <MTPContext.Provider
+      value={{
+        send: sendQueued,
+        subscribe,
+        addInterceptor,
+        readyState,
+        identified,
+        freshContacts,
+        freshCommunities,
+        freshCalls,
+        contextReady,
+        loadingDescription,
+      }}
+    >
+      {props.children}
+    </MTPContext.Provider>
+  );
+}