diff --git a/keys.go b/keys.go index 998663a..64bde0d 100644 --- a/keys.go +++ b/keys.go @@ -611,3 +611,28 @@ func (c *Context) GetPubAttribute(key interface{}, attribute AttributeType) (a * return set[attribute], nil } + +// WithPKCS11Key provides a custom extension point that exposes the underlying PKCS#11 object handle and an active session to the +// supplied callback. This can be used to e.g. extend this package non-intrusively and provide custom HSM interaction code. +func (c *Context) WithPKCS11Key(key interface{}, f func(*pkcs11.Ctx, pkcs11.SessionHandle, pkcs11.ObjectHandle) error) error { + if c.closed.Get() { + return errClosed + } + var objectHandle pkcs11.ObjectHandle + switch k := (key).(type) { + case *pkcs11PrivateKeyDSA: + objectHandle = k.handle + case *pkcs11PrivateKeyRSA: + objectHandle = k.handle + case *pkcs11PrivateKeyECDSA: + objectHandle = k.handle + case *SecretKey: + objectHandle = k.handle + default: + return errors.Errorf("not a PKCS#11 key") + } + // invoke f with a valid session. + return c.withSession(func(session *pkcs11Session) error { + return f(session.ctx, session.handle, objectHandle) + }) +} diff --git a/keys_test.go b/keys_test.go index ca8ab6a..b29398b 100644 --- a/keys_test.go +++ b/keys_test.go @@ -221,3 +221,19 @@ func TestGettingUnsupportedKeyTypeAttributes(t *testing.T) { require.Error(t, err) }) } + +func TestWithPKCS11Key(t *testing.T) { + withContext(t, func(ctx *Context) { + id := randomBytes() + key, err := ctx.GenerateSecretKey(id, 128, CipherAES) + require.NoError(t, err) + + ctx.WithPKCS11Key(key, func(ctx *pkcs11.Ctx, sessionHandle pkcs11.SessionHandle, keyHandle pkcs11.ObjectHandle) error { + // PKCS#11 2.20 spec: "Valid session handles and object handles in Cryptoki always have nonzero values." + assert.True(t, sessionHandle != 0) + // Verify the key's internal handle matches what is passed down. + assert.Equal(t, keyHandle, key.handle) + return nil + }) + }) +}