ucanto integration + fixing tests

package.json

@@ -35,6 +35,9 @@
   "author": "Alan Shaw",
   "license": "Apache-2.0 OR MIT",
   "dependencies": {
+    "@ucanto/client": "^9.0.1",
+    "@ucanto/principal": "^8.1.0",
+    "@ucanto/transport": "^9.1.1",
     "@web3-storage/blob-fetcher": "^2.2.0",
     "@web3-storage/gateway-lib": "^5.1.2",
     "dagula": "^8.0.0",
@@ -48,7 +51,6 @@
     "@types/mocha": "^10.0.9",
     "@types/node-fetch": "^2.6.11",
     "@types/sinon": "^17.0.3",
-    "@ucanto/principal": "^8.1.0",
     "@web3-storage/content-claims": "^5.0.0",
     "@web3-storage/public-bucket": "^1.1.0",
     "@web3-storage/upload-client": "^16.1.1",

src/middleware/withEgressTracker.js

@@ -27,9 +27,8 @@ export function withEgressTracker (handler) {
       createByteCountStream((totalBytesServed) => {
         // Non-blocking call to the accounting service to record egress
         if (totalBytesServed > 0) {
-          const { space, dataCid: resource } = ctx
           ctx.waitUntil(
-            ctx.accountingService.record(space, resource, totalBytesServed, new Date().toISOString())
+            ctx.ucantoClient.record(ctx.space, ctx.dataCid, totalBytesServed, new Date())
           )
         }
       })
@@ -44,15 +43,14 @@ export function withEgressTracker (handler) {
 }
 
 /**
- * Creates a TransformStream to count bytes served to the client.
- * It records egress when the stream is finalized without an error.
+ * Creates a TransformStream to count bytes in the response body.
  *
- * @param {(totalBytesServed: number) => void} onClose
+ * @param {(totalBytes: number) => void} onClose
  * @template {Uint8Array} T
  * @returns {TransformStream<T, T>} - The created TransformStream.
  */
 function createByteCountStream (onClose) {
-  let totalBytesServed = 0
+  let totalBytes = 0
 
   return new TransformStream({
     /**
@@ -64,7 +62,7 @@ function createByteCountStream (onClose) {
     async transform (chunk, controller) {
       try {
         controller.enqueue(chunk)
-        totalBytesServed += chunk.byteLength
+        totalBytes += chunk.byteLength
       } catch (error) {
         console.error('Error while counting bytes:', error)
         controller.error(error)
@@ -79,7 +77,7 @@ function createByteCountStream (onClose) {
      * NOTE: The flush function is NOT called in case of a stream error.
      */
     async flush () {
-      onClose(totalBytesServed)
+      onClose(totalBytes)
     }
   })
 }

src/middleware/withEgressTracker.types.ts

@@ -1,12 +1,12 @@
 import { IpfsUrlContext, Environment as MiddlewareEnvironment } from '@web3-storage/gateway-lib'
-import { AccountingService } from './withAccountingService.types.js'
-import { DIDKey, UnknownLink } from '@ucanto/client'
+import { UCantoClient } from './withUcantoClient.types.js'
+import { DIDKey } from '@ucanto/principal/ed25519'
 
 export interface Environment extends MiddlewareEnvironment {
   FF_EGRESS_TRACKER_ENABLED: string
 }
 
 export interface Context extends IpfsUrlContext {
   space: DIDKey
-  accountingService: AccountingService
+  ucantoClient: UCantoClient
 }

src/middleware/withRateLimit.js

@@ -99,7 +99,7 @@ async function isRateLimited (rateLimitAPI, cid) {
  * @param {Environment} env
  * @param {string} authToken
  * @param {RateLimiterContext} ctx
- * @returns {Promise<import('./withAccountingService.types.js').TokenMetadata | null>}
+ * @returns {Promise<import('./withUcantoClient.types.ts').TokenMetadata | null>}
  */
 async function getTokenMetadata (env, authToken, ctx) {
   const cachedValue = await env.AUTH_TOKEN_METADATA.get(authToken)
@@ -109,7 +109,7 @@ async function getTokenMetadata (env, authToken, ctx) {
     return decode(cachedValue)
   }
 
-  const tokenMetadata = findTokenMetadata(authToken)
+  const tokenMetadata = await ctx.ucantoClient.getTokenMetadata(authToken)
   if (tokenMetadata) {
     // NOTE: non-blocking call to the auth token metadata cache
     ctx.waitUntil(env.AUTH_TOKEN_METADATA.put(authToken, encode(tokenMetadata)))
@@ -119,27 +119,17 @@ async function getTokenMetadata (env, authToken, ctx) {
   return null
 }
 
-/**
- * @param {string} authToken
- * @returns {import('./withAccountingService.types.js').TokenMetadata | null}
- */
-function findTokenMetadata (authToken) {
-  // TODO I think this needs to check the content claims service (?) for any claims relevant to this token
-  // TODO do we have a plan for this? need to ask Hannah if the indexing service covers this?
-  return null
-}
-
 /**
  * @param {string} s
- * @returns {import('./withAccountingService.types.js').TokenMetadata}
+ * @returns {import('./withUcantoClient.types.ts').TokenMetadata}
  */
 function decode (s) {
   // TODO should this be dag-json?
   return JSON.parse(s)
 }
 
 /**
- * @param {import('./withAccountingService.types.js').TokenMetadata} m
+ * @param {import('./withUcantoClient.types.ts').TokenMetadata} m
  * @returns {string}
  */
 function encode (m) {

src/middleware/withRateLimit.types.ts

@@ -2,6 +2,7 @@ import { CID } from '@web3-storage/gateway-lib/handlers'
 import { IpfsUrlContext, Environment as MiddlewareEnvironment } from '@web3-storage/gateway-lib'
 import { KVNamespace, RateLimit } from '@cloudflare/workers-types'
 import { RATE_LIMIT_EXCEEDED } from '../constants.js'
+import { UCantoClient } from './withUcantoClient.types.js'
 
 export interface Environment extends MiddlewareEnvironment {
   RATE_LIMITER: RateLimit
@@ -11,6 +12,7 @@ export interface Environment extends MiddlewareEnvironment {
 
 export interface Context extends IpfsUrlContext {
   authToken: string | null
+  ucantoClient: UCantoClient
 }
 
 export type RateLimitExceeded = typeof RATE_LIMIT_EXCEEDED[keyof typeof RATE_LIMIT_EXCEEDED]

src/middleware/withUcantoClient.capabilities.js

@@ -0,0 +1,21 @@
+import { SpaceDID, equalWith } from '@web3-storage/capabilities/utils'
+import { capability, Schema } from '@ucanto/validator'
+
+export const Usage = {
+  /**
+   * Capability can be invoked by an agent to record egress data for a given resource.
+   */
+  record: capability({
+    can: 'usage/record',
+    with: SpaceDID,
+    nb: Schema.struct({
+      /** CID of the resource that was served. */
+      resource: Schema.link(),
+      /** Amount of bytes served. */
+      bytes: Schema.integer().greaterThan(0),
+      /** Timestamp of the event in seconds after Unix epoch. */
+      servedAt: Schema.integer().greaterThan(-1)
+    }),
+    derives: equalWith
+  })
+}

src/middleware/withUcantoClient.js

@@ -0,0 +1,92 @@
+import * as UCantoClient from '@ucanto/client'
+import * as CAR from '@ucanto/transport/car'
+import { SpaceDID } from '@web3-storage/capabilities/utils'
+import { ed25519 } from '@ucanto/principal'
+import { HTTP } from '@ucanto/transport'
+import { Usage } from './withUcantoClient.capabilities.js'
+
+/**
+ * @import { Middleware } from '@web3-storage/gateway-lib'
+ * @typedef {import('./withUcantoClient.types.ts').UcantoClientContext} UcantoClientContext
+ * @typedef {import('./withUcantoClient.types.ts').Environment} Environment
+ */
+
+/**
+ * The UCantoClient handler exposes the methods to invoke capabilities on the Upload API.
+ *
+ * @type {Middleware<UcantoClientContext, UcantoClientContext, Environment>}
+ */
+export function withUcantoClient (handler) {
+  return async (req, env, ctx) => {
+    const ucantoClient = await create(env)
+
+    return handler(req, env, { ...ctx, ucantoClient })
+  }
+}
+
+/**
+ * Creates a UCantoClient instance with the given environment.
+ *
+ * @param {Environment} env
+ * @returns {Promise<import('./withUcantoClient.types.ts').UCantoClient>}
+ */
+async function create (env) {
+  const service = ed25519.Verifier.parse(env.SERVICE_ID)
+  const principal = ed25519.Signer.parse(env.SIGNER_PRINCIPAL_KEY)
+
+  const { connection } = await connectUcantoClient(env.UPLOAD_API_URL, principal)
+
+  return {
+    /**
+     * @param {import('@ucanto/principal/ed25519').DIDKey} space - The Space DID where the content was served
+     * @param {import('@ucanto/principal/ed25519').UnknownLink} resource - The link to the resource that was served
+     * @param {number} bytes - The number of bytes served
+     * @param {Date} servedAt - The timestamp of when the content was served
+     */
+    record: async (space, resource, bytes, servedAt) => {
+      const res = await Usage.record.invoke({
+        issuer: principal,
+        audience: service,
+        with: SpaceDID.from(space),
+        nb: {
+          resource,
+          bytes,
+          servedAt: Math.floor(servedAt.getTime() / 1000)
+        }
+      }).execute(connection)
+
+      if (res.out.error) {
+        console.error('Failed to record egress', res.out.error)
+      }
+    },
+
+    /**
+     * TODO: implement this function
+     *
+     * @param {string} authToken
+     * @returns {Promise<import('./withUcantoClient.types.ts').TokenMetadata | undefined>}
+     */
+    getTokenMetadata: async (authToken) => {
+      // TODO I think this needs to check the content claims service (?) for any claims relevant to this token
+      // TODO do we have a plan for this? need to ask Hannah if the indexing service covers this?
+      return undefined
+    }
+  }
+}
+
+/**
+ * Creates a connection with the UCanto Server at the provided server URL.
+ *
+ * @param {string} serverUrl
+ * @param {import('@ucanto/principal/ed25519').EdSigner} principal
+ *
+ */
+async function connectUcantoClient (serverUrl, principal) {
+  const connection = await UCantoClient.connect({
+    id: principal,
+    codec: CAR.outbound,
+    channel: HTTP.open({ url: new URL(serverUrl) })
+  })
+
+  return { connection }
+}

src/middleware/withAccountingService.types.ts → src/middleware/withUcantoClient.types.ts

@@ -2,18 +2,23 @@ import { Environment as MiddlewareEnvironment, Context as MiddlewareContext } fr
 import { DIDKey, UnknownLink } from '@ucanto/principal/ed25519'
 
 export interface Environment extends MiddlewareEnvironment {
-  //TODO: ucanto signer principal key
+  SERVICE_ID: string
+  SIGNER_PRINCIPAL_KEY: string
+  UPLOAD_API_URL: string
 }
 
-export interface AccountingServiceContext extends MiddlewareContext {
-  accountingService?: AccountingService
+export interface UcantoClientContext extends MiddlewareContext {
+  ucantoClient?: UCantoClient
 }
 
 export interface TokenMetadata {
   locationClaim?: unknown // TODO: figure out the right type to use for this - we probably need it for the private data case to verify auth
   invalid?: boolean
 }
 
-export interface AccountingService {
-  record: (space: DIDKey, resource: UnknownLink, bytes: number, servedAt: string) => Promise<void>
+export interface UCantoClient {
+  record: (space: DIDKey, resource: UnknownLink, bytes: number, servedAt: Date) => Promise<void>
+  getTokenMetadata: (token: string) => Promise<TokenMetadata | undefined>
 }
+
+

test/helpers/builder.js

@@ -19,7 +19,9 @@ export class Builder {
    * @returns {Promise<{ root: import('multiformats').UnknownLink, shards: import('multiformats').Link[]}>}
    */
   async add (input, options = {}) {
-    console.log('Adding ' + (Array.isArray(input) ? `${input.length} file${input.length === 1 ? '' : 's'}` : '1 blob') + '...')
+    if (process.env.DEBUG) {
+      console.log('Adding ' + (Array.isArray(input) ? `${input.length} file${input.length === 1 ? '' : 's'}` : '1 blob') + '...')
+    }
     const unixFsEncoder = Array.isArray(input)
       ? UnixFS.createDirectoryEncoderStream(input)
       : UnixFS.createFileEncoderStream(input)