Skip to content

Commit

Permalink
refactor(askar)!: short-lived askar sessions (#1743)
Browse files Browse the repository at this point in the history 
Signed-off-by: Timo Glastra <[email protected]>
  • Loading branch information
@TimoGlastra
TimoGlastra authored Feb 7, 2024
1 parent 81ff63c commit 2cb9ba8
Show file tree
Hide file tree
Showing 9 changed files with 225 additions and 141 deletions.
24 changes: 12 additions & 12 deletions packages/askar/src/storage/AskarStorageService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,16 @@ export class AskarStorageService<T extends BaseRecord> implements StorageService
/** @inheritDoc */
public async save(agentContext: AgentContext, record: T) {
assertAskarWallet(agentContext.wallet)
const session = agentContext.wallet.session

record.updatedAt = new Date()

const value = JsonTransformer.serialize(record)
const tags = transformFromRecordTagValues(record.getTags()) as Record<string, string>

try {
await session.insert({ category: record.type, name: record.id, value, tags })
await agentContext.wallet.withSession((session) =>
session.insert({ category: record.type, name: record.id, value, tags })
)
} catch (error) {
if (isAskarError(error, AskarErrorCode.Duplicate)) {
throw new RecordDuplicateError(`Record with id ${record.id} already exists`, { recordType: record.type })
Expand All @@ -34,15 +35,16 @@ export class AskarStorageService<T extends BaseRecord> implements StorageService
/** @inheritDoc */
public async update(agentContext: AgentContext, record: T): Promise<void> {
assertAskarWallet(agentContext.wallet)
const session = agentContext.wallet.session

record.updatedAt = new Date()

const value = JsonTransformer.serialize(record)
const tags = transformFromRecordTagValues(record.getTags()) as Record<string, string>

try {
await session.replace({ category: record.type, name: record.id, value, tags })
await agentContext.wallet.withSession((session) =>
session.replace({ category: record.type, name: record.id, value, tags })
)
} catch (error) {
if (isAskarError(error, AskarErrorCode.NotFound)) {
throw new RecordNotFoundError(`record with id ${record.id} not found.`, {
Expand All @@ -58,10 +60,9 @@ export class AskarStorageService<T extends BaseRecord> implements StorageService
/** @inheritDoc */
public async delete(agentContext: AgentContext, record: T) {
assertAskarWallet(agentContext.wallet)
const session = agentContext.wallet.session

try {
await session.remove({ category: record.type, name: record.id })
await agentContext.wallet.withSession((session) => session.remove({ category: record.type, name: record.id }))
} catch (error) {
if (isAskarError(error, AskarErrorCode.NotFound)) {
throw new RecordNotFoundError(`record with id ${record.id} not found.`, {
Expand All @@ -80,10 +81,9 @@ export class AskarStorageService<T extends BaseRecord> implements StorageService
id: string
): Promise<void> {
assertAskarWallet(agentContext.wallet)
const session = agentContext.wallet.session

try {
await session.remove({ category: recordClass.type, name: id })
await agentContext.wallet.withSession((session) => session.remove({ category: recordClass.type, name: id }))
} catch (error) {
if (isAskarError(error, AskarErrorCode.NotFound)) {
throw new RecordNotFoundError(`record with id ${id} not found.`, {
Expand All @@ -98,10 +98,11 @@ export class AskarStorageService<T extends BaseRecord> implements StorageService
/** @inheritDoc */
public async getById(agentContext: AgentContext, recordClass: BaseRecordConstructor<T>, id: string): Promise<T> {
assertAskarWallet(agentContext.wallet)
const session = agentContext.wallet.session

try {
const record = await session.fetch({ category: recordClass.type, name: id })
const record = await agentContext.wallet.withSession((session) =>
session.fetch({ category: recordClass.type, name: id })
)
if (!record) {
throw new RecordNotFoundError(`record with id ${id} not found.`, {
recordType: recordClass.type,
Expand All @@ -117,9 +118,8 @@ export class AskarStorageService<T extends BaseRecord> implements StorageService
/** @inheritDoc */
public async getAll(agentContext: AgentContext, recordClass: BaseRecordConstructor<T>): Promise<T[]> {
assertAskarWallet(agentContext.wallet)
const session = agentContext.wallet.session

const records = await session.fetchAll({ category: recordClass.type })
const records = await agentContext.wallet.withSession((session) => session.fetchAll({ category: recordClass.type }))

const instances = []
for (const record of records) {
Expand Down
69 changes: 37 additions & 32 deletions packages/askar/src/storage/__tests__/AskarStorageService.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,13 +59,15 @@ describe('AskarStorageService', () => {
},
})

const retrieveRecord = await ariesAskar.sessionFetch({
category: record.type,
name: record.id,
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
sessionHandle: wallet.session.handle!,
forUpdate: false,
})
const retrieveRecord = await wallet.withSession((session) =>
ariesAskar.sessionFetch({
category: record.type,
name: record.id,
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
sessionHandle: session.handle!,
forUpdate: false,
})
)

expect(JSON.parse(retrieveRecord?.getTags(0) ?? '{}')).toEqual({
someBoolean: '1',
Expand All @@ -81,31 +83,34 @@ describe('AskarStorageService', () => {
})

it('should correctly transform tag values from string after retrieving', async () => {
await ariesAskar.sessionUpdate({
category: TestRecord.type,
name: 'some-id',
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
sessionHandle: wallet.session.handle!,
value: TypedArrayEncoder.fromString('{}'),
tags: {
someBoolean: '1',
someOtherBoolean: '0',
someStringValue: 'string',
// Before 0.5.0, there was a bug where array values that contained a : would be incorrectly
// parsed back into a record as we would split on ':' and thus only the first part would be included
// in the record as the tag value. If the record was never re-saved it would work well, as well as if the
// record tag was generated dynamically before save (as then the incorrectly transformed back value would be
// overwritten again on save).
'anArrayValueWhereValuesContainColon:foo:bar:test': '1',
'anArrayValueWhereValuesContainColon:https://google.com': '1',
'anArrayValue:foo': '1',
'anArrayValue:bar': '1',
// booleans are stored as '1' and '0' so we store the string values '1' and '0' as 'n__1' and 'n__0'
someStringNumberValue: 'n__1',
anotherStringNumberValue: 'n__0',
},
operation: 0, // EntryOperation.Insert
})
await wallet.withSession(
async (session) =>
await ariesAskar.sessionUpdate({
category: TestRecord.type,
name: 'some-id',
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
sessionHandle: session.handle!,
value: TypedArrayEncoder.fromString('{}'),
tags: {
someBoolean: '1',
someOtherBoolean: '0',
someStringValue: 'string',
// Before 0.5.0, there was a bug where array values that contained a : would be incorrectly
// parsed back into a record as we would split on ':' and thus only the first part would be included
// in the record as the tag value. If the record was never re-saved it would work well, as well as if the
// record tag was generated dynamically before save (as then the incorrectly transformed back value would be
// overwritten again on save).
'anArrayValueWhereValuesContainColon:foo:bar:test': '1',
'anArrayValueWhereValuesContainColon:https://google.com': '1',
'anArrayValue:foo': '1',
'anArrayValue:bar': '1',
// booleans are stored as '1' and '0' so we store the string values '1' and '0' as 'n__1' and 'n__0'
someStringNumberValue: 'n__1',
anotherStringNumberValue: 'n__0',
},
operation: 0, // EntryOperation.Insert
})
)

const record = await storageService.getById(agentContext, TestRecord, 'some-id')

Expand Down
135 changes: 102 additions & 33 deletions packages/askar/src/wallet/AskarBaseWallet.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,6 @@ import { didcommV1Pack, didcommV1Unpack } from './didcommV1'
const isError = (error: unknown): error is Error => error instanceof Error

export abstract class AskarBaseWallet implements Wallet {
protected _session?: Session

protected logger: Logger
protected signingKeyProviderRegistry: SigningProviderRegistry

Expand All @@ -67,12 +65,54 @@ export abstract class AskarBaseWallet implements Wallet {
public abstract dispose(): void | Promise<void>
public abstract profile: string

public get session() {
if (!this._session) {
throw new CredoError('No Wallet Session is opened')
protected abstract store: Store

/**
* Run callback with the session provided, the session will
* be closed once the callback resolves or rejects if it is not closed yet.
*
* TODO: update to new `using` syntax so we don't have to use a callback
*/
public async withSession<Return>(callback: (session: Session) => Return): Promise<Awaited<Return>> {
let session: Session | undefined = undefined
try {
session = await this.store.session(this.profile).open()

const result = await callback(session)

return result
} finally {
if (session?.handle) {
await session.close()
}
}
}

/**
* Run callback with a transaction. If the callback resolves the transaction
* will be committed if the transaction is not closed yet. If the callback rejects
* the transaction will be rolled back if the transaction is not closed yet.
*
* TODO: update to new `using` syntax so we don't have to use a callback
*/
public async withTransaction<Return>(callback: (transaction: Session) => Return): Promise<Awaited<Return>> {
let session: Session | undefined = undefined
try {
session = await this.store.transaction(this.profile).open()

const result = await callback(session)

return this._session
if (session.handle) {
await session.commit()
}
return result
} catch (error) {
if (session?.handle) {
await session?.rollback()
}

throw error
}
}

public get supportedKeyTypes() {
Expand Down Expand Up @@ -105,15 +145,23 @@ export abstract class AskarBaseWallet implements Wallet {
// Create key
let key: AskarKey | undefined
try {
key = privateKey
const _key = privateKey
? AskarKey.fromSecretBytes({ secretKey: privateKey, algorithm })
: seed
? AskarKey.fromSeed({ seed, algorithm })
: AskarKey.generate(algorithm)

// FIXME: we need to create a separate const '_key' so TS definitely knows _key is defined in the session callback.
// This will be fixed once we use the new 'using' syntax
key = _key

const keyPublicBytes = key.publicBytes

// Store key
await this.session.insertKey({ key, name: TypedArrayEncoder.toBase58(keyPublicBytes) })
await this.withSession((session) =>
session.insertKey({ key: _key, name: TypedArrayEncoder.toBase58(keyPublicBytes) })
)

key.handle.free()
return Key.fromPublicKey(keyPublicBytes, keyType)
} catch (error) {
Expand Down Expand Up @@ -162,7 +210,9 @@ export abstract class AskarBaseWallet implements Wallet {

try {
if (isKeyTypeSupportedByAskarForPurpose(key.keyType, AskarKeyTypePurpose.KeyManagement)) {
askarKey = (await this.session.fetchKey({ name: key.publicKeyBase58 }))?.key
askarKey = await this.withSession(
async (session) => (await session.fetchKey({ name: key.publicKeyBase58 }))?.key
)
}

// FIXME: remove the custom KeyPair record now that we deprecate Indy SDK.
Expand All @@ -171,19 +221,25 @@ export abstract class AskarBaseWallet implements Wallet {
// Fallback to fetching key from the non-askar storage, this is to handle the case
// where a key wasn't supported at first by the wallet, but now is
if (!askarKey) {
// TODO: we should probably make retrieveKeyPair + insertKey + deleteKeyPair a transaction
keyPair = await this.retrieveKeyPair(key.publicKeyBase58)

// If we have the key stored in a custom record, but it is now supported by Askar,
// we 'import' the key into askar storage and remove the custom key record
if (keyPair && isKeyTypeSupportedByAskarForPurpose(keyPair.keyType, AskarKeyTypePurpose.KeyManagement)) {
askarKey = AskarKey.fromSecretBytes({
const _askarKey = AskarKey.fromSecretBytes({
secretKey: TypedArrayEncoder.fromBase58(keyPair.privateKeyBase58),
algorithm: keyAlgFromString(keyPair.keyType),
})
await this.session.insertKey({
name: key.publicKeyBase58,
key: askarKey,
})
askarKey = _askarKey

await this.withSession((session) =>
session.insertKey({
name: key.publicKeyBase58,
key: _askarKey,
})
)

// Now we can remove it from the custom record as we have imported it into Askar
await this.deleteKeyPair(key.publicKeyBase58)
keyPair = undefined
Expand Down Expand Up @@ -313,7 +369,9 @@ export abstract class AskarBaseWallet implements Wallet {
recipientKeys: string[],
senderVerkey?: string // in base58
): Promise<EncryptedMessage> {
const senderKey = senderVerkey ? await this.session.fetchKey({ name: senderVerkey }) : undefined
const senderKey = senderVerkey
? await this.withSession((session) => session.fetchKey({ name: senderVerkey }))
: undefined

try {
if (senderVerkey && !senderKey) {
Expand All @@ -339,18 +397,25 @@ export abstract class AskarBaseWallet implements Wallet {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const recipientKids: string[] = protectedJson.recipients.map((r: any) => r.header.kid)

for (const recipientKid of recipientKids) {
const recipientKeyEntry = await this.session.fetchKey({ name: recipientKid })
try {
if (recipientKeyEntry) {
return didcommV1Unpack(messagePackage, recipientKeyEntry.key)
// TODO: how long should sessions last? Just for the duration of the unpack? Or should each item in the recipientKids get a separate session?
const returnValue = await this.withSession(async (session) => {
for (const recipientKid of recipientKids) {
const recipientKeyEntry = await session.fetchKey({ name: recipientKid })
try {
if (recipientKeyEntry) {
return didcommV1Unpack(messagePackage, recipientKeyEntry.key)
}
} finally {
recipientKeyEntry?.key.handle.free()
}
} finally {
recipientKeyEntry?.key.handle.free()
}
})

if (!returnValue) {
throw new WalletError('No corresponding recipient key found')
}

throw new WalletError('No corresponding recipient key found')
return returnValue
}

public async generateNonce(): Promise<string> {
Expand All @@ -376,7 +441,9 @@ export abstract class AskarBaseWallet implements Wallet {

private async retrieveKeyPair(publicKeyBase58: string): Promise<KeyPair | null> {
try {
const entryObject = await this.session.fetch({ category: 'KeyPairRecord', name: `key-${publicKeyBase58}` })
const entryObject = await this.withSession((session) =>
session.fetch({ category: 'KeyPairRecord', name: `key-${publicKeyBase58}` })
)

if (!entryObject) return null

Expand All @@ -388,22 +455,24 @@ export abstract class AskarBaseWallet implements Wallet {

private async deleteKeyPair(publicKeyBase58: string): Promise<void> {
try {
await this.session.remove({ category: 'KeyPairRecord', name: `key-${publicKeyBase58}` })
await this.withSession((session) => session.remove({ category: 'KeyPairRecord', name: `key-${publicKeyBase58}` }))
} catch (error) {
throw new WalletError('Error removing KeyPair record', { cause: error })
}
}

private async storeKeyPair(keyPair: KeyPair): Promise<void> {
try {
await this.session.insert({
category: 'KeyPairRecord',
name: `key-${keyPair.publicKeyBase58}`,
value: JSON.stringify(keyPair),
tags: {
keyType: keyPair.keyType,
},
})
await this.withSession((session) =>
session.insert({
category: 'KeyPairRecord',
name: `key-${keyPair.publicKeyBase58}`,
value: JSON.stringify(keyPair),
tags: {
keyType: keyPair.keyType,
},
})
)
} catch (error) {
if (isAskarError(error, AskarErrorCode.Duplicate)) {
throw new WalletKeyExistsError('Key already exists')
Expand Down
Loading

0 comments on commit 2cb9ba8

Please sign in to comment.