Skip to content

Commit 3bb04c6

Browse files
rishav-karanjitrishav-karanjit
committed
m
1 parent 4a67f17 commit 3bb04c6

13 files changed

Lines changed: 712 additions & 82 deletions

File tree

‎src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/DynamoDbEncryption.java‎

Lines changed: 31 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -123,12 +123,23 @@ private static final class DdbBranchKeyIdSupplierAdapter implements IBranchKeyId
123123

124124
@Override
125125
public GetBranchKeyIdOutput GetBranchKeyId(GetBranchKeyIdInput input) {
126-
// The encryption context contains the DDB key attributes
127126
Map<String, String> ec = input.encryptionContext();
128-
// Convert EC string values to DDB AttributeValues for the DDB supplier
127+
// Extract DDB key attributes from the encryption context
129128
Map<String, AttributeValue> ddbKey = new java.util.HashMap<>();
130-
for (Map.Entry<String, String> entry : ec.entrySet()) {
131-
ddbKey.put(entry.getKey(), AttributeValue.fromS(entry.getValue()));
129+
String partitionKeyName = ec.get("aws-crypto-partition-name");
130+
String sortKeyName = ec.get("aws-crypto-sort-name");
131+
String attrPrefix = "aws-crypto-attr.";
132+
if (partitionKeyName != null) {
133+
String encoded = ec.get(attrPrefix + partitionKeyName);
134+
if (encoded != null) {
135+
ddbKey.put(partitionKeyName, decodeAttrFromContext(encoded));
136+
}
137+
}
138+
if (sortKeyName != null) {
139+
String encoded = ec.get(attrPrefix + sortKeyName);
140+
if (encoded != null) {
141+
ddbKey.put(sortKeyName, decodeAttrFromContext(encoded));
142+
}
132143
}
133144
software.amazon.cryptography.dbencryptionsdk.dynamodb.model.GetBranchKeyIdFromDdbKeyOutput result =
134145
ddbSupplier.GetBranchKeyIdFromDdbKey(
@@ -139,6 +150,22 @@ public GetBranchKeyIdOutput GetBranchKeyId(GetBranchKeyIdInput input) {
139150
.branchKeyId(result.branchKeyId())
140151
.build();
141152
}
153+
154+
private static AttributeValue decodeAttrFromContext(String encoded) {
155+
byte[] decoded = java.util.Base64.getDecoder().decode(encoded);
156+
// First 2 bytes are typeId, rest is value
157+
byte[] typeId = new byte[] { decoded[0], decoded[1] };
158+
byte[] value = new byte[decoded.length - 2];
159+
System.arraycopy(decoded, 2, value, 0, value.length);
160+
// 0x0001 = String, 0x0002 = Number, 0x00FF = Binary
161+
if (typeId[0] == 0x00 && typeId[1] == 0x01) {
162+
return AttributeValue.fromS(new String(value, StandardCharsets.UTF_8));
163+
} else if (typeId[0] == 0x00 && typeId[1] == 0x02) {
164+
return AttributeValue.fromN(new String(value, StandardCharsets.UTF_8));
165+
} else {
166+
return AttributeValue.fromB(software.amazon.awssdk.core.SdkBytes.fromByteArray(value));
167+
}
168+
}
142169
}
143170

144171
public interface Builder {

‎src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/internal/BeaconConfigResolver.java‎

Lines changed: 92 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import java.util.ArrayList;
66
import java.util.Collections;
77
import java.util.List;
8+
import java.util.Map;
89
import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.BeaconVersion;
910
import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.CompoundBeacon;
1011
import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.EncryptedPart;
@@ -24,12 +25,35 @@ public static final class SearchInfo {
2425
public final List<StandardBeaconImpl> standardBeacons;
2526
public final List<CompoundBeaconImpl> compoundBeacons;
2627
public final BeaconVersion beaconVersion;
28+
public final software.amazon.cryptography.keystore.KeyStore keyStore;
29+
public final String beaconKeyId;
2730

2831
public SearchInfo(List<StandardBeaconImpl> standardBeacons, List<CompoundBeaconImpl> compoundBeacons,
2932
BeaconVersion beaconVersion) {
3033
this.standardBeacons = standardBeacons;
3134
this.compoundBeacons = compoundBeacons;
3235
this.beaconVersion = beaconVersion;
36+
this.keyStore = beaconVersion.keyStore();
37+
// Extract key ID from key source
38+
if (beaconVersion.keySource() != null && beaconVersion.keySource().single() != null) {
39+
this.beaconKeyId = beaconVersion.keySource().single().keyId();
40+
} else {
41+
this.beaconKeyId = null;
42+
}
43+
}
44+
45+
/** Get the HMAC key bytes from the beacon key store. Returns null if not configured. */
46+
public byte[] getHmacKey() {
47+
if (keyStore == null || beaconKeyId == null) return null;
48+
software.amazon.cryptography.keystore.model.GetBeaconKeyOutput output =
49+
keyStore.GetBeaconKey(software.amazon.cryptography.keystore.model.GetBeaconKeyInput.builder()
50+
.branchKeyIdentifier(beaconKeyId).build());
51+
java.nio.ByteBuffer keyBuf = output.beaconKeyMaterials().beaconKey();
52+
if (keyBuf == null) return null;
53+
keyBuf = keyBuf.asReadOnlyBuffer();
54+
byte[] key = new byte[keyBuf.remaining()];
55+
keyBuf.get(key);
56+
return key;
3357
}
3458
}
3559

@@ -53,7 +77,11 @@ private static List<StandardBeaconImpl> resolveStandardBeacons(BeaconVersion ver
5377
if (version.standardBeacons() == null) return Collections.emptyList();
5478
List<StandardBeaconImpl> result = new ArrayList<>();
5579
for (StandardBeacon sb : version.standardBeacons()) {
56-
result.add(new StandardBeaconImpl(sb.name(), sb.length()));
80+
boolean isSet = false;
81+
if (sb.style() != null) {
82+
isSet = sb.style().asSet() != null || sb.style().sharedSet() != null;
83+
}
84+
result.add(new StandardBeaconImpl(sb.name(), sb.length(), isSet));
5785
}
5886
return result;
5987
}
@@ -62,25 +90,84 @@ private static List<CompoundBeaconImpl> resolveCompoundBeacons(
6290
BeaconVersion version, List<StandardBeaconImpl> standards
6391
) {
6492
if (version.compoundBeacons() == null) return Collections.emptyList();
93+
94+
// Build a map of encrypted parts from both version-level and compound-level
95+
Map<String, EncryptedPart> versionEncParts = new java.util.HashMap<>();
96+
if (version.encryptedParts() != null) {
97+
for (EncryptedPart ep : version.encryptedParts()) {
98+
versionEncParts.put(ep.name(), ep);
99+
}
100+
}
101+
Map<String, SignedPart> versionSigParts = new java.util.HashMap<>();
102+
if (version.signedParts() != null) {
103+
for (SignedPart sp : version.signedParts()) {
104+
versionSigParts.put(sp.name(), sp);
105+
}
106+
}
107+
65108
List<CompoundBeaconImpl> result = new ArrayList<>();
66109
for (CompoundBeacon cb : version.compoundBeacons()) {
67110
char split = cb.split().charAt(0);
68111
List<CompoundBeaconImpl.EncryptedPart> encParts = new ArrayList<>();
69-
if (cb.encrypted() != null) {
112+
List<CompoundBeaconImpl.SignedPart> sigParts = new ArrayList<>();
113+
114+
// Use compound-level parts if defined, otherwise use version-level parts via constructors
115+
if (cb.encrypted() != null && !cb.encrypted().isEmpty()) {
70116
for (EncryptedPart ep : cb.encrypted()) {
71117
StandardBeaconImpl beacon = findBeacon(standards, ep.name());
72118
if (beacon != null) {
73119
encParts.add(new CompoundBeaconImpl.EncryptedPart(ep.prefix(), beacon));
74120
}
75121
}
76122
}
77-
List<CompoundBeaconImpl.SignedPart> sigParts = new ArrayList<>();
78-
if (cb.signed() != null) {
123+
if (cb.signed() != null && !cb.signed().isEmpty()) {
79124
for (SignedPart sp : cb.signed()) {
80125
sigParts.add(new CompoundBeaconImpl.SignedPart(sp.prefix(), sp.name()));
81126
}
82127
}
83-
result.add(new CompoundBeaconImpl(cb.name(), split, encParts, sigParts));
128+
129+
// If no parts from compound-level, resolve from constructors using version-level parts
130+
if (encParts.isEmpty() && sigParts.isEmpty() && cb.constructors() != null) {
131+
java.util.Set<String> seenEncParts = new java.util.HashSet<>();
132+
java.util.Set<String> seenSigParts = new java.util.HashSet<>();
133+
for (software.amazon.cryptography.dbencryptionsdk.dynamodb.model.Constructor ctor : cb.constructors()) {
134+
if (ctor.parts() != null) {
135+
for (software.amazon.cryptography.dbencryptionsdk.dynamodb.model.ConstructorPart cp : ctor.parts()) {
136+
String partName = cp.name();
137+
// Check if it's an encrypted part
138+
EncryptedPart ep = versionEncParts.get(partName);
139+
if (ep != null && seenEncParts.add(partName)) {
140+
StandardBeaconImpl beacon = findBeacon(standards, ep.name());
141+
if (beacon != null) {
142+
encParts.add(new CompoundBeaconImpl.EncryptedPart(ep.prefix(), beacon));
143+
}
144+
} else {
145+
// Check if it's a signed part
146+
SignedPart sp = versionSigParts.get(partName);
147+
if (sp != null && seenSigParts.add(partName)) {
148+
sigParts.add(new CompoundBeaconImpl.SignedPart(sp.prefix(), sp.name()));
149+
}
150+
}
151+
}
152+
}
153+
}
154+
}
155+
156+
// Build constructor part name lists for ordered construction
157+
List<List<String>> ctorPartNames = new ArrayList<>();
158+
if (cb.constructors() != null) {
159+
for (software.amazon.cryptography.dbencryptionsdk.dynamodb.model.Constructor ctor : cb.constructors()) {
160+
if (ctor.parts() != null) {
161+
List<String> partNames = new ArrayList<>();
162+
for (software.amazon.cryptography.dbencryptionsdk.dynamodb.model.ConstructorPart cp : ctor.parts()) {
163+
partNames.add(cp.name());
164+
}
165+
ctorPartNames.add(partNames);
166+
}
167+
}
168+
}
169+
170+
result.add(new CompoundBeaconImpl(cb.name(), split, encParts, sigParts, ctorPartNames));
84171
}
85172
return result;
86173
}

0 commit comments

Comments
 (0)