Perf: Avoid fetching account data twice in putRecord (#4107)

* Perf: Avoid fetching account data twice in `putRecord`

* Apply same changes in `createRecord`/`deleteRecord`/`applyWrites`

* tidy

Co-authored-by: Copilot <[email protected]>

* switch validation order

---------

Co-authored-by: devin ivy <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

.changeset/fair-olives-boil.md

@@ -0,0 +1,5 @@
+---
+"@atproto/pds": patch
+---
+
+Allow use of handle in `applyWrites`'s `repo`

.changeset/fast-hairs-attack.md

@@ -0,0 +1,5 @@
+---
+"@atproto/pds": patch
+---
+
+Perf: Avoid fetching account data twice in `putRecord`

packages/pds/src/api/com/atproto/repo/applyWrites.ts

@@ -39,8 +39,15 @@ const ratelimitPoints = ({ input }: { input: HandlerInput }) => {
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.repo.applyWrites({
     auth: ctx.authVerifier.authorization({
-      checkTakedown: true,
-      checkDeactivated: true,
+      // @NOTE the "checkTakedown" and "checkDeactivated" checks are typically
+      // performed during auth. However, since this method's "repo" parameter
+      // can be a handle, we will need to fetch the account again to ensure that
+      // the handle matches the DID from the request's credentials. In order to
+      // avoid fetching the account twice (during auth, and then again in the
+      // controller), the checks are disabled here:
+
+      // checkTakedown: true,
+      // checkDeactivated: true,
       authorize: () => {
         // Performed in the handler as it is based on the request body
       },
@@ -62,10 +69,16 @@ export default function (server: Server, ctx: AppContext) {
     handler: async ({ input, auth }) => {
       const { repo, validate, swapCommit, writes } = input.body
 
-      const { did } = auth.credentials
-      if (repo !== did) {
+      const account = await ctx.authVerifier.findAccount(repo, {
+        checkDeactivated: true,
+        checkTakedown: true,
+      })
+
+      const did = account.did
+      if (did !== auth.credentials.did) {
         throw new AuthRequiredError()
       }
+
       if (writes.length > 200) {
         throw new InvalidRequestError('Too many writes. Max: 200')
       }

packages/pds/src/api/com/atproto/repo/createRecord.ts

@@ -15,8 +15,15 @@ import {
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.repo.createRecord({
     auth: ctx.authVerifier.authorization({
-      checkTakedown: true,
-      checkDeactivated: true,
+      // @NOTE the "checkTakedown" and "checkDeactivated" checks are typically
+      // performed during auth. However, since this method's "repo" parameter
+      // can be a handle, we will need to fetch the account again to ensure that
+      // the handle matches the DID from the request's credentials. In order to
+      // avoid fetching the account twice (during auth, and then again in the
+      // controller), the checks are disabled here:
+
+      // checkTakedown: true,
+      // checkDeactivated: true,
       authorize: () => {
         // Performed in the handler as it requires the request body
       },
@@ -37,26 +44,23 @@ export default function (server: Server, ctx: AppContext) {
       const { repo, collection, rkey, record, swapCommit, validate } =
         input.body
 
+      const account = await ctx.authVerifier.findAccount(repo, {
+        checkDeactivated: true,
+        checkTakedown: true,
+      })
+
+      const did = account.did
+      if (did !== auth.credentials.did) {
+        throw new AuthRequiredError()
+      }
+
       if (auth.credentials.type === 'oauth') {
         auth.credentials.permissions.assertRepo({
           action: 'create',
           collection,
         })
       }
 
-      const account = await ctx.accountManager.getAccount(repo, {
-        includeDeactivated: true,
-      })
-
-      if (!account) {
-        throw new InvalidRequestError(`Could not find repo: ${repo}`)
-      } else if (account.deactivatedAt) {
-        throw new InvalidRequestError('Account is deactivated')
-      }
-      const did = account.did
-      if (did !== auth.credentials.did) {
-        throw new AuthRequiredError()
-      }
       const swapCommitCid = swapCommit ? CID.parse(swapCommit) : undefined
 
       let write: PreparedCreate

packages/pds/src/api/com/atproto/repo/deleteRecord.ts

@@ -12,8 +12,15 @@ import {
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.repo.deleteRecord({
     auth: ctx.authVerifier.authorization({
-      checkTakedown: true,
-      checkDeactivated: true,
+      // @NOTE the "checkTakedown" and "checkDeactivated" checks are typically
+      // performed during auth. However, since this method's "repo" parameter
+      // can be a handle, we will need to fetch the account again to ensure that
+      // the handle matches the DID from the request's credentials. In order to
+      // avoid fetching the account twice (during auth, and then again in the
+      // controller), the checks are disabled here:
+
+      // checkTakedown: true,
+      // checkDeactivated: true,
       authorize: () => {
         // Performed in the handler as it requires the request body
       },
@@ -33,6 +40,16 @@ export default function (server: Server, ctx: AppContext) {
     handler: async ({ input, auth }) => {
       const { repo, collection, rkey, swapCommit, swapRecord } = input.body
 
+      const account = await ctx.authVerifier.findAccount(repo, {
+        checkDeactivated: true,
+        checkTakedown: true,
+      })
+
+      const did = account.did
+      if (did !== auth.credentials.did) {
+        throw new AuthRequiredError()
+      }
+
       // We can't compute permissions based on the request payload ("input") in
       // the 'auth' phase, so we do it here.
       if (auth.credentials.type === 'oauth') {
@@ -42,20 +59,6 @@ export default function (server: Server, ctx: AppContext) {
         })
       }
 
-      const account = await ctx.accountManager.getAccount(repo, {
-        includeDeactivated: true,
-      })
-
-      if (!account) {
-        throw new InvalidRequestError(`Could not find repo: ${repo}`)
-      } else if (account.deactivatedAt) {
-        throw new InvalidRequestError('Account is deactivated')
-      }
-      const did = account.did
-      if (did !== auth.credentials.did) {
-        throw new AuthRequiredError()
-      }
-
       const swapCommitCid = swapCommit ? CID.parse(swapCommit) : undefined
       const swapRecordCid = swapRecord ? CID.parse(swapRecord) : undefined
 

packages/pds/src/api/com/atproto/repo/putRecord.ts

@@ -21,8 +21,15 @@ import {
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.repo.putRecord({
     auth: ctx.authVerifier.authorization({
-      checkTakedown: true,
-      checkDeactivated: true,
+      // @NOTE the "checkTakedown" and "checkDeactivated" checks are typically
+      // performed during auth. However, since this method's "repo" parameter
+      // can be a handle, we will need to fetch the account again to ensure that
+      // the handle matches the DID from the request's credentials. In order to
+      // avoid fetching the account twice (during auth, and then again in the
+      // controller), the checks are disabled here:
+
+      // checkTakedown: true,
+      // checkDeactivated: true,
       authorize: () => {
         // Performed in the handler as it requires the request body
       },
@@ -50,6 +57,16 @@ export default function (server: Server, ctx: AppContext) {
         swapRecord,
       } = input.body
 
+      const account = await ctx.authVerifier.findAccount(repo, {
+        checkDeactivated: true,
+        checkTakedown: true,
+      })
+
+      const did = account.did
+      if (did !== auth.credentials.did) {
+        throw new AuthRequiredError()
+      }
+
       // We can't compute permissions based on the request payload ("input") in
       // the 'auth' phase, so we do it here.
       if (auth.credentials.type === 'oauth') {
@@ -63,20 +80,6 @@ export default function (server: Server, ctx: AppContext) {
         })
       }
 
-      const account = await ctx.accountManager.getAccount(repo, {
-        includeDeactivated: true,
-      })
-
-      if (!account) {
-        throw new InvalidRequestError(`Could not find repo: ${repo}`)
-      } else if (account.deactivatedAt) {
-        throw new InvalidRequestError('Account is deactivated')
-      }
-      const did = account.did
-      if (did !== auth.credentials.did) {
-        throw new AuthRequiredError()
-      }
-
       const uri = AtUri.make(did, collection, rkey)
       const swapCommitCid = swapCommit ? CID.parse(swapCommit) : undefined
       const swapRecordCid =

packages/pds/src/auth-verifier.ts

@@ -24,6 +24,7 @@ import {
   verifyJwt as verifyServiceJwt,
 } from '@atproto/xrpc-server'
 import { AccountManager } from './account-manager/account-manager'
+import { ActorAccount } from './account-manager/helpers/account'
 import {
   AccessOutput,
   AdminTokenOutput,
@@ -424,30 +425,43 @@ export class AuthVerifier {
 
   protected async verifyStatus(
     did: string,
-    { checkTakedown = false, checkDeactivated = false }: VerifiedOptions,
+    options: VerifiedOptions,
   ): Promise<void> {
-    if (checkTakedown || checkDeactivated) {
-      const found = await this.accountManager.getAccount(did, {
-        includeDeactivated: true,
-        includeTakenDown: true,
-      })
-      if (!found) {
-        // will be turned into ExpiredToken for the client if proxied by entryway
-        throw new ForbiddenError('Account not found', 'AccountNotFound')
-      }
-      if (checkTakedown && softDeleted(found)) {
-        throw new AuthRequiredError(
-          'Account has been taken down',
-          'AccountTakedown',
-        )
-      }
-      if (checkDeactivated && found.deactivatedAt) {
-        throw new AuthRequiredError(
-          'Account is deactivated',
-          'AccountDeactivated',
-        )
-      }
+    if (options.checkDeactivated || options.checkTakedown) {
+      await this.findAccount(did, options)
+    }
+  }
+
+  /**
+   * Finds an account by its handle or DID, returning possibly deactivated or
+   * taken down accounts (unless `options.checkDeactivated` or
+   * `options.checkTakedown` are set to true, respectively).
+   */
+  public async findAccount(
+    handleOrDid: string,
+    options: VerifiedOptions,
+  ): Promise<ActorAccount> {
+    const account = await this.accountManager.getAccount(handleOrDid, {
+      includeDeactivated: true,
+      includeTakenDown: true,
+    })
+    if (!account) {
+      // will be turned into ExpiredToken for the client if proxied by entryway
+      throw new ForbiddenError('Account not found', 'AccountNotFound')
+    }
+    if (options.checkTakedown && softDeleted(account)) {
+      throw new AuthRequiredError(
+        'Account has been taken down',
+        'AccountTakedown',
+      )
+    }
+    if (options.checkDeactivated && account.deactivatedAt) {
+      throw new AuthRequiredError(
+        'Account is deactivated',
+        'AccountDeactivated',
+      )
     }
+    return account
   }
 
   /**