Skip to content

Commit 2bb7eb0

Browse files
authored
feat: add AES protected key interface and implementation (#2599)
### Proposed Changes Add ProtectedKey and Encapsulator interfaces to lib/ocrypto package, along with AESProtectedKey implementation. This exposes previously internal cryptographic functionality for external consumption. - Add interfaces.go with ProtectedKey and Encapsulator interfaces - Add protected_key.go with AESProtectedKey implementation - Add comprehensive test suite in protected_key_test.go - Use standard error types without logging dependencies ### Checklist - [ ] I have added or updated unit tests - [ ] I have added or updated integration tests (if appropriate) - [ ] I have added or updated documentation ### Testing Instructions
1 parent 6a233c1 commit 2bb7eb0

File tree

3 files changed

+328
-0
lines changed

3 files changed

+328
-0
lines changed

lib/ocrypto/interfaces.go

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package ocrypto
2+
3+
import (
4+
"context"
5+
)
6+
7+
// Encapsulator enables key encapsulation with a public key
8+
type Encapsulator interface {
9+
// Encapsulate wraps a secret key with the encapsulation key
10+
Encapsulate(dek ProtectedKey) ([]byte, error)
11+
12+
// Encrypt wraps a secret key with the encapsulation key
13+
Encrypt(data []byte) ([]byte, error)
14+
15+
// PublicKeyAsPEM exports the public key, used to encapsulate the value, in Privacy-Enhanced Mail format,
16+
// or the empty string if not present.
17+
PublicKeyAsPEM() (string, error)
18+
19+
// For EC schemes, this method returns the public part of the ephemeral key.
20+
// Otherwise, it returns nil.
21+
EphemeralKey() []byte
22+
}
23+
24+
// ProtectedKey represents a decrypted key with operations that can be performed on it
25+
type ProtectedKey interface {
26+
// VerifyBinding checks if the policy binding matches the given policy data
27+
VerifyBinding(ctx context.Context, policy, policyBinding []byte) error
28+
29+
// Export returns the raw key data, optionally encrypting it with the provided encapsulator
30+
Export(encapsulator Encapsulator) ([]byte, error)
31+
32+
// DecryptAESGCM decrypts encrypted policies and metadata
33+
DecryptAESGCM(iv []byte, body []byte, tagSize int) ([]byte, error)
34+
}

lib/ocrypto/protected_key.go

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
package ocrypto
2+
3+
import (
4+
"context"
5+
"crypto/hmac"
6+
"crypto/sha256"
7+
"errors"
8+
"fmt"
9+
)
10+
11+
var (
12+
// ErrEmptyKeyData is returned when the key data is empty
13+
ErrEmptyKeyData = errors.New("key data is empty")
14+
// ErrPolicyHMACMismatch is returned when policy binding verification fails
15+
ErrPolicyHMACMismatch = errors.New("policy HMAC mismatch")
16+
)
17+
18+
// AESProtectedKey implements the ProtectedKey interface with an in-memory secret key
19+
type AESProtectedKey struct {
20+
rawKey []byte
21+
aesGcm AesGcm
22+
}
23+
24+
var _ ProtectedKey = (*AESProtectedKey)(nil)
25+
26+
// NewAESProtectedKey creates a new instance of AESProtectedKey
27+
func NewAESProtectedKey(rawKey []byte) (*AESProtectedKey, error) {
28+
if len(rawKey) == 0 {
29+
return nil, ErrEmptyKeyData
30+
}
31+
// Create a defensive copy of the key
32+
keyCopy := append([]byte{}, rawKey...)
33+
34+
// Pre-initialize the AES-GCM cipher for performance
35+
aesGcm, err := NewAESGcm(keyCopy)
36+
if err != nil {
37+
return nil, fmt.Errorf("failed to initialize AES-GCM cipher: %w", err)
38+
}
39+
40+
return &AESProtectedKey{
41+
rawKey: keyCopy,
42+
aesGcm: aesGcm,
43+
}, nil
44+
}
45+
46+
// DecryptAESGCM decrypts data using AES-GCM with the protected key
47+
func (k *AESProtectedKey) DecryptAESGCM(iv []byte, body []byte, tagSize int) ([]byte, error) {
48+
// Use the pre-initialized AES-GCM cipher for better performance
49+
decryptedData, err := k.aesGcm.DecryptWithIVAndTagSize(iv, body, tagSize)
50+
if err != nil {
51+
return nil, fmt.Errorf("AES-GCM decryption failed: %w", err)
52+
}
53+
54+
return decryptedData, nil
55+
}
56+
57+
// Export returns the raw key data, optionally encrypting it with the provided Encapsulator
58+
func (k *AESProtectedKey) Export(encapsulator Encapsulator) ([]byte, error) {
59+
if encapsulator == nil {
60+
// Return error if encapsulator is nil
61+
return nil, errors.New("encapsulator cannot be nil")
62+
}
63+
64+
// Encrypt the key data before returning
65+
encryptedKey, err := encapsulator.Encrypt(k.rawKey)
66+
if err != nil {
67+
return nil, fmt.Errorf("failed to encrypt key data for export: %w", err)
68+
}
69+
70+
return encryptedKey, nil
71+
}
72+
73+
// VerifyBinding checks if the policy binding matches the given policy data
74+
func (k *AESProtectedKey) VerifyBinding(_ context.Context, policy, policyBinding []byte) error {
75+
actualHMAC := k.generateHMACDigest(policy)
76+
77+
if !hmac.Equal(actualHMAC, policyBinding) {
78+
return ErrPolicyHMACMismatch
79+
}
80+
81+
return nil
82+
}
83+
84+
// generateHMACDigest is a helper to generate an HMAC digest from a message using the key
85+
func (k *AESProtectedKey) generateHMACDigest(msg []byte) []byte {
86+
mac := hmac.New(sha256.New, k.rawKey)
87+
mac.Write(msg)
88+
return mac.Sum(nil)
89+
}

lib/ocrypto/protected_key_test.go

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
package ocrypto
2+
3+
import (
4+
"context"
5+
"crypto/rand"
6+
"testing"
7+
8+
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/require"
10+
)
11+
12+
const testKey = "test-key-12345678901234567890123"
13+
14+
func TestNewAESProtectedKey(t *testing.T) {
15+
key := make([]byte, 32)
16+
_, err := rand.Read(key)
17+
require.NoError(t, err)
18+
19+
protectedKey, err := NewAESProtectedKey(key)
20+
require.NoError(t, err)
21+
assert.NotNil(t, protectedKey)
22+
assert.Equal(t, key, protectedKey.rawKey)
23+
}
24+
25+
func TestAESProtectedKey_DecryptAESGCM(t *testing.T) {
26+
// Generate a random 256-bit key
27+
key := make([]byte, 32)
28+
_, err := rand.Read(key)
29+
require.NoError(t, err)
30+
31+
protectedKey, err := NewAESProtectedKey(key)
32+
require.NoError(t, err)
33+
34+
// Test data
35+
plaintext := []byte("Hello, World!")
36+
37+
// Encrypt the data first using the same key
38+
aesGcm, err := NewAESGcm(key)
39+
require.NoError(t, err)
40+
41+
encrypted, err := aesGcm.Encrypt(plaintext)
42+
require.NoError(t, err)
43+
44+
// Extract IV and ciphertext (first 12 bytes are IV for GCM standard nonce size)
45+
iv := encrypted[:GcmStandardNonceSize]
46+
ciphertext := encrypted[GcmStandardNonceSize:]
47+
48+
// Test decryption
49+
decrypted, err := protectedKey.DecryptAESGCM(iv, ciphertext, 16) // 16 is standard GCM tag size
50+
require.NoError(t, err)
51+
assert.Equal(t, plaintext, decrypted)
52+
}
53+
54+
func TestAESProtectedKey_DecryptAESGCM_InvalidKey(t *testing.T) {
55+
// Empty key should fail
56+
_, err := NewAESProtectedKey([]byte{})
57+
require.Error(t, err)
58+
assert.ErrorIs(t, err, ErrEmptyKeyData)
59+
}
60+
61+
func TestAESProtectedKey_Export_NoEncapsulator(t *testing.T) {
62+
key := []byte(testKey) // 32 bytes
63+
protectedKey, err := NewAESProtectedKey(key)
64+
require.NoError(t, err)
65+
66+
exported, err := protectedKey.Export(nil)
67+
require.Error(t, err)
68+
require.ErrorContains(t, err, "encapsulator cannot be nil")
69+
assert.Nil(t, exported)
70+
}
71+
72+
func TestAESProtectedKey_Export_WithEncapsulator(t *testing.T) {
73+
key := []byte(testKey) // 32 bytes
74+
protectedKey, err := NewAESProtectedKey(key)
75+
require.NoError(t, err)
76+
77+
// Mock encapsulator
78+
mockEncapsulator := &mockEncapsulator{
79+
encryptFunc: func(data []byte) ([]byte, error) {
80+
// Simple XOR encryption for testing
81+
result := make([]byte, len(data))
82+
for i, b := range data {
83+
result[i] = b ^ 0xFF
84+
}
85+
return result, nil
86+
},
87+
}
88+
89+
exported, err := protectedKey.Export(mockEncapsulator)
90+
require.NoError(t, err)
91+
92+
// Verify it was encrypted (should be different from original)
93+
assert.NotEqual(t, key, exported)
94+
assert.Len(t, exported, len(key))
95+
96+
// Verify we can decrypt it back
97+
for i, b := range exported {
98+
assert.Equal(t, key[i], b^0xFF)
99+
}
100+
}
101+
102+
func TestAESProtectedKey_Export_EncapsulatorError(t *testing.T) {
103+
key := []byte(testKey) // 32 bytes
104+
protectedKey, err := NewAESProtectedKey(key)
105+
require.NoError(t, err)
106+
107+
mockEncapsulator := &mockEncapsulator{
108+
encryptFunc: func(_ []byte) ([]byte, error) {
109+
return nil, assert.AnError
110+
},
111+
}
112+
113+
_, err = protectedKey.Export(mockEncapsulator)
114+
require.Error(t, err)
115+
assert.Contains(t, err.Error(), "failed to encrypt key data for export")
116+
}
117+
118+
func TestAESProtectedKey_VerifyBinding(t *testing.T) {
119+
key := []byte(testKey) // 32 bytes
120+
protectedKey, err := NewAESProtectedKey(key)
121+
require.NoError(t, err)
122+
123+
policy := []byte("test-policy-data")
124+
ctx := context.Background()
125+
126+
// Generate the expected HMAC
127+
expectedHMAC := protectedKey.generateHMACDigest(policy)
128+
129+
// Verify binding should succeed with correct HMAC
130+
err = protectedKey.VerifyBinding(ctx, policy, expectedHMAC)
131+
assert.NoError(t, err)
132+
}
133+
134+
func TestAESProtectedKey_VerifyBinding_Mismatch(t *testing.T) {
135+
key := []byte(testKey) // 32 bytes
136+
protectedKey, err := NewAESProtectedKey(key)
137+
require.NoError(t, err)
138+
139+
policy := []byte("test-policy-data")
140+
wrongBinding := []byte("wrong-binding-data")
141+
ctx := context.Background()
142+
143+
err = protectedKey.VerifyBinding(ctx, policy, wrongBinding)
144+
require.Error(t, err)
145+
assert.Equal(t, ErrPolicyHMACMismatch, err)
146+
}
147+
148+
func TestAESProtectedKey_VerifyBinding_DifferentPolicyData(t *testing.T) {
149+
key := []byte(testKey) // 32 bytes
150+
protectedKey, err := NewAESProtectedKey(key)
151+
require.NoError(t, err)
152+
153+
ctx := context.Background()
154+
155+
// Generate HMAC for first policy
156+
policy1 := []byte("policy-data-1")
157+
hmac1 := protectedKey.generateHMACDigest(policy1)
158+
159+
// Try to verify with different policy data
160+
policy2 := []byte("policy-data-2")
161+
err = protectedKey.VerifyBinding(ctx, policy2, hmac1)
162+
require.Error(t, err)
163+
assert.Equal(t, ErrPolicyHMACMismatch, err)
164+
}
165+
166+
func TestAESProtectedKey_InterfaceCompliance(t *testing.T) {
167+
key := make([]byte, 32)
168+
protectedKey, err := NewAESProtectedKey(key)
169+
require.NoError(t, err)
170+
171+
// Ensure it implements the ProtectedKey interface
172+
assert.Implements(t, (*ProtectedKey)(nil), protectedKey)
173+
}
174+
175+
// Mock encapsulator for testing
176+
type mockEncapsulator struct {
177+
encryptFunc func([]byte) ([]byte, error)
178+
publicKeyPEMFunc func() (string, error)
179+
ephemeralKeyFunc func() []byte
180+
}
181+
182+
func (m *mockEncapsulator) Encapsulate(_ ProtectedKey) ([]byte, error) {
183+
return nil, nil
184+
}
185+
186+
func (m *mockEncapsulator) Encrypt(data []byte) ([]byte, error) {
187+
if m.encryptFunc != nil {
188+
return m.encryptFunc(data)
189+
}
190+
return data, nil
191+
}
192+
193+
func (m *mockEncapsulator) PublicKeyAsPEM() (string, error) {
194+
if m.publicKeyPEMFunc != nil {
195+
return m.publicKeyPEMFunc()
196+
}
197+
return "", nil
198+
}
199+
200+
func (m *mockEncapsulator) EphemeralKey() []byte {
201+
if m.ephemeralKeyFunc != nil {
202+
return m.ephemeralKeyFunc()
203+
}
204+
return nil
205+
}

0 commit comments

Comments
 (0)