PDS proxy to appview performance (#2773)

* accept entryway session tokens * extra check + tests * build * build * pr feedback --------- Co-authored-by: Devin Ivy <[email protected]>
bluesky-social · Sep 6, 2024 · 6c1ec14 · 6c1ec14
1 parent 71785d3
commit 6c1ec14
Show file tree

Hide file tree

Showing 7 changed files with 270 additions and 2 deletions.
diff --git a/.github/workflows/build-and-push-bsky-ghcr.yaml b/.github/workflows/build-and-push-bsky-ghcr.yaml
@@ -3,6 +3,7 @@ on:
   push:
     branches:
       - main
+      - bsky-tweaks
 env:
   REGISTRY: ghcr.io
   USERNAME: ${{ github.actor }}

diff --git a/packages/bsky/package.json b/packages/bsky/package.json
@@ -48,6 +48,7 @@
     "http-terminator": "^3.2.0",
     "ioredis": "^5.3.2",
     "jose": "^5.0.1",
+    "key-encoder": "^2.0.3",
     "kysely": "^0.22.0",
     "multiformats": "^9.9.0",
     "p-queue": "^6.6.2",

diff --git a/packages/bsky/src/auth-verifier.ts b/packages/bsky/src/auth-verifier.ts
@@ -1,9 +1,12 @@
+import { KeyObject, createPublicKey } from 'node:crypto'
 import {
   AuthRequiredError,
   parseReqNsid,
   verifyJwt as verifyServiceJwt,
 } from '@atproto/xrpc-server'
+import KeyEncoder from 'key-encoder'
 import * as ui8 from 'uint8arrays'
+import * as jose from 'jose'
 import express from 'express'
 import {
   Code,
@@ -59,18 +62,26 @@ type ModServiceOutput = {
   }
 }
 
+const ALLOWED_AUTH_SCOPES = new Set([
+  'com.atproto.access',
+  'com.atproto.appPass',
+  'com.atproto.appPassPrivileged',
+])
+
 export type AuthVerifierOpts = {
   ownDid: string
   alternateAudienceDids: string[]
   modServiceDid: string
   adminPasses: string[]
+  entrywayJwtPublicKey?: KeyObject
 }
 
 export class AuthVerifier {
   public ownDid: string
   public standardAudienceDids: Set<string>
   public modServiceDid: string
   private adminPasses: Set<string>
+  private entrywayJwtPublicKey?: KeyObject
 
   constructor(
     public dataplane: DataPlaneClient,
@@ -83,6 +94,7 @@ export class AuthVerifier {
     ])
     this.modServiceDid = opts.modServiceDid
     this.adminPasses = new Set(opts.adminPasses)
+    this.entrywayJwtPublicKey = opts.entrywayJwtPublicKey
   }
 
   // verifiers (arrow fns to preserve scope)
@@ -103,6 +115,17 @@ export class AuthVerifier {
           credentials: { type: 'standard', iss, aud },
         }
       } else if (isBearerToken(ctx.req)) {
+        // @NOTE temporarily accept entryway session tokens to shed load from PDS instances
+        const token = bearerTokenFromReq(ctx.req)
+        const header = token ? jose.decodeProtectedHeader(token) : undefined
+        if (header?.typ === 'at+jwt') {
+          // we should never use entryway session tokens in the case of flexible auth audiences (namely in the case of getFeed)
+          if (opts.skipAudCheck) {
+            throw new AuthRequiredError('Malformed token', 'InvalidToken')
+          }
+          return this.entrywaySession(ctx)
+        }
+
         const { iss, aud } = await this.verifyServiceJwt(ctx, {
           lxmCheck: opts.lxmCheck,
           iss: null,
@@ -182,6 +205,54 @@ export class AuthVerifier {
     }
   }
 
+  // @NOTE this auth verifier method is not recommended to be implemented by most appviews
+  // this is a short term fix to remove proxy load from Bluesky's PDS and in line with possible
+  // future plans to have the client talk directly with the appview
+  entrywaySession = async (reqCtx: ReqCtx): Promise<StandardOutput> => {
+    const token = bearerTokenFromReq(reqCtx.req)
+    if (!token) {
+      throw new AuthRequiredError(undefined, 'AuthMissing')
+    }
+
+    // if entryway jwt key not configured then do not parsed these tokens
+    if (!this.entrywayJwtPublicKey) {
+      throw new AuthRequiredError('Malformed token', 'InvalidToken')
+    }
+
+    const res = await jose
+      .jwtVerify(token, this.entrywayJwtPublicKey)
+      .catch((err) => {
+        if (err?.['code'] === 'ERR_JWT_EXPIRED') {
+          throw new AuthRequiredError('Token has expired', 'ExpiredToken')
+        }
+        throw new AuthRequiredError(
+          'Token could not be verified',
+          'InvalidToken',
+        )
+      })
+
+    const { sub, aud, scope } = res.payload
+    if (typeof sub !== 'string' || !sub.startsWith('did:')) {
+      throw new AuthRequiredError('Malformed token', 'InvalidToken')
+    } else if (
+      typeof aud !== 'string' ||
+      !aud.startsWith('did:web:') ||
+      !aud.endsWith('.bsky.network')
+    ) {
+      throw new AuthRequiredError('Bad token aud', 'InvalidToken')
+    } else if (typeof scope !== 'string' || !ALLOWED_AUTH_SCOPES.has(scope)) {
+      throw new AuthRequiredError('Bad token scope', 'InvalidToken')
+    }
+
+    return {
+      credentials: {
+        type: 'standard',
+        aud: this.ownDid,
+        iss: sub,
+      },
+    }
+  }
+
   modService = async (reqCtx: ReqCtx): Promise<ModServiceOutput> => {
     const { iss, aud } = await this.verifyServiceJwt(reqCtx, {
       aud: this.ownDid,
@@ -365,3 +436,9 @@ export const buildBasicAuth = (username: string, password: string): string => {
     ui8.toString(ui8.fromString(`${username}:${password}`, 'utf8'), 'base64pad')
   )
 }
+
+const keyEncoder = new KeyEncoder('secp256k1')
+export const createPublicKeyObject = (publicKeyHex: string): KeyObject => {
+  const key = keyEncoder.encodePublic(publicKeyHex, 'raw', 'pem')
+  return createPublicKey({ format: 'pem', key })
+}
diff --git a/packages/bsky/src/config.ts b/packages/bsky/src/config.ts
@@ -8,6 +8,7 @@ export interface ServerConfigValues {
   publicUrl?: string
   serverDid: string
   alternateAudienceDids: string[]
+  entrywayJwtPublicKeyHex?: string
   // external services
   dataplaneUrls: string[]
   dataplaneHttpVersion?: '1.1' | '2'
@@ -56,6 +57,8 @@ export class ServerConfig {
     const alternateAudienceDids = process.env.BSKY_ALT_AUDIENCE_DIDS
       ? process.env.BSKY_ALT_AUDIENCE_DIDS.split(',')
       : []
+    const entrywayJwtPublicKeyHex =
+      process.env.BSKY_ENTRYWAY_JWT_PUBLIC_KEY_HEX || undefined
     const handleResolveNameservers = process.env.BSKY_HANDLE_RESOLVE_NAMESERVERS
       ? process.env.BSKY_HANDLE_RESOLVE_NAMESERVERS.split(',')
       : []
@@ -126,6 +129,7 @@ export class ServerConfig {
       publicUrl,
       serverDid,
       alternateAudienceDids,
+      entrywayJwtPublicKeyHex,
       dataplaneUrls,
       dataplaneHttpVersion,
       dataplaneIgnoreBadTls,
@@ -194,6 +198,10 @@ export class ServerConfig {
     return this.cfg.alternateAudienceDids
   }
 
+  get entrywayJwtPublicKeyHex() {
+    return this.cfg.entrywayJwtPublicKeyHex
+  }
+
   get dataplaneUrls() {
     return this.cfg.dataplaneUrls
   }

diff --git a/packages/bsky/src/index.ts b/packages/bsky/src/index.ts
@@ -20,7 +20,7 @@ import { Keypair } from '@atproto/crypto'
 import { createDataPlaneClient } from './data-plane/client'
 import { Hydrator } from './hydration/hydrator'
 import { Views } from './views'
-import { AuthVerifier } from './auth-verifier'
+import { AuthVerifier, createPublicKeyObject } from './auth-verifier'
 import { authWithApiKey as bsyncAuth, createBsyncClient } from './bsync'
 import { authWithApiKey as courierAuth, createCourierClient } from './courier'
 import { FeatureGates } from './feature-gates'
@@ -119,11 +119,15 @@ export class BskyAppView {
         : [],
     })
 
+    const entrywayJwtPublicKey = config.entrywayJwtPublicKeyHex
+      ? createPublicKeyObject(config.entrywayJwtPublicKeyHex)
+      : undefined
     const authVerifier = new AuthVerifier(dataplane, {
       ownDid: config.serverDid,
       alternateAudienceDids: config.alternateAudienceDids,
       modServiceDid: config.modServiceDid,
       adminPasses: config.adminPasswords,
+      entrywayJwtPublicKey,
     })
 
     const featureGates = new FeatureGates({

diff --git a/packages/bsky/tests/entryway-auth.test.ts b/packages/bsky/tests/entryway-auth.test.ts
@@ -0,0 +1,174 @@
+import * as nodeCrypto from 'node:crypto'
+import KeyEncoder from 'key-encoder'
+import * as ui8 from 'uint8arrays'
+import * as jose from 'jose'
+import * as crypto from '@atproto/crypto'
+import { AtpAgent, AtUri } from '@atproto/api'
+import { basicSeed, SeedClient, TestNetwork } from '@atproto/dev-env'
+import assert from 'node:assert'
+import { MINUTE } from '@atproto/common'
+
+const keyEncoder = new KeyEncoder('secp256k1')
+
+const derivePrivKey = async (
+  keypair: crypto.ExportableKeypair,
+): Promise<nodeCrypto.KeyObject> => {
+  const privKeyRaw = await keypair.export()
+  const privKeyEncoded = keyEncoder.encodePrivate(
+    ui8.toString(privKeyRaw, 'hex'),
+    'raw',
+    'pem',
+  )
+  return nodeCrypto.createPrivateKey(privKeyEncoded)
+}
+
+// @NOTE temporary measure, see note on entrywaySession in bsky/src/auth-verifier.ts
+describe('entryway auth', () => {
+  let network: TestNetwork
+  let agent: AtpAgent
+  let sc: SeedClient
+  let alice: string
+  let jwtPrivKey: nodeCrypto.KeyObject
+
+  beforeAll(async () => {
+    const keypair = await crypto.Secp256k1Keypair.create({ exportable: true })
+    jwtPrivKey = await derivePrivKey(keypair)
+    const entrywayJwtPublicKeyHex = ui8.toString(
+      keypair.publicKeyBytes(),
+      'hex',
+    )
+
+    network = await TestNetwork.create({
+      dbPostgresSchema: 'bsky_entryway_auth',
+      bsky: {
+        entrywayJwtPublicKeyHex,
+      },
+    })
+    agent = network.bsky.getClient()
+    sc = network.getSeedClient()
+    await basicSeed(sc)
+    await network.processAll()
+    alice = sc.dids.alice
+  })
+
+  afterAll(async () => {
+    await network.close()
+  })
+
+  it('works', async () => {
+    const signer = new jose.SignJWT({ scope: 'com.atproto.access' })
+      .setSubject(alice)
+      .setIssuedAt()
+      .setExpirationTime('60mins')
+      .setAudience('did:web:fake.server.bsky.network')
+      .setProtectedHeader({
+        typ: 'at+jwt', // https://www.rfc-editor.org/rfc/rfc9068.html
+        alg: 'ES256K',
+      })
+    const token = await signer.sign(jwtPrivKey)
+    const res = await agent.app.bsky.actor.getProfile(
+      { actor: sc.dids.bob },
+      { headers: { authorization: `Bearer ${token}` } },
+    )
+    expect(res.data.did).toEqual(sc.dids.bob)
+    // ensure this request is personalized for alice
+    const followingUri = res.data.viewer?.following
+    assert(followingUri)
+    const parsed = new AtUri(followingUri)
+    expect(parsed.hostname).toEqual(alice)
+  })
+
+  it('does not work on bad scopes', async () => {
+    const signer = new jose.SignJWT({ scope: 'com.atproto.refresh' })
+      .setSubject(alice)
+      .setIssuedAt()
+      .setExpirationTime('60mins')
+      .setAudience('did:web:fake.server.bsky.network')
+      .setProtectedHeader({
+        typ: 'at+jwt', // https://www.rfc-editor.org/rfc/rfc9068.html
+        alg: 'ES256K',
+      })
+    const token = await signer.sign(jwtPrivKey)
+    const attempt = agent.app.bsky.actor.getProfile(
+      { actor: sc.dids.bob },
+      { headers: { authorization: `Bearer ${token}` } },
+    )
+    await expect(attempt).rejects.toThrow('Bad token scope')
+  })
+
+  it('does not work on expired tokens', async () => {
+    const time = Math.floor((Date.now() - 5 * MINUTE) / 1000)
+    const signer = new jose.SignJWT({ scope: 'com.atproto.access' })
+      .setSubject(alice)
+      .setIssuedAt()
+      .setExpirationTime(time)
+      .setAudience('did:web:fake.server.bsky.network')
+      .setProtectedHeader({
+        typ: 'at+jwt', // https://www.rfc-editor.org/rfc/rfc9068.html
+        alg: 'ES256K',
+      })
+    const token = await signer.sign(jwtPrivKey)
+    const attempt = agent.app.bsky.actor.getProfile(
+      { actor: sc.dids.bob },
+      { headers: { authorization: `Bearer ${token}` } },
+    )
+    await expect(attempt).rejects.toThrow('Token has expired')
+  })
+
+  it('does not work on bad auds', async () => {
+    const signer = new jose.SignJWT({ scope: 'com.atproto.access' })
+      .setSubject(alice)
+      .setIssuedAt()
+      .setExpirationTime('60mins')
+      .setAudience('did:web:my.personal.pds.com')
+      .setProtectedHeader({
+        typ: 'at+jwt', // https://www.rfc-editor.org/rfc/rfc9068.html
+        alg: 'ES256K',
+      })
+    const token = await signer.sign(jwtPrivKey)
+    const attempt = agent.app.bsky.actor.getProfile(
+      { actor: sc.dids.bob },
+      { headers: { authorization: `Bearer ${token}` } },
+    )
+    await expect(attempt).rejects.toThrow('Bad token aud')
+  })
+
+  it('does not work with bad signatures', async () => {
+    const fakeKey = await crypto.Secp256k1Keypair.create({ exportable: true })
+    const fakeJwtKey = await derivePrivKey(fakeKey)
+    const signer = new jose.SignJWT({ scope: 'com.atproto.access' })
+      .setSubject(alice)
+      .setIssuedAt()
+      .setExpirationTime('60mins')
+      .setAudience('did:web:my.personal.pds.com')
+      .setProtectedHeader({
+        typ: 'at+jwt', // https://www.rfc-editor.org/rfc/rfc9068.html
+        alg: 'ES256K',
+      })
+    const token = await signer.sign(fakeJwtKey)
+    const attempt = agent.app.bsky.actor.getProfile(
+      { actor: sc.dids.bob },
+      { headers: { authorization: `Bearer ${token}` } },
+    )
+    await expect(attempt).rejects.toThrow('Token could not be verified')
+  })
+
+  it('does not work on flexible aud routes', async () => {
+    const signer = new jose.SignJWT({ scope: 'com.atproto.access' })
+      .setSubject(alice)
+      .setIssuedAt()
+      .setExpirationTime('60mins')
+      .setAudience('did:web:fake.server.bsky.network')
+      .setProtectedHeader({
+        typ: 'at+jwt', // https://www.rfc-editor.org/rfc/rfc9068.html
+        alg: 'ES256K',
+      })
+    const token = await signer.sign(jwtPrivKey)
+    const feedUri = AtUri.make(alice, 'app.bsky.feed.generator', 'fake-feed')
+    const attempt = agent.app.bsky.feed.getFeed(
+      { feed: feedUri.toString() },
+      { headers: { authorization: `Bearer ${token}` } },
+    )
+    await expect(attempt).rejects.toThrow('Malformed token')
+  })
+})