Facilitate authing w/ PDS based on DID doc

lexicons/com/atproto/server/createAccount.json

@@ -29,7 +29,8 @@
             "accessJwt": { "type": "string" },
             "refreshJwt": { "type": "string" },
             "handle": { "type": "string", "format": "handle" },
-            "did": { "type": "string", "format": "did" }
+            "did": { "type": "string", "format": "did" },
+            "didDoc": { "type": "unknown" }
           }
         }
       },

lexicons/com/atproto/server/createSession.json

@@ -29,6 +29,7 @@
             "refreshJwt": { "type": "string" },
             "handle": { "type": "string", "format": "handle" },
             "did": { "type": "string", "format": "did" },
+            "didDoc": { "type": "unknown" },
             "email": { "type": "string" },
             "emailConfirmed": { "type": "boolean" }
           }

lexicons/com/atproto/server/refreshSession.json

@@ -14,7 +14,8 @@
             "accessJwt": { "type": "string" },
             "refreshJwt": { "type": "string" },
             "handle": { "type": "string", "format": "handle" },
-            "did": { "type": "string", "format": "did" }
+            "did": { "type": "string", "format": "did" },
+            "didDoc": { "type": "unknown" }
           }
         }
       },

packages/api/package.json

@@ -36,7 +36,8 @@
     "@atproto/xrpc": "workspace:^",
     "multiformats": "^9.9.0",
     "tlds": "^1.234.0",
-    "typed-emitter": "^2.1.0"
+    "typed-emitter": "^2.1.0",
+    "zod": "^3.21.4"
   },
   "devDependencies": {
     "@atproto/lex-cli": "workspace:^",

packages/api/src/agent.ts

@@ -18,6 +18,7 @@ import {
   AtpPersistSessionHandler,
   AtpAgentOpts,
 } from './types'
+import { getPdsEndpoint } from './did/did-doc'
 
 const REFRESH_SESSION = 'com.atproto.server.refreshSession'
 
@@ -30,6 +31,11 @@ export class AtpAgent {
   api: AtpServiceClient
   session?: AtpSessionData
 
+  /**
+   * The PDS URL, driven by the did doc. May be undefined.
+   */
+  pdsUrl: URL | undefined
+
   private _baseClient: AtpBaseClient
   private _persistSession?: AtpPersistSessionHandler
   private _refreshSessionPromise: Promise<void> | undefined
@@ -97,6 +103,7 @@ export class AtpAgent {
         email: opts.email,
         emailConfirmed: false,
       }
+      this._updateApiEndpoint(res.data.didDoc)
       return res
     } catch (e) {
       this.session = undefined
@@ -129,6 +136,7 @@ export class AtpAgent {
         email: res.data.email,
         emailConfirmed: res.data.emailConfirmed,
       }
+      this._updateApiEndpoint(res.data.didDoc)
       return res
     } catch (e) {
       this.session = undefined
@@ -157,6 +165,7 @@ export class AtpAgent {
       this.session.email = res.data.email
       this.session.handle = res.data.handle
       this.session.emailConfirmed = res.data.emailConfirmed
+      this._updateApiEndpoint(res.data.didDoc)
       return res
     } catch (e) {
       this.session = undefined
@@ -253,7 +262,7 @@ export class AtpAgent {
     }
 
     // send the refresh request
-    const url = new URL(this.service.origin)
+    const url = new URL((this.pdsUrl || this.service).origin)
     url.pathname = `/xrpc/${REFRESH_SESSION}`
     const res = await AtpAgent.fetch(
       url.toString(),
@@ -277,6 +286,7 @@ export class AtpAgent {
         handle: res.body.handle,
         did: res.body.did,
       }
+      this._updateApiEndpoint(res.body.didDoc)
       this._persistSession?.('update', this.session)
     }
     // else: other failures should be ignored - the issue will
@@ -311,6 +321,21 @@ export class AtpAgent {
    */
   createModerationReport: typeof this.api.com.atproto.moderation.createReport =
     (data, opts) => this.api.com.atproto.moderation.createReport(data, opts)
+
+  /**
+   * Helper to update the pds endpoint dynamically.
+   *
+   * The session methods (create, resume, refresh) may respond with the user's
+   * did document which contains the user's canonical PDS endpoint. That endpoint
+   * may differ from the endpoint used to contact the server. We capture that
+   * PDS endpoint and update the client to use that given endpoint for future
+   * requests. (This helps ensure smooth migrations between PDSes, especially
+   * when the PDSes are operated by a single org.)
+   */
+  private _updateApiEndpoint(didDoc: unknown) {
+    this.pdsUrl = getPdsEndpoint(didDoc)
+    this.api.xrpc.uri = this.pdsUrl || this.service
+  }
 }
 
 function isErrorObject(v: unknown): v is ErrorResponseBody {

packages/api/src/client/lexicons.ts

@@ -2338,6 +2338,9 @@ export const schemaDict = {
                 type: 'string',
                 format: 'did',
               },
+              didDoc: {
+                type: 'unknown',
+              },
             },
           },
         },
@@ -2563,6 +2566,9 @@ export const schemaDict = {
                 type: 'string',
                 format: 'did',
               },
+              didDoc: {
+                type: 'unknown',
+              },
               email: {
                 type: 'string',
               },
@@ -2879,6 +2885,9 @@ export const schemaDict = {
                 type: 'string',
                 format: 'did',
               },
+              didDoc: {
+                type: 'unknown',
+              },
             },
           },
         },

packages/api/src/client/types/com/atproto/server/createAccount.ts

@@ -24,6 +24,7 @@ export interface OutputSchema {
   refreshJwt: string
   handle: string
   did: string
+  didDoc?: {}
   [k: string]: unknown
 }
 

packages/api/src/client/types/com/atproto/server/createSession.ts

@@ -21,6 +21,7 @@ export interface OutputSchema {
   refreshJwt: string
   handle: string
   did: string
+  didDoc?: {}
   email?: string
   emailConfirmed?: boolean
   [k: string]: unknown

packages/api/src/client/types/com/atproto/server/refreshSession.ts

@@ -16,6 +16,7 @@ export interface OutputSchema {
   refreshJwt: string
   handle: string
   did: string
+  didDoc?: {}
   [k: string]: unknown
 }
 

packages/api/src/did/did-doc.ts

@@ -0,0 +1,40 @@
+import { z } from 'zod'
+
+export const verificationMethod = z.object({
+  id: z.string(),
+  type: z.string(),
+  controller: z.string(),
+  publicKeyMultibase: z.string().optional(),
+})
+
+export const service = z.object({
+  id: z.string(),
+  type: z.string(),
+  serviceEndpoint: z.union([z.string(), z.record(z.unknown())]),
+})
+
+export const didDocument = z.object({
+  id: z.string(),
+  alsoKnownAs: z.array(z.string()).optional(),
+  verificationMethod: z.array(verificationMethod).optional(),
+  service: z.array(service).optional(),
+})
+
+export type DidDocument = z.infer<typeof didDocument>
+
+export function isValidDidDoc(doc: unknown): doc is DidDocument {
+  return didDocument.safeParse(doc).success
+}
+
+export function getPdsEndpoint(doc: unknown): URL | undefined {
+  if (isValidDidDoc(doc)) {
+    const pds = doc.service?.find((s) => s.type === 'AtprotoPersonalDataServer')
+    if (pds && typeof pds.serviceEndpoint === 'string') {
+      try {
+        return new URL(pds.serviceEndpoint)
+      } catch {
+        return undefined
+      }
+    }
+  }
+}

packages/api/tests/agent.test.ts

@@ -13,6 +13,9 @@ describe('agent', () => {
   beforeAll(async () => {
     network = await TestNetworkNoAppView.create({
       dbPostgresSchema: 'api_agent',
+      pds: {
+        enableDidDocWithSession: true,
+      },
     })
   })
 

packages/bsky/src/lexicon/lexicons.ts

@@ -2338,6 +2338,9 @@ export const schemaDict = {
                 type: 'string',
                 format: 'did',
               },
+              didDoc: {
+                type: 'unknown',
+              },
             },
           },
         },
@@ -2563,6 +2566,9 @@ export const schemaDict = {
                 type: 'string',
                 format: 'did',
               },
+              didDoc: {
+                type: 'unknown',
+              },
               email: {
                 type: 'string',
               },
@@ -2879,6 +2885,9 @@ export const schemaDict = {
                 type: 'string',
                 format: 'did',
               },
+              didDoc: {
+                type: 'unknown',
+              },
             },
           },
         },

packages/bsky/src/lexicon/types/com/atproto/server/createAccount.ts

@@ -25,6 +25,7 @@ export interface OutputSchema {
   refreshJwt: string
   handle: string
   did: string
+  didDoc?: {}
   [k: string]: unknown
 }
 

packages/bsky/src/lexicon/types/com/atproto/server/createSession.ts

@@ -22,6 +22,7 @@ export interface OutputSchema {
   refreshJwt: string
   handle: string
   did: string
+  didDoc?: {}
   email?: string
   emailConfirmed?: boolean
   [k: string]: unknown

packages/bsky/src/lexicon/types/com/atproto/server/refreshSession.ts

@@ -17,6 +17,7 @@ export interface OutputSchema {
   refreshJwt: string
   handle: string
   did: string
+  didDoc?: {}
   [k: string]: unknown
 }
 

packages/dev-env/src/bin.ts

@@ -19,6 +19,7 @@ const run = async () => {
       port: 2583,
       hostname: 'localhost',
       dbPostgresSchema: 'pds',
+      enableDidDocWithSession: true,
     },
     bsky: {
       dbPostgresSchema: 'bsky',

packages/pds/src/api/com/atproto/server/createAccount.ts

@@ -1,16 +1,17 @@
+import { MINUTE } from '@atproto/common'
+import { AtprotoData } from '@atproto/identity'
 import { InvalidRequestError } from '@atproto/xrpc-server'
+import * as plc from '@did-plc/lib'
 import disposable from 'disposable-email'
 import { normalizeAndValidateHandle } from '../../../../handle'
-import * as plc from '@did-plc/lib'
 import * as scrypt from '../../../../db/scrypt'
 import { Server } from '../../../../lexicon'
 import { InputSchema as CreateAccountInput } from '../../../../lexicon/types/com/atproto/server/createAccount'
 import { countAll } from '../../../../db/util'
 import { UserAlreadyExistsError } from '../../../../services/account'
 import AppContext from '../../../../context'
 import Database from '../../../../db'
-import { AtprotoData } from '@atproto/identity'
-import { MINUTE } from '@atproto/common'
+import { didDocForSession } from './util'
 
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.server.createAccount({
@@ -118,11 +119,14 @@ export default function (server: Server, ctx: AppContext) {
         }
       })
 
+      const didDoc = await didDocForSession(ctx, result.did, true)
+
       return {
         encoding: 'application/json',
         body: {
           handle,
           did: result.did,
+          didDoc,
           accessJwt: result.accessJwt,
           refreshJwt: result.refreshJwt,
         },

packages/pds/src/api/com/atproto/server/createSession.ts

@@ -1,9 +1,10 @@
+import { DAY, MINUTE } from '@atproto/common'
 import { AuthRequiredError } from '@atproto/xrpc-server'
 import AppContext from '../../../../context'
 import { softDeleted } from '../../../../db/util'
 import { Server } from '../../../../lexicon'
 import { AuthScope } from '../../../../auth'
-import { DAY, MINUTE } from '@atproto/common'
+import { didDocForSession } from './util'
 
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.server.createSession({
@@ -60,12 +61,16 @@ export default function (server: Server, ctx: AppContext) {
         scope: appPasswordName === null ? AuthScope.Access : AuthScope.AppPass,
       })
       const refresh = ctx.auth.createRefreshToken({ did: user.did })
-      await authService.grantRefreshToken(refresh.payload, appPasswordName)
+      const [didDoc] = await Promise.all([
+        didDocForSession(ctx, user.did),
+        authService.grantRefreshToken(refresh.payload, appPasswordName),
+      ])
 
       return {
         encoding: 'application/json',
         body: {
           did: user.did,
+          didDoc,
           handle: user.handle,
           email: user.email,
           emailConfirmed: !!user.emailConfirmedAt,

packages/pds/src/api/com/atproto/server/refreshSession.ts

@@ -3,6 +3,7 @@ import AppContext from '../../../../context'
 import { softDeleted } from '../../../../db/util'
 import { Server } from '../../../../lexicon'
 import { AuthScope } from '../../../../auth'
+import { didDocForSession } from './util'
 
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.server.refreshSession({
@@ -50,10 +51,13 @@ export default function (server: Server, ctx: AppContext) {
         scope: res.appPassName === null ? AuthScope.Access : AuthScope.AppPass,
       })
 
+      const didDoc = await didDocForSession(ctx, user.did)
+
       return {
         encoding: 'application/json',
         body: {
           did: user.did,
+          didDoc,
           handle: user.handle,
           accessJwt: access.jwt,
           refreshJwt: res.refresh.jwt,