Armory now supports multiple datastore keys; engine already supported…

… this (#507)
narval-xyz · Aug 16, 2024 · d7b1f0e · d7b1f0e
1 parent 86b1b2a
commit d7b1f0e
Show file tree

Hide file tree

Showing 23 changed files with 308 additions and 90 deletions.
diff --git a/apps/armory/Makefile b/apps/armory/Makefile
@@ -72,7 +72,8 @@ armory/db/create-migration:
 	npx dotenv -e ${ARMORY_PROJECT_DIR}/.env -- \
 		prisma migrate dev \
 			--schema ${ARMORY_DATABASE_SCHEMA} \
-			--name ${NAME}
+			--name ${NAME} \
+			--create-only
 
 # To maintain seed data within their respective modules and then import them
 # into the main seed.ts file for execution, it's necessary to compile the

diff --git a/apps/armory/src/client/__test__/e2e/client.spec.ts b/apps/armory/src/client/__test__/e2e/client.spec.ts
@@ -47,6 +47,7 @@ const mockPolicyEngineServer = (url: string, clientId: string) => {
   }
 
   nock(url).post('/v1/clients').reply(HttpStatus.CREATED, createClientResponse)
+  nock(url).post('/v1/clients/sync').reply(HttpStatus.OK, { success: true })
 }
 
 describe('Client', () => {
@@ -64,6 +65,7 @@ describe('Client', () => {
   const adminApiKey = 'test-admin-api-key'
 
   const entityStorePublicKey = getPublicKey(privateKeyToJwk(generatePrivateKey()))
+  const entityStorePublicKey2 = getPublicKey(privateKeyToJwk(generatePrivateKey()))
 
   const policyStorePublicKey = getPublicKey(privateKeyToJwk(generatePrivateKey()))
 
@@ -117,7 +119,7 @@ describe('Client', () => {
         entity: {
           data: dataStoreSource,
           signature: dataStoreSource,
-          keys: [entityStorePublicKey]
+          keys: [entityStorePublicKey, entityStorePublicKey2] // test w/ 2 keys
         },
         policy: {
           data: dataStoreSource,
@@ -151,15 +153,15 @@ describe('Client', () => {
         updatedAt: actualClient?.updatedAt.toISOString()
       })
 
-      expect(actualClient?.dataStore.entityPublicKey).toEqual(createClientPayload.dataStore.entity.keys[0])
-      expect(actualClient?.dataStore.policyPublicKey).toEqual(createClientPayload.dataStore.policy.keys[0])
+      expect(actualClient?.dataStore.entityPublicKeys).toEqual(createClientPayload.dataStore.entity.keys)
+      expect(actualClient?.dataStore.policyPublicKeys).toEqual(createClientPayload.dataStore.policy.keys)
 
       expect(status).toEqual(HttpStatus.CREATED)
     })
 
     it('creates a new client with given policy engines', async () => {
       mockPolicyEngineServer(policyEngineNodeUrl, clientId)
-      mockPolicyEngineServer(policyEngineNodeUrl, clientId)
+      mockPolicyEngineServer(policyEngineNodeUrl, clientId) // second mock is required because it gets called twice since we duplicate the url; test will fail if you delete this.
 
       const createClientWithGivenPolicyEngine: CreateClientRequestDto = {
         ...createClientPayload,

diff --git a/apps/armory/src/client/core/service/client.service.ts b/apps/armory/src/client/core/service/client.service.ts
@@ -46,8 +46,8 @@ export class ClientService {
       dataSecret: null,
       name: input.name,
       dataStore: {
-        entityPublicKey: entityDataStore.keys[0],
-        policyPublicKey: policyDataStore.keys[0]
+        entityPublicKeys: entityDataStore.keys,
+        policyPublicKeys: policyDataStore.keys
       },
       policyEngine: { nodes: [] },
       createdAt: now,
@@ -78,6 +78,9 @@ export class ClientService {
 
     const createdClient = await this.clientRepository.save(client)
 
+    // Trigger a cluster data sync, since we likely _just_ created the datastore.
+    await this.clusterService.sync(clientId)
+
     return {
       ...this.buildPublicClient({
         client: createdClient,

diff --git a/apps/armory/src/client/core/type/client.type.ts b/apps/armory/src/client/core/type/client.type.ts
@@ -33,8 +33,8 @@ export const Client = z.object({
   createdAt: z.date(),
   updatedAt: z.date(),
   dataStore: z.object({
-    entityPublicKey: jwkSchema,
-    policyPublicKey: jwkSchema
+    entityPublicKeys: z.array(jwkSchema).min(1),
+    policyPublicKeys: z.array(jwkSchema).min(1)
   }),
   policyEngine: z.object({
     nodes: z.array(PolicyEngineNode)

diff --git a/apps/armory/src/client/persistence/repository/client.repository.ts b/apps/armory/src/client/persistence/repository/client.repository.ts
@@ -1,6 +1,6 @@
-import { publicKeySchema } from '@narval/signature'
+import { PublicKey, publicKeySchema } from '@narval/signature'
 import { Injectable } from '@nestjs/common'
-import { Client as Model, Prisma } from '@prisma/client/armory'
+import { DataStoreKey, Client as Model, Prisma } from '@prisma/client/armory'
 import { PrismaService } from '../../../shared/module/persistence/service/prisma.service'
 import { Client } from '../../core/type/client.type'
 
@@ -10,7 +10,14 @@ export class ClientRepository {
 
   async findById(id: string): Promise<Client | null> {
     const model = await this.prismaService.client.findUnique({
-      where: { id }
+      where: { id },
+      include: {
+        dataStoreKeys: {
+          where: {
+            deletedAt: null
+          }
+        }
+      }
     })
 
     if (model) {
@@ -21,31 +28,54 @@ export class ClientRepository {
   }
 
   async save(client: Client): Promise<Client> {
+    const dataStoreKeys = this.encodeDataStoreKeys(client)
+
+    const clientData = this.encode(client)
+
     await this.prismaService.client.create({
-      data: this.encode(client)
+      data: {
+        ...clientData,
+        dataStoreKeys: {
+          createMany: {
+            data: dataStoreKeys
+          }
+        }
+      }
     })
 
     return client
   }
 
-  private decode(model: Model): Client {
+  private decode(
+    model: Model & {
+      dataStoreKeys: DataStoreKey[]
+    }
+  ): Client {
     return {
       id: model.id,
       clientSecret: model.clientSecret,
       dataSecret: model.dataSecret,
       name: model.name,
       createdAt: model.createdAt,
       updatedAt: model.updatedAt,
-      dataStore: {
-        entityPublicKey: publicKeySchema.parse(model.entityPublicKey),
-        policyPublicKey: publicKeySchema.parse(model.policyPublicKey)
-      },
+      dataStore: this.decodeDataStoreKeys(model.dataStoreKeys),
       policyEngine: {
         nodes: []
       }
     }
   }
 
+  private decodeDataStoreKeys(model: DataStoreKey[]): { entityPublicKeys: PublicKey[]; policyPublicKeys: PublicKey[] } {
+    return {
+      entityPublicKeys: model
+        .filter((key) => key.storeType === 'entity')
+        .map((key) => publicKeySchema.parse(key.publicKey)),
+      policyPublicKeys: model
+        .filter((key) => key.storeType === 'policy')
+        .map((key) => publicKeySchema.parse(key.publicKey))
+    }
+  }
+
   private encode(client: Client) {
     return {
       id: client.id,
@@ -54,9 +84,20 @@ export class ClientRepository {
       name: client.name,
       createdAt: client.createdAt,
       updatedAt: client.updatedAt,
-      enginePublicKey: publicKeySchema.parse(client.dataStore.entityPublicKey) as Prisma.InputJsonValue,
-      entityPublicKey: publicKeySchema.parse(client.dataStore.entityPublicKey) as Prisma.InputJsonValue,
-      policyPublicKey: publicKeySchema.parse(client.dataStore.policyPublicKey) as Prisma.InputJsonValue
+      enginePublicKey: publicKeySchema.parse(client.policyEngine.nodes[0].publicKey) as Prisma.InputJsonValue
     }
   }
+
+  private encodeDataStoreKeys(client: Client) {
+    const entityDataStoreKeys = client.dataStore.entityPublicKeys.map((key) => ({
+      storeType: 'entity',
+      publicKey: publicKeySchema.parse(key)
+    }))
+    const policyDataStoreKeys = client.dataStore.policyPublicKeys.map((key) => ({
+      storeType: 'policy',
+      publicKey: publicKeySchema.parse(key)
+    }))
+    const dataStoreKeys = [...entityDataStoreKeys, ...policyDataStoreKeys]
+    return dataStoreKeys
+  }
 }
diff --git a/apps/armory/src/managed-data-store/__test__/e2e/data-store.spec.ts b/apps/armory/src/managed-data-store/__test__/e2e/data-store.spec.ts
@@ -107,8 +107,8 @@ describe('Data Store', () => {
       createdAt: new Date(),
       updatedAt: new Date(),
       dataStore: {
-        entityPublicKey: getPublicKey(dataStorePrivateKey),
-        policyPublicKey: getPublicKey(dataStorePrivateKey)
+        entityPublicKeys: [getPublicKey(dataStorePrivateKey)],
+        policyPublicKeys: [getPublicKey(dataStorePrivateKey)]
       },
       policyEngine: {
         nodes: [

diff --git a/apps/armory/src/managed-data-store/core/__test__/unit/signature.service.spec.ts b/apps/armory/src/managed-data-store/core/__test__/unit/signature.service.spec.ts
@@ -21,7 +21,7 @@ describe(SignatureService.name, () => {
 
     await expect(() =>
       signatureService.verifySignature({
-        pubKey: jwk,
+        keys: [jwk],
         payload: { signature, data: FIXTURE.ENTITIES },
         date: new Date('2024-01-01')
       })
@@ -42,19 +42,19 @@ describe(SignatureService.name, () => {
 
     await expect(() =>
       signatureService.verifySignature({
-        pubKey: jwk,
+        keys: [jwk],
         payload: { signature, data: FIXTURE.ENTITIES },
         date: new Date('2023-01-01')
       })
     ).rejects.toThrow(ApplicationException)
 
     await expect(() =>
       signatureService.verifySignature({
-        pubKey: jwk,
+        keys: [jwk],
         payload: { signature, data: FIXTURE.ENTITIES },
         date: new Date('2023-01-01')
       })
-    ).rejects.toThrow('Invalid signature')
+    ).rejects.toThrow('Signature not valid for keys')
   })
 
   it('returns true if the payload iat is more recent than the db createdAt date', async () => {
@@ -68,7 +68,32 @@ describe(SignatureService.name, () => {
     const signature = await signJwt(payload, jwk, { alg: SigningAlg.EIP191 }, buildSignerEip191(DATA_STORE_PRIVATE_KEY))
 
     const result = await signatureService.verifySignature({
-      pubKey: jwk,
+      keys: [jwk],
+      payload: { signature, data: FIXTURE.ENTITIES },
+      date: new Date('2023-01-01')
+    })
+
+    expect(result).toEqual(true)
+  })
+
+  it('verifies a signature with multiple keys when kid does not match', async () => {
+    const payload: Payload = {
+      data: hash(FIXTURE.ENTITIES),
+      sub: 'test-root-user-uid',
+      iss: 'https://armory.narval.xyz',
+      iat: Math.floor(new Date('2024-01-01').getTime() / 1000) // in seconds
+    }
+
+    const secondaryDataKey = generatePrivateKey()
+    const secondJwk = secp256k1PrivateKeyToJwk(secondaryDataKey)
+    secondJwk.kid = 'secondary-data-key' // overwrite the kid to ensure no matches.
+
+    const primaryJwk = { ...jwk, kid: 'primary-data-key' } // overwrite the kid to ensure no matches
+
+    const signature = await signJwt(payload, jwk, { alg: SigningAlg.EIP191 }, buildSignerEip191(DATA_STORE_PRIVATE_KEY))
+
+    const result = await signatureService.verifySignature({
+      keys: [secondJwk, primaryJwk],
       payload: { signature, data: FIXTURE.ENTITIES },
       date: new Date('2023-01-01')
     })

diff --git a/apps/armory/src/managed-data-store/core/service/entity-data-store.service.ts b/apps/armory/src/managed-data-store/core/service/entity-data-store.service.ts
@@ -1,5 +1,4 @@
 import { EntityStore } from '@narval/policy-engine-shared'
-import { publicKeySchema } from '@narval/signature'
 import { HttpStatus, Injectable, NotFoundException } from '@nestjs/common'
 import { ClientService } from '../../../client/core/service/client.service'
 import { ClusterService } from '../../../policy-engine/core/service/cluster.service'
@@ -10,15 +9,15 @@ import { SignatureService } from './signature.service'
 @Injectable()
 export class EntityDataStoreService extends SignatureService {
   constructor(
-    private entitydataStoreRepository: EntityDataStoreRepository,
+    private entityDataStoreRepository: EntityDataStoreRepository,
     private clientService: ClientService,
     private clusterService: ClusterService
   ) {
     super()
   }
 
   async getEntities(clientId: string): Promise<EntityStore | null> {
-    const entityStore = await this.entitydataStoreRepository.getLatestDataStore(clientId)
+    const entityStore = await this.entityDataStoreRepository.getLatestDataStore(clientId)
 
     return entityStore ? EntityStore.parse(entityStore.data) : null
   }
@@ -33,15 +32,15 @@ export class EntityDataStoreService extends SignatureService {
       })
     }
 
-    const latestDataStore = await this.entitydataStoreRepository.getLatestDataStore(clientId)
+    const latestDataStore = await this.entityDataStoreRepository.getLatestDataStore(clientId)
 
     await this.verifySignature({
       payload,
-      pubKey: publicKeySchema.parse(client.dataStore.entityPublicKey),
+      keys: client.dataStore.entityPublicKeys,
       date: latestDataStore?.createdAt
     })
 
-    const { data, version } = await this.entitydataStoreRepository.setDataStore(clientId, {
+    const { data, version } = await this.entityDataStoreRepository.setDataStore(clientId, {
       version: latestDataStore?.version ? latestDataStore.version + 1 : 1,
       data: EntityStore.parse(payload)
     })

diff --git a/apps/armory/src/managed-data-store/core/service/policy-data-store.service.ts b/apps/armory/src/managed-data-store/core/service/policy-data-store.service.ts
@@ -1,5 +1,4 @@
 import { PolicyStore } from '@narval/policy-engine-shared'
-import { publicKeySchema } from '@narval/signature'
 import { HttpStatus, Injectable, NotFoundException } from '@nestjs/common'
 import { ClientService } from '../../../client/core/service/client.service'
 import { ClusterService } from '../../../policy-engine/core/service/cluster.service'
@@ -36,7 +35,7 @@ export class PolicyDataStoreService extends SignatureService {
 
     await this.verifySignature({
       payload,
-      pubKey: publicKeySchema.parse(client.dataStore.policyPublicKey),
+      keys: client.dataStore.policyPublicKeys,
       date: latestDataStore?.createdAt
     })