Skip to content

Commit 9e91b97

Browse files
committed
WIP
1 parent bc3b2ac commit 9e91b97

File tree

6 files changed

+113
-44
lines changed

6 files changed

+113
-44
lines changed

example/main.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99
"github.com/aws/aws-sdk-go-v2/service/dynamodb"
1010
"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
1111
"github.com/aws/aws-sdk-go/aws"
12-
"github.com/cloudopsy/dynamodb-encryption-go/pkg/client"
12+
"github.com/cloudopsy/dynamodb-encryption-go/pkg/encrypted"
1313
"github.com/cloudopsy/dynamodb-encryption-go/pkg/provider"
1414
"github.com/cloudopsy/dynamodb-encryption-go/pkg/provider/store"
1515
)
@@ -47,8 +47,11 @@ func main() {
4747
log.Fatalf("Failed to create cryptographic materials provider: %v", err)
4848
}
4949

50+
attributeAction := encrypted.NewAttributeActions(encrypted.AttributeActionDoNothing)
51+
attributeAction.SetAttributeAction("Password", encrypted.AttributeActionEncrypt)
52+
5053
// Initialize EncryptedClient
51-
ec := client.NewEncryptedClient(dynamoDBClient, cmp)
54+
ec := encrypted.NewEncryptedClient(dynamoDBClient, cmp, attributeAction)
5255

5356
// User credentials to encrypt and store
5457
userID := "user1"
@@ -162,7 +165,7 @@ func main() {
162165
deleteItem := &dynamodb.DeleteItemInput{
163166
TableName: &tableName,
164167
Key: map[string]types.AttributeValue{
165-
"UserID": &types.AttributeValueMemberS{Value: "user1"},
168+
"UserID": &types.AttributeValueMemberS{Value: "user2"},
166169
},
167170
}
168171

pkg/delegatedkeys/delegated_keys.go

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import (
77

88
"github.com/tink-crypto/tink-go-awskms/integration/awskms"
99
"github.com/tink-crypto/tink-go/v2/aead"
10-
"github.com/tink-crypto/tink-go/v2/daead"
1110
"github.com/tink-crypto/tink-go/v2/keyset"
1211
"github.com/tink-crypto/tink-go/v2/signature"
1312
)
@@ -61,6 +60,7 @@ func (dk *TinkDelegatedKey) AllowedForRawMaterials() bool {
6160
}
6261

6362
func (dk *TinkDelegatedKey) Encrypt(plaintext []byte, associatedData []byte) ([]byte, error) {
63+
// TODO: Support AEAD and DAEAD primitives
6464
a, err := aead.New(dk.keysetHandle)
6565
if err != nil {
6666
return nil, fmt.Errorf("failed to get AEAD primitive: %v", err)
@@ -69,6 +69,7 @@ func (dk *TinkDelegatedKey) Encrypt(plaintext []byte, associatedData []byte) ([]
6969
}
7070

7171
func (dk *TinkDelegatedKey) Decrypt(ciphertext []byte, associatedData []byte) ([]byte, error) {
72+
// TODO: Support AEAD and DAEAD primitives
7273
a, err := aead.New(dk.keysetHandle)
7374
if err != nil {
7475
return nil, fmt.Errorf("failed to get AEAD primitive: %v", err)
@@ -142,21 +143,11 @@ func UnwrapKeyset(encryptedKeyset []byte, kekUri string) (*TinkDelegatedKey, err
142143
return NewTinkDelegatedKey(handle, kekUri), nil
143144
}
144145

145-
func GenerateDataKey(keyURI, keyType string) (*TinkDelegatedKey, []byte, error) {
146-
var kh *keyset.Handle
147-
var err error
148-
149-
switch strings.ToLower(keyType) {
150-
case "aead":
151-
kh, err = keyset.NewHandle(aead.AES256GCMKeyTemplate())
152-
case "daead":
153-
kh, err = keyset.NewHandle(daead.AESSIVKeyTemplate())
154-
default:
155-
return nil, nil, fmt.Errorf("unsupported key type: %v", keyType)
156-
}
146+
func GenerateDataKey(keyURI string) (*TinkDelegatedKey, []byte, error) {
147+
kh, err := keyset.NewHandle(aead.AES256GCMKeyTemplate())
157148

158149
if err != nil {
159-
return nil, nil, fmt.Errorf("failed to generate new keyset handle for %v: %v", keyType, err)
150+
return nil, nil, fmt.Errorf("failed to generate new keyset handle: %v", err)
160151
}
161152

162153
delegatedKey := NewTinkDelegatedKey(kh, keyURI)

pkg/encrypted/action.go

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package encrypted
2+
3+
type AttributeAction int
4+
5+
const (
6+
AttributeActionDoNothing AttributeAction = iota
7+
AttributeActionEncrypt
8+
AttributeActionEncryptDeterministically
9+
AttributeActionSign
10+
)
11+
12+
type AttributeActions struct {
13+
defaultAction AttributeAction
14+
attributeActions map[string]AttributeAction
15+
}
16+
17+
func NewAttributeActions(defaultAction AttributeAction) *AttributeActions {
18+
return &AttributeActions{
19+
defaultAction: defaultAction,
20+
attributeActions: make(map[string]AttributeAction),
21+
}
22+
}
23+
24+
func (aa *AttributeActions) SetDefaultAction(action AttributeAction) {
25+
aa.defaultAction = action
26+
}
27+
28+
func (aa *AttributeActions) SetAttributeAction(attributeName string, action AttributeAction) {
29+
aa.attributeActions[attributeName] = action
30+
}
31+
32+
func (aa *AttributeActions) GetAttributeAction(attributeName string) AttributeAction {
33+
action, ok := aa.attributeActions[attributeName]
34+
if !ok {
35+
return aa.defaultAction
36+
}
37+
return action
38+
}

pkg/client/client.go renamed to pkg/encrypted/client.go

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package client
1+
package encrypted
22

33
import (
44
"context"
@@ -35,15 +35,18 @@ type EncryptedClient struct {
3535
client DynamoDBClientInterface
3636
materialsProvider provider.CryptographicMaterialsProvider
3737
primaryKeyCache map[string]*PrimaryKeyInfo
38-
lock sync.RWMutex
38+
attributeActions *AttributeActions
39+
40+
lock sync.RWMutex
3941
}
4042

4143
// NewEncryptedClient creates a new instance of EncryptedClient.
42-
func NewEncryptedClient(client DynamoDBClientInterface, materialsProvider provider.CryptographicMaterialsProvider) *EncryptedClient {
44+
func NewEncryptedClient(client DynamoDBClientInterface, materialsProvider provider.CryptographicMaterialsProvider, attributeActions *AttributeActions) *EncryptedClient {
4345
return &EncryptedClient{
4446
client: client,
4547
materialsProvider: materialsProvider,
4648
primaryKeyCache: make(map[string]*PrimaryKeyInfo),
49+
attributeActions: attributeActions,
4750
lock: sync.RWMutex{},
4851
}
4952
}
@@ -289,12 +292,18 @@ func (ec *EncryptedClient) encryptItem(ctx context.Context, tableName string, it
289292
return nil, fmt.Errorf("error converting attribute value to bytes: %v", err)
290293
}
291294

292-
encryptedData, err := encryptionMaterials.EncryptionKey().Encrypt(rawData, []byte(key))
293-
if err != nil {
294-
return nil, fmt.Errorf("error encrypting attribute value: %v", err)
295+
action := ec.attributeActions.GetAttributeAction(key)
296+
switch action {
297+
case AttributeActionEncrypt, AttributeActionEncryptDeterministically:
298+
// TODO: Implement deterministic encryption
299+
encryptedData, err := encryptionMaterials.EncryptionKey().Encrypt(rawData, []byte(key))
300+
if err != nil {
301+
return nil, fmt.Errorf("error encrypting attribute value: %v", err)
302+
}
303+
encryptedItem[key] = &types.AttributeValueMemberB{Value: encryptedData}
304+
case AttributeActionDoNothing:
305+
encryptedItem[key] = value
295306
}
296-
297-
encryptedItem[key] = &types.AttributeValueMemberB{Value: encryptedData}
298307
}
299308

300309
return encryptedItem, nil
@@ -327,20 +336,29 @@ func (ec *EncryptedClient) decryptItem(ctx context.Context, tableName string, it
327336

328337
encryptedData, ok := value.(*types.AttributeValueMemberB)
329338
if !ok {
330-
return nil, fmt.Errorf("expected binary data for encrypted attribute value")
339+
// If the attribute is not encrypted, copy it as is
340+
decryptedItem[key] = value
341+
continue
331342
}
332343

333-
rawData, err := decryptionMaterials.DecryptionKey().Decrypt(encryptedData.Value, []byte(key))
334-
if err != nil {
335-
return nil, fmt.Errorf("error decrypting attribute value: %v", err)
336-
}
344+
action := ec.attributeActions.GetAttributeAction(key)
345+
switch action {
346+
case AttributeActionEncrypt, AttributeActionEncryptDeterministically:
347+
// TODO: Implement deterministic encryption
348+
rawData, err := decryptionMaterials.DecryptionKey().Decrypt(encryptedData.Value, []byte(key))
349+
if err != nil {
350+
return nil, fmt.Errorf("error decrypting attribute value: %v", err)
351+
}
352+
decryptedValue, err := utils.BytesToAttributeValue(rawData)
353+
if err != nil {
354+
return nil, fmt.Errorf("error converting bytes to attribute value: %v", err)
355+
}
337356

338-
decryptedValue, err := utils.BytesToAttributeValue(rawData)
339-
if err != nil {
340-
return nil, fmt.Errorf("error converting bytes to attribute value: %v", err)
357+
decryptedItem[key] = decryptedValue
358+
case AttributeActionDoNothing:
359+
decryptedItem[key] = value
341360
}
342361

343-
decryptedItem[key] = decryptedValue
344362
}
345363

346364
return decryptedItem, nil

pkg/client/client_test.go renamed to pkg/encrypted/client_test.go

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package client
1+
package encrypted
22

33
import (
44
"context"
@@ -162,7 +162,10 @@ func TestEncryptedClient_PutItem(t *testing.T) {
162162
nil,
163163
), nil)
164164

165-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
165+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
166+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
167+
168+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
166169

167170
item := map[string]types.AttributeValue{
168171
"PK": &types.AttributeValueMemberS{Value: "123"},
@@ -186,7 +189,10 @@ func TestEncryptedClient_PutItem(t *testing.T) {
186189
func TestEncryptedClient_PutItem_Failure(t *testing.T) {
187190
mockDynamoDBClient := new(MockDynamoDBClient)
188191
mockCMProvider := new(MockCryptographicMaterialsProvider)
189-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
192+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
193+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
194+
195+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
190196

191197
// Mock the DescribeTable call to simulate fetching table primary key schema.
192198
mockDynamoDBClient.On("DescribeTable", mock.Anything, mock.AnythingOfType("*dynamodb.DescribeTableInput"), mock.Anything).Return(&dynamodb.DescribeTableOutput{
@@ -229,7 +235,10 @@ func TestEncryptedClient_PutItem_Failure(t *testing.T) {
229235
func TestEncryptedClient_GetItem_Success(t *testing.T) {
230236
mockDynamoDBClient := new(MockDynamoDBClient)
231237
mockCMProvider := new(MockCryptographicMaterialsProvider)
232-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
238+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
239+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
240+
241+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
233242

234243
// Mock DescribeTable call to simulate fetching table primary key schema.
235244
mockDynamoDBClient.On("DescribeTable", mock.Anything, mock.AnythingOfType("*dynamodb.DescribeTableInput"), mock.Anything).Return(&dynamodb.DescribeTableOutput{
@@ -277,7 +286,9 @@ func TestEncryptedClient_GetItem_Success(t *testing.T) {
277286
func TestEncryptedClient_Query(t *testing.T) {
278287
mockDynamoDBClient := new(MockDynamoDBClient)
279288
mockCMProvider := new(MockCryptographicMaterialsProvider)
280-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
289+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
290+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
291+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
281292

282293
// Mock DescribeTable call to simulate fetching table primary key schema.
283294
mockDynamoDBClient.On("DescribeTable", mock.Anything, mock.AnythingOfType("*dynamodb.DescribeTableInput"), mock.Anything).Return(&dynamodb.DescribeTableOutput{
@@ -331,7 +342,9 @@ func TestEncryptedClient_Query(t *testing.T) {
331342
func TestEncryptedClient_Scan(t *testing.T) {
332343
mockDynamoDBClient := new(MockDynamoDBClient)
333344
mockCMProvider := new(MockCryptographicMaterialsProvider)
334-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
345+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
346+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
347+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
335348

336349
// Mock DescribeTable call to simulate fetching table primary key schema.
337350
mockDynamoDBClient.On("DescribeTable", mock.Anything, mock.AnythingOfType("*dynamodb.DescribeTableInput"), mock.Anything).Return(&dynamodb.DescribeTableOutput{
@@ -381,7 +394,9 @@ func TestEncryptedClient_Scan(t *testing.T) {
381394
func TestEncryptedClient_BatchGetItem(t *testing.T) {
382395
mockDynamoDBClient := new(MockDynamoDBClient)
383396
mockCMProvider := new(MockCryptographicMaterialsProvider)
384-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
397+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
398+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
399+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
385400

386401
// Mock DescribeTable call to simulate fetching table primary key schema.
387402
mockDynamoDBClient.On("DescribeTable", mock.Anything, mock.AnythingOfType("*dynamodb.DescribeTableInput"), mock.Anything).Return(&dynamodb.DescribeTableOutput{
@@ -446,7 +461,9 @@ func TestEncryptedClient_BatchGetItem(t *testing.T) {
446461
func TestEncryptedClient_BatchWriteItem(t *testing.T) {
447462
mockDynamoDBClient := new(MockDynamoDBClient)
448463
mockCMProvider := new(MockCryptographicMaterialsProvider)
449-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
464+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
465+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
466+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
450467

451468
// Mock DescribeTable call to simulate fetching table primary key schema.
452469
mockDynamoDBClient.On("DescribeTable", mock.Anything, mock.AnythingOfType("*dynamodb.DescribeTableInput"), mock.Anything).Return(&dynamodb.DescribeTableOutput{
@@ -503,7 +520,9 @@ func TestEncryptedClient_BatchWriteItem(t *testing.T) {
503520
func TestEncryptedClient_DeleteItem(t *testing.T) {
504521
mockDynamoDBClient := new(MockDynamoDBClient)
505522
mockCMProvider := new(MockCryptographicMaterialsProvider)
506-
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider)
523+
attributeActions := NewAttributeActions(AttributeActionDoNothing)
524+
attributeActions.SetAttributeAction("SensitiveAttribute", AttributeActionEncrypt)
525+
encryptedClient := NewEncryptedClient(mockDynamoDBClient, mockCMProvider, attributeActions)
507526

508527
// Mock DescribeTable call to simulate fetching table primary key schema.
509528
mockDynamoDBClient.On("DescribeTable", mock.Anything, mock.AnythingOfType("*dynamodb.DescribeTableInput"), mock.Anything).Return(&dynamodb.DescribeTableOutput{

pkg/provider/kms.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ func NewAwsKmsCryptographicMaterialsProvider(keyID string, encryptionContext map
3030

3131
// GenerateDataKey generates a new data key using AWS KMS and wraps the Tink keyset.
3232
func (p *AwsKmsCryptographicMaterialsProvider) GenerateDataKey() (*delegatedkeys.TinkDelegatedKey, []byte, error) {
33-
delegatedKey, wrappedKeyset, err := delegatedkeys.GenerateDataKey(p.KeyID, "aead")
33+
delegatedKey, wrappedKeyset, err := delegatedkeys.GenerateDataKey(p.KeyID)
3434
if err != nil {
3535
return nil, nil, fmt.Errorf("failed to generate data key: %v", err)
3636
}

0 commit comments

Comments
 (0)