From 0c785831e745d14d2a2979e1e9225e7c7a4d5477 Mon Sep 17 00:00:00 2001
From: scanash00 <scan@scanash.com>
Date: Thu, 11 Jun 2026 22:34:14 -0800
Subject: [PATCH] extension and auth fixes

---
 backend/internal/oauth/client.go              |  36 +++++
 backend/internal/oauth/handler.go             |   8 +
 extension/src/components/popup/App.tsx        |  11 +-
 extension/src/entrypoints/background.ts       |   4 +
 .../src/entrypoints/margin-auth.content.ts    |   3 +
 extension/src/utils/api.ts                    | 153 +++++++++++-------
 extension/src/utils/messaging.ts              |   1 +
 extension/src/utils/storage.ts                |   6 +
 8 files changed, 164 insertions(+), 58 deletions(-)

diff --git a/backend/internal/oauth/client.go b/backend/internal/oauth/client.go
index 843c1a0..c0b0768 100644
--- a/backend/internal/oauth/client.go
+++ b/backend/internal/oauth/client.go
@@ -180,6 +180,42 @@ func (c *Client) ResolveDIDToPDS(ctx context.Context, did string) (string, error
 	return "", fmt.Errorf("no PDS found in DID document")
 }
 
+func (c *Client) ResolveDIDToHandle(ctx context.Context, did string) (string, error) {
+	var docURL string
+	if strings.HasPrefix(did, "did:plc:") {
+		docURL = config.Get().PLCResolveURL(did)
+	} else if strings.HasPrefix(did, "did:web:") {
+		domain := strings.TrimPrefix(did, "did:web:")
+		docURL = fmt.Sprintf("https://%s/.well-known/did.json", domain)
+	} else {
+		return "", fmt.Errorf("unsupported DID method: %s", did)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, docURL, nil)
+	if err != nil {
+		return "", err
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	var doc struct {
+		AlsoKnownAs []string `json:"alsoKnownAs"`
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&doc); err != nil {
+		return "", err
+	}
+
+	for _, aka := range doc.AlsoKnownAs {
+		if handle := strings.TrimPrefix(aka, "at://"); handle != aka {
+			return handle, nil
+		}
+	}
+	return "", fmt.Errorf("no handle found in DID document")
+}
+
 func (c *Client) GetAuthServerMetadata(ctx context.Context, pds string) (*AuthServerMetadata, error) {
 	resourceURL := fmt.Sprintf("%s/.well-known/oauth-protected-resource", strings.TrimSuffix(pds, "/"))
 	resp, err := http.Get(resourceURL)
diff --git a/backend/internal/oauth/handler.go b/backend/internal/oauth/handler.go
index 8eb60f7..8abeb9c 100644
--- a/backend/internal/oauth/handler.go
+++ b/backend/internal/oauth/handler.go
@@ -446,6 +446,14 @@ func (h *Handler) HandleCallback(w http.ResponseWriter, r *http.Request) {
 
 	_ = newNonce
 
+	if pending.Handle == "" {
+		if resolved, herr := client.ResolveDIDToHandle(ctx, tokenResp.Sub); herr == nil {
+			pending.Handle = resolved
+		} else {
+			logger.Error("Failed to resolve handle for %s: %v", tokenResp.Sub, herr)
+		}
+	}
+
 	sessionID := generateSessionID()
 	expiresAt := time.Now().Add(7 * 24 * time.Hour)
 
diff --git a/extension/src/components/popup/App.tsx b/extension/src/components/popup/App.tsx
index 0a342fc..f3c04a7 100644
--- a/extension/src/components/popup/App.tsx
+++ b/extension/src/components/popup/App.tsx
@@ -171,6 +171,15 @@ export function App() {
     }
   }
 
+  async function handleSignIn() {
+    try {
+      await browser.permissions.request({ origins: [`${new URL(apiUrl).origin}/*`] });
+    } catch {
+      /* empty */
+    }
+    await browser.tabs.create({ url: `${apiUrl}/login` });
+  }
+
   async function loadCurrentTab() {
     const [tab] = await browser.tabs.query({ active: true, currentWindow: true });
     if (tab?.url) {
@@ -474,7 +483,7 @@ export function App() {
             Annotate, highlight, and bookmark the web with your AT Protocol identity.
           </p>
           <button
-            onClick={() => browser.tabs.create({ url: `${apiUrl}/login` })}
+            onClick={handleSignIn}
             className="w-full max-w-[280px] px-6 py-2.5 bg-[var(--accent)] text-white rounded-xl text-sm font-semibold hover:bg-[var(--accent-hover)] transition-colors"
           >
             Sign In
diff --git a/extension/src/entrypoints/background.ts b/extension/src/entrypoints/background.ts
index 95e51a6..bfab13a 100644
--- a/extension/src/entrypoints/background.ts
+++ b/extension/src/entrypoints/background.ts
@@ -55,6 +55,10 @@ export default defineBackground(() => {
     return await checkSession();
   });
 
+  onMessage('sessionPing', async () => {
+    await checkSession();
+  });
+
   onMessage('getAnnotations', async ({ data, sender }) => {
     let citedUrls: string[] = data.citedUrls ?? [];
 
diff --git a/extension/src/entrypoints/margin-auth.content.ts b/extension/src/entrypoints/margin-auth.content.ts
index 7edc3ec..c50f071 100644
--- a/extension/src/entrypoints/margin-auth.content.ts
+++ b/extension/src/entrypoints/margin-auth.content.ts
@@ -1,4 +1,5 @@
 import { defineContentScript } from 'wxt/sandbox';
+import { sendMessage } from '@/utils/messaging';
 
 export default defineContentScript({
   matches: ['https://margin.at/*'],
@@ -17,5 +18,7 @@ export default defineContentScript({
           .catch((err) => ({ ok: false, status: 0, body: String(err) }));
       }
     });
+
+    sendMessage('sessionPing', undefined).catch(() => {});
   },
 });
diff --git a/extension/src/utils/api.ts b/extension/src/utils/api.ts
index dfc06bd..dc88ab9 100644
--- a/extension/src/utils/api.ts
+++ b/extension/src/utils/api.ts
@@ -1,25 +1,41 @@
 import type { MarginSession, TextSelector, Annotation } from './types';
-import { apiUrlItem } from './storage';
+import { apiUrlItem, cachedSessionItem } from './storage';
 
-async function safariFetch(url: string, options: RequestInit = {}): Promise<Response> {
+async function safariProxyFetch(url: string, options: RequestInit = {}): Promise<Response | null> {
   try {
-    const tabs = await browser.tabs.query({ url: 'https://margin.at/*' });
-    const tab = tabs[0];
-    if (!tab?.id) return new Response('{}', { status: 401 });
-
-    const result = (await browser.tabs.sendMessage(tab.id, {
-      type: 'SAFARI_FETCH',
-      url,
-      options: {
-        method: options.method || 'GET',
-        headers: options.headers || {},
-        body: options.body,
-      },
-    })) as { ok: boolean; status: number; body: string };
-
-    return new Response(result.body, { status: result.status });
+    const apiUrl = await getApiUrl();
+    const tabs = await browser.tabs.query({ url: `${new URL(apiUrl).origin}/*` });
+
+    for (const tab of tabs) {
+      if (!tab.id) continue;
+      try {
+        const result = (await browser.tabs.sendMessage(tab.id, {
+          type: 'SAFARI_FETCH',
+          url,
+          options: {
+            method: options.method || 'GET',
+            headers: options.headers || {},
+            body: options.body,
+          },
+        })) as { ok: boolean; status: number; body: string } | undefined;
+
+        if (result) return new Response(result.body, { status: result.status });
+      } catch {
+        continue;
+      }
+    }
+    return null;
   } catch {
-    return new Response('{}', { status: 500 });
+    return null;
+  }
+}
+
+async function hasApiHostPermission(): Promise<boolean> {
+  try {
+    const apiUrl = await getApiUrl();
+    return await browser.permissions.contains({ origins: [`${new URL(apiUrl).origin}/*`] });
+  } catch {
+    return false;
   }
 }
 
@@ -76,44 +92,58 @@ async function getSessionCookie(): Promise<string | null> {
   }
 }
 
+async function sessionFromResponse(res: Response): Promise<MarginSession> {
+  if (!res.ok) {
+    await cachedSessionItem.setValue(null);
+    return { authenticated: false };
+  }
+
+  const sessionData = await res.json();
+  if (!sessionData.did || !sessionData.handle) {
+    await cachedSessionItem.setValue(null);
+    return { authenticated: false };
+  }
+
+  const session: MarginSession = {
+    authenticated: true,
+    did: sessionData.did,
+    handle: sessionData.handle,
+    accessJwt: sessionData.accessJwt,
+    refreshJwt: sessionData.refreshJwt,
+  };
+  await cachedSessionItem.setValue(session);
+  return session;
+}
+
 export async function checkSession(): Promise<MarginSession> {
   try {
     const apiUrl = await getApiUrl();
-
-    if (import.meta.env.BROWSER === 'safari') {
-      const res = await safariFetch(`${apiUrl}/auth/session`);
-      if (!res.ok) return { authenticated: false };
-      const sessionData = await res.json();
-      if (!sessionData.did || !sessionData.handle) return { authenticated: false };
-      return {
-        authenticated: true,
-        did: sessionData.did,
-        handle: sessionData.handle,
-        accessJwt: sessionData.accessJwt,
-        refreshJwt: sessionData.refreshJwt,
-      };
-    }
-
     const cookie = await getSessionCookie();
 
-    if (!cookie) return { authenticated: false };
-
-    const res = await fetch(`${apiUrl}/auth/session`, {
-      headers: { 'X-Session-Token': cookie },
-    });
+    if (cookie) {
+      try {
+        const res = await fetch(`${apiUrl}/auth/session`, {
+          headers: { 'X-Session-Token': cookie },
+        });
+        return await sessionFromResponse(res);
+      } catch (error) {
+        if (import.meta.env.BROWSER !== 'safari') throw error;
+      }
+    }
 
-    if (!res.ok) return { authenticated: false };
+    if (import.meta.env.BROWSER === 'safari') {
+      const res = await safariProxyFetch(`${apiUrl}/auth/session`);
+      if (res) return await sessionFromResponse(res);
 
-    const sessionData = await res.json();
-    if (!sessionData.did || !sessionData.handle) return { authenticated: false };
+      if (!cookie && typeof browser.cookies?.get === 'function' && (await hasApiHostPermission())) {
+        await cachedSessionItem.setValue(null);
+        return { authenticated: false };
+      }
+      const cached = await cachedSessionItem.getValue();
+      if (cached?.authenticated) return cached;
+    }
 
-    return {
-      authenticated: true,
-      did: sessionData.did,
-      handle: sessionData.handle,
-      accessJwt: sessionData.accessJwt,
-      refreshJwt: sessionData.refreshJwt,
-    };
+    return { authenticated: false };
   } catch (error) {
     console.error('Session check error:', error);
     return { authenticated: false };
@@ -123,14 +153,7 @@ export async function checkSession(): Promise<MarginSession> {
 async function apiRequest(path: string, options: RequestInit = {}): Promise<Response> {
   const apiUrl = await getApiUrl();
   const apiPath = path.startsWith('/api') ? path : `/api${path}`;
-
-  if (import.meta.env.BROWSER === 'safari') {
-    const headers: Record<string, string> = {
-      'Content-Type': 'application/json',
-      ...(options.headers as Record<string, string>),
-    };
-    return safariFetch(`${apiUrl}${apiPath}`, { ...options, headers });
-  }
+  const url = `${apiUrl}${apiPath}`;
 
   const cookie = await getSessionCookie();
 
@@ -143,7 +166,23 @@ async function apiRequest(path: string, options: RequestInit = {}): Promise<Resp
     headers['X-Session-Token'] = cookie;
   }
 
-  return fetch(`${apiUrl}${apiPath}`, {
+  if (import.meta.env.BROWSER === 'safari') {
+    if (cookie) {
+      try {
+        return await fetch(url, { ...options, headers, credentials: 'include' });
+      } catch {
+        /* empty */
+      }
+    }
+    const res = await safariProxyFetch(url, { ...options, headers });
+    if (res) return res;
+    if (!cookie) {
+      return fetch(url, { ...options, headers, credentials: 'include' });
+    }
+    return new Response(JSON.stringify({ error: 'Could not reach margin.at' }), { status: 503 });
+  }
+
+  return fetch(url, {
     ...options,
     headers,
     credentials: 'include',
diff --git a/extension/src/utils/messaging.ts b/extension/src/utils/messaging.ts
index 8e82660..70236f7 100644
--- a/extension/src/utils/messaging.ts
+++ b/extension/src/utils/messaging.ts
@@ -10,6 +10,7 @@ import type {
 
 interface ProtocolMap {
   checkSession(): MarginSession;
+  sessionPing(): void;
 
   getAnnotations(data: { url: string; citedUrls?: string[]; cacheBust?: boolean }): Annotation[];
   activateOnPdf(data: { tabId: number; url: string }): { redirected: boolean };
diff --git a/extension/src/utils/storage.ts b/extension/src/utils/storage.ts
index e1c2dc9..899c0de 100644
--- a/extension/src/utils/storage.ts
+++ b/extension/src/utils/storage.ts
@@ -1,7 +1,13 @@
+import type { MarginSession } from './types';
+
 export const apiUrlItem = storage.defineItem<string>('local:apiUrl', {
   fallback: 'https://margin.at',
 });
 
+export const cachedSessionItem = storage.defineItem<MarginSession | null>('local:cachedSession', {
+  fallback: null,
+});
+
 export const overlayEnabledItem = storage.defineItem<boolean>('local:overlayEnabled', {
   fallback: true,
 });
-- 
2.51.2