Skip to content

Commit 5243ef8

Browse files
ThirdEyeSqueegeeThirdEyeSqueegee
committed
feat: get secrets by versionId or versionStage (#19)
1 parent df10cc3 commit 5243ef8

File tree

5 files changed

+190
-9
lines changed

5 files changed

+190
-9
lines changed

‎src/main/java/com/amazonaws/secretsmanager/caching/SecretCache.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,26 @@ public String getSecretString(final String secretId) {
136136
return gsv.secretString();
137137
}
138138

139+
/**
140+
* Retrieve and cache a secret string from AWS Secrets Manager.
141+
*
142+
* @param secretId the secret ID of the desired secret.
143+
* @param versionId the version ID of the desired secret (optional, can be <c>null</c>).
144+
* @param versionStage the version stage of the desired secret (optional, can be <c>null</c>).
145+
*
146+
* @return The secret string for the desired secret.
147+
*/
148+
public String getSecretString(final String secretId, final String versionId, final String versionStage) {
149+
SecretCacheItem secret = this.getCachedSecret(secretId);
150+
GetSecretValueResponse gsv = secret.getSecretValue(versionId, versionStage);
151+
152+
if (gsv == null) {
153+
return null;
154+
}
155+
156+
return gsv.secretString();
157+
}
158+
139159
/**
140160
* Method to retrieve a binary secret from AWS Secrets Manager.
141161
*
@@ -151,6 +171,26 @@ public ByteBuffer getSecretBinary(final String secretId) {
151171
return gsv.secretBinary().asByteBuffer();
152172
}
153173

174+
/**
175+
* Retrieve and cache a secret binary from AWS Secrets Manager.
176+
*
177+
* @param secretId the secret ID of the desired secret.
178+
* @param versionId the version ID of the desired secret (optional, can be <c>null</c>).
179+
* @param versionStage the version stage of the desired secret (optional, can be <c>null</c>).
180+
*
181+
* @return The secret binary for the desired secret.
182+
*/
183+
public ByteBuffer getSecretBinary(final String secretId, final String versionId, final String versionStage) {
184+
SecretCacheItem secret = this.getCachedSecret(secretId);
185+
GetSecretValueResponse gsv = secret.getSecretValue(versionId, versionStage);
186+
187+
if (gsv == null) {
188+
return null;
189+
}
190+
191+
return gsv.secretBinary().asByteBuffer();
192+
}
193+
154194
/**
155195
* Method to force the refresh of a cached secret state.
156196
*

‎src/main/java/com/amazonaws/secretsmanager/caching/cache/SecretCacheItem.java‎

Lines changed: 48 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313

1414
package com.amazonaws.secretsmanager.caching.cache;
1515

16+
import java.util.List;
17+
import java.util.Map;
1618
import java.util.Objects;
1719
import java.util.Optional;
1820
import java.util.concurrent.ThreadLocalRandom;
@@ -115,16 +117,32 @@ protected DescribeSecretResponse executeRefresh() {
115117
* The result of the Describe Secret request to AWS Secrets Manager.
116118
* @return The cached secret version.
117119
*/
118-
private SecretCacheVersion getVersion(DescribeSecretResponse describeResponse) {
120+
private SecretCacheVersion getVersion(DescribeSecretResponse describeResponse, String versionId, String versionStage) {
119121
if (null == describeResponse) { return null; }
120122
if (null == describeResponse.versionIdsToStages()) { return null; }
121-
Optional<String> currentVersionId = describeResponse.versionIdsToStages().entrySet()
122-
.stream()
123-
.filter(Objects::nonNull)
124-
.filter(x -> x.getValue() != null)
125-
.filter(x -> x.getValue().contains(this.config.getVersionStage()))
126-
.map(x -> x.getKey())
127-
.findFirst();
123+
124+
Optional<String> currentVersionId = Optional.empty();
125+
126+
for (Map.Entry<String, List<String>> entry : describeResponse.versionIdsToStages().entrySet()) {
127+
if (entry == null) {
128+
continue;
129+
}
130+
131+
if (entry.getValue() == null) {
132+
continue;
133+
}
134+
135+
if (versionId != null && entry.getKey() == versionId) {
136+
currentVersionId = Optional.of(versionId);
137+
break;
138+
}
139+
140+
if ((versionStage != null && entry.getValue().contains(versionStage)) || entry.getValue().contains(config.getVersionStage())) {
141+
currentVersionId = Optional.of(entry.getKey());
142+
break;
143+
}
144+
}
145+
128146
if (currentVersionId.isPresent()) {
129147
SecretCacheVersion version = versions.get(currentVersionId.get());
130148
if (null == version) {
@@ -134,6 +152,7 @@ private SecretCacheVersion getVersion(DescribeSecretResponse describeResponse) {
134152
}
135153
return version;
136154
}
155+
137156
return null;
138157
}
139158

@@ -146,9 +165,29 @@ private SecretCacheVersion getVersion(DescribeSecretResponse describeResponse) {
146165
*/
147166
@Override
148167
protected GetSecretValueResponse getSecretValue(DescribeSecretResponse describeResponse) {
149-
SecretCacheVersion version = getVersion(describeResponse);
168+
SecretCacheVersion version = getVersion(describeResponse, null, null);
150169
if (null == version) { return null; }
151170
return version.getSecretValue();
152171
}
153172

173+
/**
174+
* Return the cached GetSecretValue result.
175+
*
176+
* @param describeResponse the DescribeSecret result.
177+
* @param versionId the version ID of the desired secret (optional, can be <c>null</c>).
178+
* @param versionStage the version stage of the desired secret (optional, can be <c>null</c>).
179+
*
180+
* @return The cached GetSecretValue result.
181+
*/
182+
@Override
183+
protected GetSecretValueResponse getSecretValue(DescribeSecretResponse describeResponse, String versionId, String versionStage) {
184+
SecretCacheVersion version = getVersion(describeResponse, versionId, versionStage);
185+
186+
if (version == null) {
187+
return null;
188+
}
189+
190+
return version.getSecretValue();
191+
}
192+
154193
}

‎src/main/java/com/amazonaws/secretsmanager/caching/cache/SecretCacheObject.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,17 @@ public SecretCacheObject(final String secretId,
118118
*/
119119
protected abstract GetSecretValueResponse getSecretValue(T result);
120120

121+
/**
122+
* Execute the actual refresh of the cached secret state.
123+
*
124+
* @param result the GetSecretValue or DescribeSecret result.
125+
* @param versionId the version ID of the desired secret (optional, can be <c>null</c>).
126+
* @param versionStage the version stage of the desired secret (optional, can be <c>null</c>).
127+
*
128+
* @return The cached GetSecretValue result based on the current cached state.
129+
*/
130+
protected abstract GetSecretValueResponse getSecretValue(T result, String versionId, String versionStage);;
131+
121132
public abstract boolean equals(Object obj);
122133
public abstract int hashCode();
123134
public abstract String toString();
@@ -252,4 +263,26 @@ public GetSecretValueResponse getSecretValue() {
252263
}
253264
}
254265

266+
/**
267+
* Return the cached GetSecretValue result.
268+
*
269+
* @param versionId the version ID of the desired secret (optional, can be <c>null</c>).
270+
* @param versionStage the version stage of the desired secret (optional, can be <c>null</c>).
271+
*
272+
* @return The cached GetSecretValue result.
273+
*/
274+
public GetSecretValueResponse getSecretValue(String versionId, String versionStage) {
275+
synchronized (lock) {
276+
refresh();
277+
278+
if (this.data == null) {
279+
if (this.exception != null) {
280+
throw this.exception;
281+
}
282+
}
283+
284+
return this.getSecretValue(this.getResult(), versionId, versionStage);
285+
}
286+
}
287+
255288
}

‎src/main/java/com/amazonaws/secretsmanager/caching/cache/SecretCacheVersion.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,4 +98,18 @@ protected GetSecretValueResponse getSecretValue(GetSecretValueResponse gsvResult
9898
return gsvResult;
9999
}
100100

101+
/**
102+
* Return the cached GetSecretValue result.
103+
*
104+
* @param gsvResult the GetSecretValue or DescribeSecret result.
105+
* @param versionId the version ID of the desired secret (optional, can be <c>null</c>).
106+
* @param versionStage the version stage of the desired secret (optional, can be <c>null</c>).
107+
*
108+
* @return The cached GetSecretValue result.
109+
*/
110+
@Override
111+
protected GetSecretValueResponse getSecretValue(GetSecretValueResponse gsvResult, String versionId, String versionStage) {
112+
return gsvResult;
113+
}
114+
101115
}

‎src/test/java/com/amazonaws/secretsmanager/caching/SecretCacheTest.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,61 @@ public void basicSecretCacheTest() {
111111
sc.close();
112112
}
113113

114+
@Test
115+
public void basicSecretCacheVersionIdTest() {
116+
final String secret = "basicSecretCacheTest";
117+
Map<String, List<String>> versionMap = new HashMap<String, List<String>>();
118+
versionMap.put("versionId", Arrays.asList("AWSCURRENT"));
119+
versionMap.put("otherVersionId", Arrays.asList("AWSCURRENT"));
120+
Mockito.when(describeSecretResponse.versionIdsToStages()).thenReturn(versionMap);
121+
GetSecretValueResponse.Builder resBuilder = GetSecretValueResponse.builder().secretString(secret)
122+
.secretBinary(SdkBytes.fromByteArray(secret.getBytes()));
123+
getSecretValueResponse = resBuilder.build();
124+
125+
Mockito.when(asm.describeSecret(Mockito.any(DescribeSecretRequest.class))).thenReturn(describeSecretResponse);
126+
Mockito.when(asm.getSecretValue(Mockito.any(GetSecretValueRequest.class))).thenReturn(getSecretValueResponse);
127+
128+
SecretCache sc = new SecretCache(asm);
129+
130+
// Request the secret multiple times and verify the correct result
131+
repeat(10, n -> Assert.assertEquals(sc.getSecretString("", "otherVersionId", null), secret));
132+
133+
// Verify that multiple requests did not call the API
134+
Mockito.verify(asm, Mockito.times(1)).describeSecret(Mockito.any(DescribeSecretRequest.class));
135+
Mockito.verify(asm, Mockito.times(1)).getSecretValue(Mockito.any(GetSecretValueRequest.class));
136+
137+
repeat(10, n -> Assert.assertEquals(sc.getSecretBinary("", "otherVersionId", null),
138+
ByteBuffer.wrap(secret.getBytes())));
139+
sc.close();
140+
}
141+
142+
@Test
143+
public void basicSecretCacheVersionStageTest() {
144+
final String secret = "basicSecretCacheTest";
145+
Map<String, List<String>> versionMap = new HashMap<String, List<String>>();
146+
versionMap.put("versionId", Arrays.asList("AWSCURRENT"));
147+
Mockito.when(describeSecretResponse.versionIdsToStages()).thenReturn(versionMap);
148+
GetSecretValueResponse.Builder resBuilder = GetSecretValueResponse.builder().secretString(secret)
149+
.secretBinary(SdkBytes.fromByteArray(secret.getBytes()));
150+
getSecretValueResponse = resBuilder.build();
151+
152+
Mockito.when(asm.describeSecret(Mockito.any(DescribeSecretRequest.class))).thenReturn(describeSecretResponse);
153+
Mockito.when(asm.getSecretValue(Mockito.any(GetSecretValueRequest.class))).thenReturn(getSecretValueResponse);
154+
155+
SecretCache sc = new SecretCache(asm);
156+
157+
// Request the secret multiple times and verify the correct result
158+
repeat(10, n -> Assert.assertEquals(sc.getSecretString("", null, "AWSCURRENT"), secret));
159+
160+
// Verify that multiple requests did not call the API
161+
Mockito.verify(asm, Mockito.times(1)).describeSecret(Mockito.any(DescribeSecretRequest.class));
162+
Mockito.verify(asm, Mockito.times(1)).getSecretValue(Mockito.any(GetSecretValueRequest.class));
163+
164+
repeat(10, n -> Assert.assertEquals(sc.getSecretBinary("", null, "AWSCURRENT"),
165+
ByteBuffer.wrap(secret.getBytes())));
166+
sc.close();
167+
}
168+
114169
@Test
115170
public void hookSecretCacheTest() {
116171
final String secret = "hookSecretCacheTest";

0 commit comments

Comments
 (0)