Skip to content

Commit ccb1ede

Browse files
author
Bryan Donlan
committed
Restore logic to add in UA suffix
1 parent 8401ba3 commit ccb1ede

File tree

3 files changed

+82
-6
lines changed

3 files changed

+82
-6
lines changed

‎src/main/java/com/amazonaws/encryptionsdk/kms/KmsMasterKey.java‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.Map;
2424

2525
import com.amazonaws.AmazonServiceException;
26+
import com.amazonaws.AmazonWebServiceRequest;
2627
import com.amazonaws.auth.AWSCredentials;
2728
import com.amazonaws.auth.AWSCredentialsProvider;
2829
import com.amazonaws.encryptionsdk.AwsCrypto;
@@ -33,6 +34,7 @@
3334
import com.amazonaws.encryptionsdk.MasterKeyProvider;
3435
import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
3536
import com.amazonaws.encryptionsdk.exception.UnsupportedProviderException;
37+
import com.amazonaws.encryptionsdk.internal.VersionInfo;
3638
import com.amazonaws.services.kms.AWSKMS;
3739
import com.amazonaws.services.kms.model.DecryptRequest;
3840
import com.amazonaws.services.kms.model.DecryptResult;
@@ -51,6 +53,12 @@ public final class KmsMasterKey extends MasterKey<KmsMasterKey> implements KmsMe
5153
private final String id_;
5254
private final List<String> grantTokens_ = new ArrayList<>();
5355

56+
private <T extends AmazonWebServiceRequest> T updateUserAgent(T request) {
57+
request.getRequestClientOptions().appendUserAgent(VersionInfo.USER_AGENT);
58+
59+
return request;
60+
}
61+
5462
/**
5563
*
5664
* @deprecated Use a {@link KmsMasterKeyProvider} to obtain {@link KmsMasterKey}s.
@@ -93,13 +101,13 @@ public String getKeyId() {
93101
@Override
94102
public DataKey<KmsMasterKey> generateDataKey(final CryptoAlgorithm algorithm,
95103
final Map<String, String> encryptionContext) {
96-
final GenerateDataKeyResult gdkResult = kms_.generateDataKey(
104+
final GenerateDataKeyResult gdkResult = kms_.generateDataKey(updateUserAgent(
97105
new GenerateDataKeyRequest()
98106
.withKeyId(getKeyId())
99107
.withNumberOfBytes(algorithm.getDataKeyLength())
100108
.withEncryptionContext(encryptionContext)
101109
.withGrantTokens(grantTokens_)
102-
);
110+
));
103111
final byte[] rawKey = new byte[algorithm.getDataKeyLength()];
104112
gdkResult.getPlaintext().get(rawKey);
105113
if (gdkResult.getPlaintext().remaining() > 0) {
@@ -137,12 +145,12 @@ public DataKey<KmsMasterKey> encryptDataKey(final CryptoAlgorithm algorithm,
137145
throw new IllegalArgumentException("Only RAW encoded keys are supported");
138146
}
139147
try {
140-
final EncryptResult encryptResult = kms_.encrypt(
148+
final EncryptResult encryptResult = kms_.encrypt(updateUserAgent(
141149
new EncryptRequest()
142150
.withKeyId(id_)
143151
.withPlaintext(ByteBuffer.wrap(key.getEncoded()))
144152
.withEncryptionContext(encryptionContext)
145-
.withGrantTokens(grantTokens_));
153+
.withGrantTokens(grantTokens_)));
146154
final byte[] edk = new byte[encryptResult.getCiphertextBlob().remaining()];
147155
encryptResult.getCiphertextBlob().get(edk);
148156
return new DataKey<>(dataKey.getKey(), edk, encryptResult.getKeyId().getBytes(StandardCharsets.UTF_8), this);
@@ -159,11 +167,11 @@ public DataKey<KmsMasterKey> decryptDataKey(final CryptoAlgorithm algorithm,
159167
final List<Exception> exceptions = new ArrayList<>();
160168
for (final EncryptedDataKey edk : encryptedDataKeys) {
161169
try {
162-
final DecryptResult decryptResult = kms_.decrypt(
170+
final DecryptResult decryptResult = kms_.decrypt(updateUserAgent(
163171
new DecryptRequest()
164172
.withCiphertextBlob(ByteBuffer.wrap(edk.getEncryptedDataKey()))
165173
.withEncryptionContext(encryptionContext)
166-
.withGrantTokens(grantTokens_));
174+
.withGrantTokens(grantTokens_)));
167175
if (decryptResult.getKeyId().equals(id_)) {
168176
final byte[] rawKey = new byte[algorithm.getDataKeyLength()];
169177
decryptResult.getPlaintext().get(rawKey);

‎src/test/java/com/amazonaws/services/kms/KMSProviderBuilderIntegrationTests.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,12 @@
1313
import java.util.Arrays;
1414

1515
import org.junit.Test;
16+
import org.mockito.ArgumentCaptor;
1617

1718
import com.amazonaws.AbortedException;
19+
import com.amazonaws.AmazonWebServiceRequest;
1820
import com.amazonaws.ClientConfiguration;
21+
import com.amazonaws.Request;
1922
import com.amazonaws.auth.AWSCredentials;
2023
import com.amazonaws.auth.AWSCredentialsProvider;
2124
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
@@ -24,6 +27,7 @@
2427
import com.amazonaws.encryptionsdk.CryptoResult;
2528
import com.amazonaws.encryptionsdk.MasterKeyProvider;
2629
import com.amazonaws.encryptionsdk.exception.CannotUnwrapDataKeyException;
30+
import com.amazonaws.encryptionsdk.internal.VersionInfo;
2731
import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
2832
import com.amazonaws.handlers.RequestHandler2;
2933
import com.amazonaws.http.exception.HttpRequestTimeoutException;
@@ -197,6 +201,34 @@ public void whenBogusEndpointIsSet_constructionFails() throws Exception {
197201
);
198202
}
199203

204+
@Test
205+
public void whenUserAgentsOverridden_originalUAsPreserved() throws Exception {
206+
RequestHandler2 handler = spy(new RequestHandler2() {});
207+
208+
KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
209+
.withClientBuilder(
210+
AWSKMSClientBuilder.standard().withRequestHandlers(handler)
211+
.withClientConfiguration(
212+
new ClientConfiguration()
213+
.withUserAgentPrefix("TEST-UA-PREFIX")
214+
.withUserAgentSuffix("TEST-UA-SUFFIX")
215+
)
216+
)
217+
.withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
218+
.clone().build();
219+
220+
new AwsCrypto().encryptData(mkp, new byte[0]);
221+
222+
ArgumentCaptor<Request> captor = ArgumentCaptor.forClass(Request.class);
223+
verify(handler, atLeastOnce()).beforeRequest(captor.capture());
224+
225+
String ua = (String)captor.getValue().getHeaders().get("User-Agent");
226+
227+
assertTrue(ua.contains("TEST-UA-PREFIX"));
228+
assertTrue(ua.contains("TEST-UA-SUFFIX"));
229+
assertTrue(ua.contains(VersionInfo.USER_AGENT));
230+
}
231+
200232
@Test
201233
public void whenDefaultRegionSet_itIsUsedForBareKeyIds() throws Exception {
202234
// TODO: Need to set up a role to assume as bare key IDs are relative to the caller account

‎src/test/java/com/amazonaws/services/kms/KMSProviderBuilderMockTests.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,11 @@
2323
import org.junit.Test;
2424
import org.mockito.ArgumentCaptor;
2525

26+
import com.amazonaws.AmazonWebServiceRequest;
27+
import com.amazonaws.RequestClientOptions;
2628
import com.amazonaws.encryptionsdk.AwsCrypto;
2729
import com.amazonaws.encryptionsdk.MasterKeyProvider;
30+
import com.amazonaws.encryptionsdk.internal.VersionInfo;
2831
import com.amazonaws.encryptionsdk.kms.KmsMasterKey;
2932
import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
3033
import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider.RegionalClientSupplier;
@@ -157,4 +160,37 @@ public void testLegacyGrantTokenPassthrough() throws Exception {
157160
assertFalse(grantTokens.contains("x"));
158161
assertFalse(grantTokens.contains("z"));
159162
}
163+
164+
@Test
165+
public void testUserAgentPassthrough() throws Exception {
166+
MockKMSClient client = spy(new MockKMSClient());
167+
168+
String key1 = client.createKey().getKeyMetadata().getArn();
169+
String key2 = client.createKey().getKeyMetadata().getArn();
170+
171+
KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
172+
.withKeysForEncryption(key1, key2)
173+
.withCustomClientFactory(ignored -> client)
174+
.build();
175+
176+
new AwsCrypto().decryptData(mkp, new AwsCrypto().encryptData(mkp, new byte[0]).getResult());
177+
178+
ArgumentCaptor<GenerateDataKeyRequest> gdkr = ArgumentCaptor.forClass(GenerateDataKeyRequest.class);
179+
verify(client, times(1)).generateDataKey(gdkr.capture());
180+
assertTrue(getUA(gdkr.getValue()).contains(VersionInfo.USER_AGENT));
181+
182+
ArgumentCaptor<EncryptRequest> encr = ArgumentCaptor.forClass(EncryptRequest.class);
183+
verify(client, times(1)).encrypt(encr.capture());
184+
assertTrue(getUA(encr.getValue()).contains(VersionInfo.USER_AGENT));
185+
186+
ArgumentCaptor<DecryptRequest> decr = ArgumentCaptor.forClass(DecryptRequest.class);
187+
verify(client, times(1)).decrypt(decr.capture());
188+
assertTrue(getUA(decr.getValue()).contains(VersionInfo.USER_AGENT));
189+
}
190+
191+
private String getUA(AmazonWebServiceRequest request) {
192+
// Note: This test may break in future versions of the AWS SDK, as Marker is documented as being for internal
193+
// use only.
194+
return request.getRequestClientOptions().getClientMarker(RequestClientOptions.Marker.USER_AGENT);
195+
}
160196
}

0 commit comments

Comments
 (0)