@atproto/pds 0.4.59 → 0.4.60

package/src/read-after-write/util.ts

@@ -1,14 +1,26 @@
-import { Headers } from '@atproto/xrpc'
-import { readStickyLogger as log } from '../logger'
+import { jsonToLex } from '@atproto/lexicon'
+import { HeadersMap } from '@atproto/xrpc'
+import {
+  HandlerPipeThrough,
+  HandlerPipeThroughBuffer,
+  parseReqNsid,
+} from '@atproto/xrpc-server'
+import express from 'express'
+
 import AppContext from '../context'
+import { lexicons } from '../lexicon/lexicons'
+import { readStickyLogger as log } from '../logger'
+import {
+  asPipeThroughBuffer,
+  isJsonContentType,
+  pipethrough,
+} from '../pipethrough'
 import { HandlerResponse, LocalRecords, MungeFn } from './types'
 import { getRecordsSinceRev } from './viewer'
-import { HandlerPipeThrough } from '@atproto/xrpc-server'
-import { parseRes } from '../pipethrough'
 
 const REPO_REV_HEADER = 'atproto-repo-rev'
 
-export const getRepoRev = (headers: Headers): string | undefined => {
+export const getRepoRev = (headers: HeadersMap): string | undefined => {
   return headers[REPO_REV_HEADER]
 }
 
@@ -23,42 +35,56 @@ export const getLocalLag = (local: LocalRecords): number | undefined => {
   return Date.now() - new Date(oldest).getTime()
 }
 
-export const handleReadAfterWrite = async <T>(
+export const pipethroughReadAfterWrite = async <T>(
   ctx: AppContext,
-  nsid: string,
-  requester: string,
-  res: HandlerPipeThrough,
+  reqCtx: { req: express.Request; auth: { credentials: { did: string } } },
   munge: MungeFn<T>,
 ): Promise<HandlerResponse<T> | HandlerPipeThrough> => {
+  const { req, auth } = reqCtx
+  const requester = auth.credentials.did
+
+  const streamRes = await pipethrough(ctx, req, { iss: requester })
+
+  const rev = getRepoRev(streamRes.headers)
+  if (!rev) return streamRes
+
+  if (isJsonContentType(streamRes.headers['content-type']) === false) {
+    // content-type is present but not JSON, we can't munge this
+    return streamRes
+  }
+
+  // if the munging fails, we can't return the original response because the
+  // stream will already have been read. If we end-up buffering the response,
+  // we'll return the buffered response in case of an error.
+  let bufferRes: HandlerPipeThroughBuffer | undefined
+
   try {
-    return await readAfterWriteInternal(ctx, nsid, requester, res, munge)
+    const lxm = parseReqNsid(req)
+
+    return await ctx.actorStore.read(requester, async (store) => {
+      const local = await getRecordsSinceRev(store, rev)
+      if (local.count === 0) return streamRes
+
+      const { buffer } = (bufferRes = await asPipeThroughBuffer(streamRes))
+
+      const lex = jsonToLex(JSON.parse(buffer.toString('utf8')))
+
+      const parsedRes = lexicons.assertValidXrpcOutput(lxm, lex) as T
+
+      const localViewer = ctx.localViewer(store)
+
+      const data = await munge(localViewer, parsedRes, local, requester)
+      return formatMungedResponse(data, getLocalLag(local))
+    })
   } catch (err) {
+    // The error occurred while reading the stream, this is non-recoverable
+    if (!bufferRes && !streamRes.stream.readable) throw err
+
     log.warn({ err, requester }, 'error in read after write munge')
-    return res
+    return bufferRes ?? streamRes
   }
 }
 
-export const readAfterWriteInternal = async <T>(
-  ctx: AppContext,
-  nsid: string,
-  requester: string,
-  res: HandlerPipeThrough,
-  munge: MungeFn<T>,
-): Promise<HandlerResponse<T> | HandlerPipeThrough> => {
-  const rev = getRepoRev(res.headers ?? {})
-  if (!rev) return res
-  return ctx.actorStore.read(requester, async (store) => {
-    const local = await getRecordsSinceRev(store, rev)
-    if (local.count === 0) {
-      return res
-    }
-    const localViewer = ctx.localViewer(store)
-    const parsedRes = parseRes<T>(nsid, res)
-    const data = await munge(localViewer, parsedRes, local, requester)
-    return formatMungedResponse(data, getLocalLag(local))
-  })
-}
-
 export const formatMungedResponse = <T>(
   body: T,
   lag?: number,

package/tests/account-deletion.test.ts

@@ -78,7 +78,7 @@ describe('account deletion', () => {
     )
 
     expect(mail.to).toEqual(carol.email)
-    expect(mail.html).toContain('Delete your Bluesky account')
+    expect(mail.html).toContain('To permanently delete your account')
 
     token = getTokenFromMail(mail)
     if (!token) {

package/tests/account.test.ts

@@ -403,7 +403,7 @@ describe('account', () => {
     )
 
     expect(mail.to).toEqual(email)
-    expect(mail.html).toContain('Reset your password')
+    expect(mail.html).toContain('Reset password')
     expect(mail.html).toContain('alice.test')
 
     const token = getTokenFromMail(mail)
@@ -474,7 +474,7 @@ describe('account', () => {
     )
 
     expect(mail.to).toEqual(email)
-    expect(mail.html).toContain('Reset your password')
+    expect(mail.html).toContain('Reset password')
     expect(mail.html).toContain('alice.test')
 
     const token = getTokenFromMail(mail)

package/tests/email-confirmation.test.ts

@@ -92,7 +92,7 @@ describe('email confirmation', () => {
       }),
     )
     expect(mail.to).toEqual(alice.email)
-    expect(mail.html).toContain('Confirm your Bluesky email')
+    expect(mail.html).toContain('Confirm your email')
     confirmToken = getTokenFromMail(mail)
     expect(confirmToken).toBeDefined()
   })
@@ -164,7 +164,7 @@ describe('email confirmation', () => {
     }
     const mail = await getMailFrom(reqUpdate())
     expect(mail.to).toEqual(alice.email)
-    expect(mail.html).toContain('Update your Bluesky email')
+    expect(mail.html).toContain('Update your email')
     updateToken = getTokenFromMail(mail)
     expect(updateToken).toBeDefined()
   })

package/tests/plc-operations.test.ts

@@ -172,7 +172,7 @@ describe('plc operations', () => {
     )
 
     expect(mail.to).toEqual(sc.accounts[alice].email)
-    expect(mail.html).toContain('PLC Update Requested')
+    expect(mail.html).toContain('PLC update requested')
 
     const gotToken = getTokenFromMail(mail)
     assert(gotToken)

package/tests/proxied/proxy-catchall.test.ts

@@ -0,0 +1,255 @@
+import AtpAgent from '@atproto/api'
+import { Keypair } from '@atproto/crypto'
+import { TestNetworkNoAppView } from '@atproto/dev-env'
+import { LexiconDoc } from '@atproto/lexicon'
+import * as plc from '@did-plc/lib'
+import express from 'express'
+import getPort from 'get-port'
+import { once } from 'node:events'
+import http from 'node:http'
+import { setTimeout as sleep } from 'node:timers/promises'
+
+const lexicons = [
+  {
+    lexicon: 1,
+    id: 'com.example.ok',
+    defs: {
+      main: {
+        type: 'query',
+        output: {
+          encoding: 'application/json',
+          schema: { type: 'object', properties: { foo: { type: 'string' } } },
+        },
+      },
+    },
+  },
+  {
+    lexicon: 1,
+    id: 'com.example.slow',
+    defs: {
+      main: {
+        type: 'query',
+        output: {
+          encoding: 'application/json',
+          schema: { type: 'object', properties: { foo: { type: 'string' } } },
+        },
+      },
+    },
+  },
+  {
+    lexicon: 1,
+    id: 'com.example.abort',
+    defs: {
+      main: {
+        type: 'query',
+        output: {
+          encoding: 'application/json',
+          schema: { type: 'object', properties: { foo: { type: 'string' } } },
+        },
+      },
+    },
+  },
+  {
+    lexicon: 1,
+    id: 'com.example.error',
+    defs: {
+      main: {
+        type: 'query',
+        output: {
+          encoding: 'application/json',
+          schema: { type: 'object', properties: { foo: { type: 'string' } } },
+        },
+      },
+    },
+  },
+] as const satisfies LexiconDoc[]
+
+describe('proxy header', () => {
+  let network: TestNetworkNoAppView
+  let alice: AtpAgent
+
+  let proxyServer: ProxyServer
+
+  beforeAll(async () => {
+    network = await TestNetworkNoAppView.create({
+      dbPostgresSchema: 'proxy_catchall',
+    })
+
+    const serviceId = 'proxy_test'
+
+    proxyServer = await ProxyServer.create(
+      network.pds.ctx.plcClient,
+      network.pds.ctx.plcRotationKey,
+      serviceId,
+    )
+
+    alice = network.pds.getClient().withProxy(serviceId, proxyServer.did)
+
+    for (const lex of lexicons) alice.lex.add(lex)
+
+    await alice.createAccount({
+      email: 'alice@test.com',
+      handle: 'alice.test',
+      password: 'alice-pass',
+    })
+    await network.processAll()
+  })
+
+  afterAll(async () => {
+    await proxyServer?.close()
+    await network?.close()
+  })
+
+  it('rejects when upstream unavailable', async () => {
+    const serviceId = 'foo_bar'
+
+    const proxyServer = await ProxyServer.create(
+      network.pds.ctx.plcClient,
+      network.pds.ctx.plcRotationKey,
+      serviceId,
+    )
+
+    // Make sure the service is not available
+    await proxyServer.close()
+
+    const client = alice.withProxy(serviceId, proxyServer.did)
+    for (const lex of lexicons) client.lex.add(lex)
+
+    await expect(client.call('com.example.ok')).rejects.toThrow(
+      'pipethrough network error',
+    )
+  })
+
+  it('successfully proxies requests', async () => {
+    await expect(alice.call('com.example.ok')).resolves.toMatchObject({
+      data: { foo: 'ok' },
+      success: true,
+    })
+  })
+
+  it('handles cancelled upstream requests', async () => {
+    await expect(alice.call('com.example.abort')).rejects.toThrow('terminated')
+  })
+
+  it('handles failing upstream requests', async () => {
+    await expect(alice.call('com.example.error')).rejects.toThrowError(
+      expect.objectContaining({
+        status: 502,
+        error: 'FooBar',
+        message: 'My message',
+      }),
+    )
+  })
+
+  it('handles cancelled downstream requests', async () => {
+    const ac = new AbortController()
+
+    setTimeout(() => ac.abort(), 20)
+
+    await expect(
+      alice.call('com.example.slow', {}, undefined, { signal: ac.signal }),
+    ).rejects.toThrow('This operation was aborted')
+
+    await expect(alice.call('com.example.slow')).resolves.toMatchObject({
+      data: { foo: 'slow' },
+      success: true,
+    })
+  })
+})
+
+class ProxyServer {
+  constructor(
+    private server: http.Server,
+    public did: string,
+  ) {}
+
+  static async create(
+    plcClient: plc.Client,
+    keypair: Keypair,
+    serviceId: string,
+  ): Promise<ProxyServer> {
+    const app = express()
+
+    app.get('/xrpc/com.example.ok', (req, res) => {
+      res.status(200)
+      res.setHeader('content-type', 'application/json')
+      res.send('{"foo":"ok"}')
+    })
+
+    app.get('/xrpc/com.example.slow', async (req, res) => {
+      const wait = async (ms: number) => {
+        if (res.destroyed) return
+        const ac = new AbortController()
+        const abort = () => ac.abort()
+        res.on('close', abort)
+        try {
+          await sleep(ms, undefined, { signal: ac.signal })
+        } finally {
+          res.off('close', abort)
+        }
+      }
+
+      await wait(50)
+
+      res.status(200)
+      res.setHeader('content-type', 'application/json')
+      res.flushHeaders()
+
+      await wait(50)
+
+      for (const char of '{"foo":"slow"}') {
+        res.write(char)
+        await wait(10)
+      }
+
+      res.end()
+    })
+
+    app.get('/xrpc/com.example.abort', async (req, res) => {
+      res.status(200)
+      res.setHeader('content-type', 'application/json')
+      res.write('{"foo"')
+      await sleep(50)
+      res.destroy(new Error('abort'))
+    })
+
+    app.get('/xrpc/com.example.error', async (req, res) => {
+      res.status(500).json({ error: 'FooBar', message: 'My message' })
+    })
+
+    const port = await getPort()
+    const server = app.listen(port)
+    server.keepAliveTimeout = 30 * 1000
+    server.headersTimeout = 35 * 1000
+    await once(server, 'listening')
+
+    const plcOp = await plc.signOperation(
+      {
+        type: 'plc_operation',
+        rotationKeys: [keypair.did()],
+        alsoKnownAs: [],
+        verificationMethods: {},
+        services: {
+          [serviceId]: {
+            type: 'TestAtprotoService',
+            endpoint: `http://localhost:${port}`,
+          },
+        },
+        prev: null,
+      },
+      keypair,
+    )
+    const did = await plc.didForCreateOp(plcOp)
+    await plcClient.sendOperation(did, plcOp)
+    return new ProxyServer(server, did)
+  }
+
+  async close() {
+    await new Promise<void>((resolve, reject) => {
+      this.server.close((err) => {
+        if (err) reject(err)
+        else resolve()
+      })
+    })
+  }
+}

package/tests/proxied/proxy-header.test.ts

@@ -7,6 +7,7 @@ import { SeedClient, TestNetworkNoAppView, usersSeed } from '@atproto/dev-env'
 import getPort from 'get-port'
 import { Keypair } from '@atproto/crypto'
 import { verifyJwt } from '@atproto/xrpc-server'
+import { parseProxyHeader } from '../../src/pipethrough'
 
 describe('proxy header', () => {
   let network: TestNetworkNoAppView
@@ -51,6 +52,35 @@ describe('proxy header', () => {
     throw new Error('no error thrown')
   }
 
+  it('parses proxy header', async () => {
+    expect(parseProxyHeader(network.pds.ctx, `#atproto_test`)).rejects.toThrow(
+      'no did specified in proxy header',
+    )
+
+    expect(
+      parseProxyHeader(network.pds.ctx, `${proxyServer.did}#atproto_test#foo`),
+    ).rejects.toThrow('invalid proxy header format')
+
+    expect(
+      parseProxyHeader(network.pds.ctx, `${proxyServer.did}#atproto_test `),
+    ).rejects.toThrow('proxy header cannot contain spaces')
+
+    expect(
+      parseProxyHeader(network.pds.ctx, ` ${proxyServer.did}#atproto_test`),
+    ).rejects.toThrow('proxy header cannot contain spaces')
+
+    expect(parseProxyHeader(network.pds.ctx, `foo#bar`)).rejects.toThrow(
+      'Poorly formatted DID: foo',
+    )
+
+    expect(
+      parseProxyHeader(network.pds.ctx, `${proxyServer.did}#atproto_test`),
+    ).resolves.toEqual({
+      did: proxyServer.did,
+      url: proxyServer.url,
+    })
+  })
+
   it('proxies requests based on header', async () => {
     const path = `/xrpc/app.bsky.actor.getProfile?actor=${alice}`
     await axios.get(`${network.pds.url}${path}`, {
@@ -93,7 +123,7 @@ describe('proxy header', () => {
         'atproto-proxy': proxyServer.did,
       },
     })
-    await assertAxiosErr(attempt, 'no service id specified')
+    await assertAxiosErr(attempt, 'no service id specified in proxy header')
     expect(proxyServer.requests.length).toBe(1)
   })