Skip to content

Commit c238e85

Browse files
netfnetf
committed
WIP
1 parent 904bac4 commit c238e85

File tree

3 files changed

+88
-24
lines changed

3 files changed

+88
-24
lines changed

‎example/main.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,7 @@ func main() {
162162
deleteItem := &dynamodb.DeleteItemInput{
163163
TableName: &tableName,
164164
Key: map[string]types.AttributeValue{
165-
"UserID": &types.AttributeValueMemberS{Value: "user2"},
165+
"UserID": &types.AttributeValueMemberS{Value: "user1"},
166166
},
167167
}
168168

‎pkg/client/client.go

Lines changed: 16 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ func NewEncryptedClient(client *dynamodb.Client, materialsProvider provider.Cryp
3333
// PutItem encrypts an item and puts it into a DynamoDB table.
3434
func (ec *EncryptedClient) PutItem(ctx context.Context, input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
3535
// Encrypt the item, excluding primary keys
36-
encryptedItem, err := ec.encryptItem(ctx, *input.TableName, input.Item)
36+
encryptedItem, err := ec.encryptItem(ctx, aws.StringValue(input.TableName), input.Item)
3737
if err != nil {
3838
return nil, fmt.Errorf("failed to encrypt item: %v", err)
3939
}
@@ -62,7 +62,7 @@ func (ec *EncryptedClient) GetItem(ctx context.Context, input *dynamodb.GetItemI
6262
}
6363

6464
// Decrypt the item, excluding primary keys
65-
decryptedItem, err := ec.decryptItem(ctx, *input.TableName, encryptedOutput.Item)
65+
decryptedItem, err := ec.decryptItem(ctx, aws.StringValue(input.TableName), encryptedOutput.Item)
6666
if err != nil {
6767
return nil, fmt.Errorf("failed to decrypt item: %v", err)
6868
}
@@ -84,7 +84,7 @@ func (ec *EncryptedClient) Query(ctx context.Context, input *dynamodb.QueryInput
8484

8585
// Decrypt the items in the response
8686
for i, item := range encryptedOutput.Items {
87-
decryptedItem, decryptErr := ec.decryptItem(ctx, *input.TableName, item)
87+
decryptedItem, decryptErr := ec.decryptItem(ctx, aws.StringValue(input.TableName), item)
8888
if decryptErr != nil {
8989
return nil, decryptErr
9090
}
@@ -103,7 +103,7 @@ func (ec *EncryptedClient) Scan(ctx context.Context, input *dynamodb.ScanInput)
103103

104104
// Decrypt the items in the response
105105
for i, item := range encryptedOutput.Items {
106-
decryptedItem, decryptErr := ec.decryptItem(ctx, *input.TableName, item)
106+
decryptedItem, decryptErr := ec.decryptItem(ctx, aws.StringValue(input.TableName), item)
107107
if decryptErr != nil {
108108
return nil, decryptErr
109109
}
@@ -168,7 +168,10 @@ func (ec *EncryptedClient) DeleteItem(ctx context.Context, input *dynamodb.Delet
168168
}
169169

170170
// Construct material name based on the primary key of the item being deleted
171-
materialName := ec.constructMaterialName(input.Key, pkInfo)
171+
materialName, err := utils.ConstructMaterialName(input.Key, pkInfo)
172+
if err != nil {
173+
return nil, fmt.Errorf("error constructing material name: %v", err)
174+
}
172175

173176
// Delete the associated metadata
174177
tableName := ec.materialsProvider.TableName()
@@ -246,7 +249,10 @@ func (ec *EncryptedClient) encryptItem(ctx context.Context, tableName string, it
246249
}
247250

248251
// Generate and fetch encryption materials
249-
materialName := ec.constructMaterialName(item, pkInfo)
252+
materialName, err := utils.ConstructMaterialName(item, pkInfo)
253+
if err != nil {
254+
return nil, fmt.Errorf("error constructing material name: %v", err)
255+
}
250256
encryptionMaterials, err := ec.materialsProvider.EncryptionMaterials(ctx, materialName)
251257
if err != nil {
252258
return nil, fmt.Errorf("failed to fetch encryption materials: %v", err)
@@ -284,7 +290,10 @@ func (ec *EncryptedClient) decryptItem(ctx context.Context, tableName string, it
284290
}
285291

286292
// Construct the material name based on primary keys
287-
materialName := ec.constructMaterialName(item, pkInfo)
293+
materialName, err := utils.ConstructMaterialName(item, pkInfo)
294+
if err != nil {
295+
return nil, fmt.Errorf("error constructing material name: %v", err)
296+
}
288297
decryptionMaterials, err := ec.materialsProvider.DecryptionMaterials(ctx, materialName, 0)
289298
if err != nil {
290299
return nil, fmt.Errorf("failed to fetch decryption materials: %v", err)
@@ -318,19 +327,3 @@ func (ec *EncryptedClient) decryptItem(ctx context.Context, tableName string, it
318327

319328
return decryptedItem, nil
320329
}
321-
322-
// constructMaterialName constructs a material name based on an item's primary key.
323-
func (ec *EncryptedClient) constructMaterialName(item map[string]types.AttributeValue, pkInfo *utils.PrimaryKeyInfo) string {
324-
partitionKeyValue := item[pkInfo.PartitionKey].(*types.AttributeValueMemberS).Value
325-
sortKeyValue := ""
326-
if pkInfo.SortKey != "" && item[pkInfo.SortKey] != nil {
327-
sortKeyValue = item[pkInfo.SortKey].(*types.AttributeValueMemberS).Value
328-
}
329-
330-
rawMaterialName := pkInfo.Table + "-" + partitionKeyValue
331-
if sortKeyValue != "" {
332-
rawMaterialName += "-" + sortKeyValue
333-
}
334-
335-
return utils.HashString(rawMaterialName)
336-
}

‎pkg/utils/utils.go

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@ package utils
33
import (
44
"context"
55
"crypto/sha256"
6+
"encoding/base64"
67
"encoding/hex"
78
"encoding/json"
89
"fmt"
10+
"strconv"
911

1012
"github.com/aws/aws-sdk-go-v2/aws"
1113
"github.com/aws/aws-sdk-go-v2/service/dynamodb"
@@ -53,6 +55,52 @@ func HashString(input string) string {
5355
return hex.EncodeToString(hasher.Sum(nil))
5456
}
5557

58+
// AttributeValueToString converts a DynamoDB AttributeValue to a string representation.
59+
func AttributeValueToString(value types.AttributeValue) (string, error) {
60+
switch v := value.(type) {
61+
case *types.AttributeValueMemberS:
62+
return v.Value, nil
63+
case *types.AttributeValueMemberN:
64+
return v.Value, nil
65+
case *types.AttributeValueMemberB:
66+
return base64.StdEncoding.EncodeToString(v.Value), nil
67+
case *types.AttributeValueMemberBOOL:
68+
return strconv.FormatBool(v.Value), nil
69+
case *types.AttributeValueMemberNULL:
70+
return "", nil
71+
case *types.AttributeValueMemberM:
72+
m := make(map[string]string)
73+
for key, value := range v.Value {
74+
convertedValue, err := AttributeValueToString(value)
75+
if err != nil {
76+
return "", err
77+
}
78+
m[key] = convertedValue
79+
}
80+
jsonStr, err := json.Marshal(m)
81+
if err != nil {
82+
return "", fmt.Errorf("failed to marshal map to JSON: %v", err)
83+
}
84+
return string(jsonStr), nil
85+
case *types.AttributeValueMemberL:
86+
var l []string
87+
for _, listItem := range v.Value {
88+
convertedItem, err := AttributeValueToString(listItem)
89+
if err != nil {
90+
return "", err
91+
}
92+
l = append(l, convertedItem)
93+
}
94+
jsonStr, err := json.Marshal(l)
95+
if err != nil {
96+
return "", fmt.Errorf("failed to marshal list to JSON: %v", err)
97+
}
98+
return string(jsonStr), nil
99+
default:
100+
return "", fmt.Errorf("unsupported AttributeValue type: %T", value)
101+
}
102+
}
103+
56104
// AttributeValueMapToBytes converts a map of DynamoDB attribute values to a JSON byte slice.
57105
func AttributeValueMapToBytes(attributes map[string]types.AttributeValue) ([]byte, error) {
58106
// DynamoDB attribute values to a generic map interface{}
@@ -95,6 +143,29 @@ func BytesToAttributeValue(data []byte) (types.AttributeValue, error) {
95143
return interfaceToAttributeValue(av)
96144
}
97145

146+
// ConstructMaterialName constructs a material name based on an item's primary key.
147+
func ConstructMaterialName(item map[string]types.AttributeValue, pkInfo *PrimaryKeyInfo) (string, error) {
148+
partitionKeyValue, err := AttributeValueToString(item[pkInfo.PartitionKey])
149+
if err != nil {
150+
return "", fmt.Errorf("invalid partition key attribute type: %v", err)
151+
}
152+
153+
sortKeyValue := ""
154+
if pkInfo.SortKey != "" {
155+
sortKeyValue, err = AttributeValueToString(item[pkInfo.SortKey])
156+
if err != nil {
157+
return "", fmt.Errorf("invalid sort key attribute type: %v", err)
158+
}
159+
}
160+
161+
rawMaterialName := pkInfo.Table + "-" + partitionKeyValue
162+
if sortKeyValue != "" {
163+
rawMaterialName += "-" + sortKeyValue
164+
}
165+
166+
return HashString(rawMaterialName), nil
167+
}
168+
98169
// attributeValueToInterface converts DynamoDB's types.AttributeValue to a generic interface{}.
99170
func attributeValueToInterface(av types.AttributeValue) (interface{}, error) {
100171
switch v := av.(type) {

0 commit comments

Comments
 (0)