Add function for accessing underlying PKCS#11 keys

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)
+	})
+}

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
+		})
+	})
+}