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
@@ -253,7 +261,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 +285,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 +320,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,24 @@
+import { DidDocument, didDocument } from '@atproto/common-web'
+
+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.id === '#atproto_pds' || s.id === `${doc.id}#atproto_pds`,
+    )
+    if (
+      pds &&
+      pds.type === 'AtprotoPersonalDataServer' &&
+      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/common-web/src/did-doc.ts

@@ -0,0 +1,23 @@
+import { z } from 'zod'
+
+const verificationMethod = z.object({
+  id: z.string(),
+  type: z.string(),
+  controller: z.string(),
+  publicKeyMultibase: z.string().optional(),
+})
+
+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>

packages/common-web/src/index.ts

@@ -9,3 +9,4 @@ export * from './ipld'
 export * from './types'
 export * from './times'
 export * from './strings'
+export * from './did-doc'

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/identity/package.json

@@ -29,8 +29,7 @@
   "dependencies": {
     "@atproto/common-web": "workspace:^",
     "@atproto/crypto": "workspace:^",
-    "axios": "^0.27.2",
-    "zod": "^3.21.4"
+    "axios": "^0.27.2"
   },
   "devDependencies": {
     "@did-plc/lib": "^0.0.1",

packages/identity/src/types.ts

@@ -1,4 +1,7 @@
-import * as z from 'zod'
+import { DidDocument } from '@atproto/common-web'
+
+export { didDocument } from '@atproto/common-web'
+export type { DidDocument } from '@atproto/common-web'
 
 export type IdentityResolverOpts = {
   timeout?: number
@@ -42,25 +45,3 @@ export interface DidCache {
   clearEntry(did: string): Promise<void>
   clear(): Promise<void>
 }
-
-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>

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,
         },