Skip to content

Commit fbcaa83

Browse files
fforbeckfforbeck
committed
refactor: accounting service
1 parent 25982f4 commit fbcaa83

File tree

11 files changed

+122
-69
lines changed

11 files changed

+122
-69
lines changed

‎src/bindings.d.ts‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,13 @@ import { Environment as CarBlockEnvironment } from './middleware/withCarBlockHan
44
import { Environment as ContentClaimsDagulaEnvironment } from './middleware/withCarBlockHandler.types.ts'
55
import { Environment as EgressTrackerEnvironment } from './middleware/withEgressTracker.types.ts'
66
import { UnknownLink } from 'multiformats'
7+
import { DIDKey } from '@ucanto/principal/ed25519'
8+
79
export interface Environment
810
extends CarBlockEnvironment,
911
RateLimiterEnvironment,
1012
ContentClaimsDagulaEnvironment,
1113
EgressTrackerEnvironment {
1214
VERSION: string
1315
CONTENT_CLAIMS_SERVICE_URL?: string
14-
ACCOUNTING_SERVICE_URL: string
15-
}
16-
17-
export interface AccountingService {
18-
record: (resource: UnknownLink, bytes: number, servedAt: string) => Promise<void>
19-
getTokenMetadata: (token: string) => Promise<TokenMetadata | null>
20-
}
21-
22-
export interface Accounting {
23-
create: ({ serviceURL }: { serviceURL: string }) => AccountingService
2416
}

‎src/middleware/withAccountingService.js‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/**
2+
* @import { Middleware } from '@web3-storage/gateway-lib'
3+
* @typedef {import('./withAccountingService.types.ts').AccountingServiceContext} AccountingServiceContext
4+
* @typedef {import('./withAccountingService.types.ts').Environment} Environment
5+
*/
6+
7+
/**
8+
* The accounting service handler exposes the method `record` to record the egress bytes for a given SpaceDID, Content CID, and servedAt timestamp.
9+
*
10+
* @type {Middleware<AccountingServiceContext, AccountingServiceContext, Environment>}
11+
*/
12+
export function withAccountingService (handler) {
13+
return async (req, env, ctx) => {
14+
const accountingService = create(env, ctx)
15+
16+
return handler(req, env, { ...ctx, accountingService })
17+
}
18+
}
19+
20+
/**
21+
* @param {Environment} env
22+
* @param {AccountingServiceContext} ctx
23+
* @returns {import('./withAccountingService.types.ts').AccountingService}
24+
*/
25+
function create (env, ctx) {
26+
return {
27+
/**
28+
* @param {import('@ucanto/principal/ed25519').DIDKey} space - The Space DID where the content was served
29+
* @param {import('@ucanto/principal/ed25519').UnknownLink} resource - The link to the resource that was served
30+
* @param {number} bytes - The number of bytes served
31+
* @param {string} servedAt - The timestamp of when the content was served
32+
*/
33+
record: async (space, resource, bytes, servedAt) => {
34+
console.log(`Record egress: ${space}, ${resource}, ${bytes}, ${servedAt}`)
35+
}
36+
}
37+
}

‎src/middleware/withAccountingService.types.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import { Environment as MiddlewareEnvironment, Context as MiddlewareContext } from '@web3-storage/gateway-lib'
2+
import { DIDKey, UnknownLink } from '@ucanto/principal/ed25519'
3+
4+
export interface Environment extends MiddlewareEnvironment {
5+
//TODO: ucanto signer principal key
6+
}
7+
8+
export interface AccountingServiceContext extends MiddlewareContext {
9+
accountingService?: AccountingService
10+
}
11+
12+
export interface TokenMetadata {
13+
locationClaim?: unknown // TODO: figure out the right type to use for this - we probably need it for the private data case to verify auth
14+
invalid?: boolean
15+
}
16+
17+
export interface AccountingService {
18+
record: (space: DIDKey, resource: UnknownLink, bytes: number, servedAt: string) => Promise<void>
19+
}

‎src/middleware/withEgressTracker.js‎

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
1-
import { Accounting } from '../services/accounting.js'
2-
31
/**
4-
* @import { Context, IpfsUrlContext, Middleware } from '@web3-storage/gateway-lib'
2+
* @import { Middleware } from '@web3-storage/gateway-lib'
53
* @import { Environment } from './withEgressTracker.types.js'
6-
* @import { AccountingService } from '../bindings.js'
7-
* @typedef {IpfsUrlContext & { ACCOUNTING_SERVICE?: AccountingService }} EgressTrackerContext
4+
* @typedef {import('./withEgressTracker.types.js').Context} EgressTrackerContext
85
*/
96

107
/**
@@ -26,17 +23,13 @@ export function withEgressTracker (handler) {
2623
return response
2724
}
2825

29-
const { dataCid } = ctx
30-
const accounting = ctx.ACCOUNTING_SERVICE ?? Accounting.create({
31-
serviceURL: env.ACCOUNTING_SERVICE_URL
32-
})
33-
3426
const responseBody = response.body.pipeThrough(
3527
createByteCountStream((totalBytesServed) => {
3628
// Non-blocking call to the accounting service to record egress
3729
if (totalBytesServed > 0) {
30+
const { space, dataCid: resource } = ctx
3831
ctx.waitUntil(
39-
accounting.record(dataCid, totalBytesServed, new Date().toISOString())
32+
ctx.accountingService.record(space, resource, totalBytesServed, new Date().toISOString())
4033
)
4134
}
4235
})

‎src/middleware/withEgressTracker.types.ts‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
1-
import { Environment as MiddlewareEnvironment } from '@web3-storage/gateway-lib'
1+
import { IpfsUrlContext, Environment as MiddlewareEnvironment } from '@web3-storage/gateway-lib'
2+
import { AccountingService } from './withAccountingService.types.js'
3+
import { DIDKey, UnknownLink } from '@ucanto/client'
24

35
export interface Environment extends MiddlewareEnvironment {
4-
ACCOUNTING_SERVICE_URL: string
56
FF_EGRESS_TRACKER_ENABLED: string
67
}
8+
9+
export interface Context extends IpfsUrlContext {
10+
space: DIDKey
11+
accountingService: AccountingService
12+
}

‎src/middleware/withRateLimit.js‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,16 @@
11
import { HttpError } from '@web3-storage/gateway-lib/util'
22
import { RATE_LIMIT_EXCEEDED } from '../constants.js'
3-
import { Accounting } from '../services/accounting.js'
43

54
/**
65
* @import { Middleware } from '@web3-storage/gateway-lib'
76
* @import { R2Bucket, KVNamespace, RateLimit } from '@cloudflare/workers-types'
87
* @import {
98
* Environment,
109
* Context,
11-
* TokenMetadata,
1210
* RateLimitService,
1311
* RateLimitExceeded
1412
* } from './withRateLimit.types.js'
15-
* @typedef {Context & { ACCOUNTING_SERVICE?: import('../bindings.js').AccountingService }} RateLimiterContext
13+
* @typedef {Context} RateLimiterContext
1614
*/
1715

1816
/**
@@ -101,7 +99,7 @@ async function isRateLimited (rateLimitAPI, cid) {
10199
* @param {Environment} env
102100
* @param {string} authToken
103101
* @param {RateLimiterContext} ctx
104-
* @returns {Promise<TokenMetadata | null>}
102+
* @returns {Promise<import('./withAccountingService.types.js').TokenMetadata | null>}
105103
*/
106104
async function getTokenMetadata (env, authToken, ctx) {
107105
const cachedValue = await env.AUTH_TOKEN_METADATA.get(authToken)
@@ -111,8 +109,7 @@ async function getTokenMetadata (env, authToken, ctx) {
111109
return decode(cachedValue)
112110
}
113111

114-
const accounting = ctx.ACCOUNTING_SERVICE ?? Accounting.create({ serviceURL: env.ACCOUNTING_SERVICE_URL })
115-
const tokenMetadata = await accounting.getTokenMetadata(authToken)
112+
const tokenMetadata = findTokenMetadata(authToken)
116113
if (tokenMetadata) {
117114
// NOTE: non-blocking call to the auth token metadata cache
118115
ctx.waitUntil(env.AUTH_TOKEN_METADATA.put(authToken, encode(tokenMetadata)))
@@ -122,17 +119,27 @@ async function getTokenMetadata (env, authToken, ctx) {
122119
return null
123120
}
124121

122+
/**
123+
* @param {string} authToken
124+
* @returns {import('./withAccountingService.types.js').TokenMetadata | null}
125+
*/
126+
function findTokenMetadata (authToken) {
127+
// TODO I think this needs to check the content claims service (?) for any claims relevant to this token
128+
// TODO do we have a plan for this? need to ask Hannah if the indexing service covers this?
129+
return null
130+
}
131+
125132
/**
126133
* @param {string} s
127-
* @returns {TokenMetadata}
134+
* @returns {import('./withAccountingService.types.js').TokenMetadata}
128135
*/
129136
function decode (s) {
130137
// TODO should this be dag-json?
131138
return JSON.parse(s)
132139
}
133140

134141
/**
135-
* @param {TokenMetadata} m
142+
* @param {import('./withAccountingService.types.js').TokenMetadata} m
136143
* @returns {string}
137144
*/
138145
function encode (m) {

‎src/middleware/withRateLimit.types.ts‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import { KVNamespace, RateLimit } from '@cloudflare/workers-types'
44
import { RATE_LIMIT_EXCEEDED } from '../constants.js'
55

66
export interface Environment extends MiddlewareEnvironment {
7-
ACCOUNTING_SERVICE_URL: string
87
RATE_LIMITER: RateLimit
98
AUTH_TOKEN_METADATA: KVNamespace
109
FF_RATE_LIMITER_ENABLED: string
@@ -14,11 +13,6 @@ export interface Context extends IpfsUrlContext {
1413
authToken: string | null
1514
}
1615

17-
export interface TokenMetadata {
18-
locationClaim?: unknown // TODO: figure out the right type to use for this - we probably need it for the private data case to verify auth
19-
invalid?: boolean
20-
}
21-
2216
export type RateLimitExceeded = typeof RATE_LIMIT_EXCEEDED[keyof typeof RATE_LIMIT_EXCEEDED]
2317

2418
export interface RateLimitService {

‎src/services/accounting.js‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,7 @@ export const Accounting = {
55
create: ({ serviceURL }) => ({
66
record: async (cid, bytes, servedAt) => {
77
console.log(`using ${serviceURL} to record egress for ${cid} with total bytes: ${bytes} and servedAt: ${servedAt}`)
8-
},
9-
10-
getTokenMetadata: async () => {
11-
// TODO I think this needs to check the content claims service (?) for any claims relevant to this token
12-
// TODO do we have a plan for this? need to ask Hannah if the indexing service covers this?
13-
return null
148
}
9+
1510
})
1611
}

‎test/unit/middleware/withEgressTracker.spec.js‎

Lines changed: 36 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -30,41 +30,38 @@ const env =
3030
/** @satisfies {import('../../../src/middleware/withEgressTracker.types.js').Environment} */
3131
({
3232
DEBUG: 'true',
33-
ACCOUNTING_SERVICE_URL: 'http://example.com',
3433
FF_EGRESS_TRACKER_ENABLED: 'true'
3534
})
3635

3736
const accountingRecordMethodStub = sinon.stub()
3837
.returns(
39-
/** @type {import('../../../src/bindings.js').AccountingService['record']} */
40-
async (cid, bytes, servedAt) => {
41-
console.log(`[mock] record called with cid: ${cid}, bytes: ${bytes}, servedAt: ${servedAt}`)
38+
/** @type {import('../../../src/middleware/withAccountingService.types.js').AccountingService['record']} */
39+
async (space, resource, bytes, servedAt) => {
40+
console.log(`[mock] record called with space: ${space}, resource: ${resource}, bytes: ${bytes}, servedAt: ${servedAt}`)
4241
})
4342

4443
/**
4544
* Mock implementation of the AccountingService.
4645
*
47-
* @param {Object} options
48-
* @param {string} options.serviceURL - The URL of the accounting service.
49-
* @returns {import('../../../src/bindings.js').AccountingService}
46+
* @returns {import('../../../src/middleware/withAccountingService.types.js').AccountingService}
5047
*/
51-
const AccountingService = ({ serviceURL }) => {
52-
console.log(`[mock] Accounting.create called with serviceURL: ${serviceURL}`)
48+
const AccountingService = () => {
49+
console.log('[mock] Accounting.create called')
5350

5451
return {
55-
record: accountingRecordMethodStub,
56-
getTokenMetadata: sinon.stub().resolves(undefined)
52+
record: accountingRecordMethodStub
5753
}
5854
}
5955

6056
const ctx =
61-
/** @satisfies {import('../../../src/middleware/withEgressTracker.js').EgressTrackerContext} */
57+
/** @satisfies {import('../../../src/middleware/withEgressTracker.types.js').Context} */
6258
({
59+
space: 'did:key:z6MkknBAHEGCWvBzAi4amdH5FXEXrdKoWF1UJuvc8Psm2Mda',
6360
dataCid: CID.parse('bafybeibv7vzycdcnydl5n5lbws6lul2omkm6a6b5wmqt77sicrwnhesy7y'),
6461
waitUntil: sinon.stub().returns(undefined),
6562
path: '',
6663
searchParams: new URLSearchParams(),
67-
ACCOUNTING_SERVICE: AccountingService({ serviceURL: env.ACCOUNTING_SERVICE_URL })
64+
accountingService: AccountingService()
6865
})
6966

7067
describe('withEgressTracker', async () => {
@@ -119,8 +116,9 @@ describe('withEgressTracker', async () => {
119116
expect(response.status).to.equal(200)
120117
expect(responseBody).to.equal('Hello, world!')
121118
expect(accountingRecordMethodStub.calledOnce, 'record should be called once').to.be.true
122-
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the cid').to.equal(ctx.dataCid)
123-
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the total bytes').to.equal(totalBytes)
119+
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the space').to.equal(ctx.space)
120+
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the cid').to.equal(ctx.dataCid)
121+
expect(accountingRecordMethodStub.args[0][2], 'third argument should be the total bytes').to.equal(totalBytes)
124122
}).timeout(10_000)
125123

126124
it('should record egress for a large file', async () => {
@@ -142,8 +140,9 @@ describe('withEgressTracker', async () => {
142140

143141
expect(response.status).to.equal(200)
144142
expect(accountingRecordMethodStub.calledOnce, 'record should be called once').to.be.true
145-
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the cid').to.equal(ctx.dataCid)
146-
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the total bytes').to.equal(totalBytes)
143+
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the space').to.equal(ctx.space)
144+
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the cid').to.equal(ctx.dataCid)
145+
expect(accountingRecordMethodStub.args[0][2], 'third argument should be the total bytes').to.equal(totalBytes)
147146
})
148147

149148
it('should correctly track egress for responses with chunked transfer encoding', async () => {
@@ -169,7 +168,9 @@ describe('withEgressTracker', async () => {
169168
expect(response.status).to.equal(200)
170169
expect(responseBody).to.equal('Hello, world!')
171170
expect(accountingRecordMethodStub.calledOnce, 'record should be called once').to.be.true
172-
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the total bytes').to.equal(totalBytes)
171+
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the space').to.equal(ctx.space)
172+
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the cid').to.equal(ctx.dataCid)
173+
expect(accountingRecordMethodStub.args[0][2], 'third argument should be the total bytes').to.equal(totalBytes)
173174
})
174175

175176
it('should record egress bytes for a CAR file request', async () => {
@@ -216,7 +217,9 @@ describe('withEgressTracker', async () => {
216217

217218
// expect(blocks[0].bytes).to.deep.equal(carBytes) - FIXME (fforbeck): how to get the correct byte count?
218219
expect(accountingRecordMethodStub.calledOnce, 'record should be called once').to.be.true
219-
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the total bytes').to.equal(expectedTotalBytes)
220+
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the space').to.equal(ctx.space)
221+
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the cid').to.equal(ctx.dataCid)
222+
expect(accountingRecordMethodStub.args[0][2], 'third argument should be the total bytes').to.equal(expectedTotalBytes)
220223
})
221224

222225
it('should correctly track egress for delayed responses', async () => {
@@ -242,7 +245,9 @@ describe('withEgressTracker', async () => {
242245
expect(response.status).to.equal(200)
243246
expect(responseBody).to.equal('Delayed response content')
244247
expect(accountingRecordMethodStub.calledOnce, 'record should be called once').to.be.true
245-
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the total bytes').to.equal(totalBytes)
248+
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the space').to.equal(ctx.space)
249+
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the cid').to.equal(ctx.dataCid)
250+
expect(accountingRecordMethodStub.args[0][2], 'third argument should be the total bytes').to.equal(totalBytes)
246251
}).timeout(5000)
247252
})
248253

@@ -330,8 +335,13 @@ describe('withEgressTracker', async () => {
330335
expect(responseBody2).to.equal('Goodbye, world!')
331336

332337
expect(accountingRecordMethodStub.calledTwice, 'record should be called twice').to.be.true
333-
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the total bytes for first request').to.equal(totalBytes1)
334-
expect(accountingRecordMethodStub.args[1][1], 'second argument should be the total bytes for second request').to.equal(totalBytes2)
338+
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the space').to.equal(ctx.space)
339+
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the cid for first request').to.equal(ctx.dataCid)
340+
expect(accountingRecordMethodStub.args[0][2], 'third argument should be the total bytes for first request').to.equal(totalBytes1)
341+
342+
expect(accountingRecordMethodStub.args[1][0], 'first argument should be the space').to.equal(ctx.space)
343+
expect(accountingRecordMethodStub.args[1][1], 'second argument should be the cid for second request').to.equal(ctx.dataCid)
344+
expect(accountingRecordMethodStub.args[1][2], 'third argument should be the total bytes for second request').to.equal(totalBytes2)
335345
}).timeout(10_000)
336346
})
337347

@@ -356,7 +366,9 @@ describe('withEgressTracker', async () => {
356366
expect(response.status).to.equal(200)
357367
expect(responseBody).to.deep.equal({ message: 'Hello, JSON!' })
358368
expect(accountingRecordMethodStub.calledOnce, 'record should be called once').to.be.true
359-
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the total bytes').to.equal(totalBytes)
369+
expect(accountingRecordMethodStub.args[0][0], 'first argument should be the space').to.equal(ctx.space)
370+
expect(accountingRecordMethodStub.args[0][1], 'second argument should be the cid').to.equal(ctx.dataCid)
371+
expect(accountingRecordMethodStub.args[0][2], 'third argument should be the total bytes').to.equal(totalBytes)
360372
}).timeout(10_000)
361373
})
362374

@@ -449,7 +461,7 @@ describe('withEgressTracker', async () => {
449461
const request = await createRequest()
450462

451463
// Simulate an error in the accounting service record method
452-
ctx.ACCOUNTING_SERVICE.record = sinon.stub().rejects(new Error('Accounting service error'))
464+
ctx.accountingService.record = sinon.stub().rejects(new Error('Accounting service error'))
453465

454466
const response = await handler(request, env, ctx)
455467
const responseBody = await response.text()

‎test/unit/middleware/withRateLimit.spec.js‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@ const env =
4040
/** @satisfies {Environment} */
4141
({
4242
DEBUG: 'false',
43-
ACCOUNTING_SERVICE_URL: 'http://example.com/accounting-service',
4443
RATE_LIMITER: {
4544
limit: strictStub(sandbox, 'limit')
4645
},

0 commit comments

Comments
 (0)