Skip to content

Commit 4937c09

Browse files
committed
Improve sigv4 and s3express
1 parent 48c7528 commit 4937c09

8 files changed

Lines changed: 344 additions & 36 deletions

File tree

aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityCache.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,17 @@ private void evictLeastRecentlyUsed() {
9696
}
9797
}
9898

99+
void invalidateAll() {
100+
writeLock.lock();
101+
try {
102+
for (var entry : entries.values()) {
103+
entry.resolver.invalidate();
104+
}
105+
} finally {
106+
writeLock.unlock();
107+
}
108+
}
109+
99110
@Override
100111
public void close() {
101112
writeLock.lock();

aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProvider.java

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
package software.amazon.smithy.java.aws.client.auth.scheme.s3express;
77

88
import java.util.Objects;
9+
import java.util.concurrent.Executors;
10+
import java.util.concurrent.ScheduledExecutorService;
11+
import java.util.concurrent.atomic.AtomicInteger;
912
import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver;
1013
import software.amazon.smithy.java.auth.api.identity.IdentityResolver;
1114
import software.amazon.smithy.java.auth.api.identity.IdentityResult;
@@ -32,6 +35,13 @@
3235
*/
3336
public final class S3ExpressIdentityProvider implements IdentityResolver<S3ExpressIdentity> {
3437

38+
private static final AtomicInteger THREAD_COUNTER = new AtomicInteger();
39+
private static final ScheduledExecutorService REFRESH_EXECUTOR = Executors.newSingleThreadScheduledExecutor(r -> {
40+
Thread t = new Thread(r, "smithy-s3express-session-refresh-" + THREAD_COUNTER.incrementAndGet());
41+
t.setDaemon(true);
42+
return t;
43+
});
44+
3545
private final IdentityResolver<AwsCredentialsIdentity> baseIdentityProvider;
3646
private final S3ExpressIdentityCache cache;
3747

@@ -50,7 +60,7 @@ public S3ExpressIdentityProvider(
5060
this.cache = new S3ExpressIdentityCache(key -> buildPerBucketResolver(key, createSession));
5161
}
5262

53-
private static CachingIdentityResolver<S3ExpressIdentity> buildPerBucketResolver(
63+
private CachingIdentityResolver<S3ExpressIdentity> buildPerBucketResolver(
5464
S3ExpressIdentityKey key,
5565
CreateSessionCallback createSession
5666
) {
@@ -69,7 +79,9 @@ public Class<S3ExpressIdentity> identityType() {
6979
return S3ExpressIdentity.class;
7080
}
7181
};
72-
return CachingIdentityResolver.builder(delegate).build();
82+
return CachingIdentityResolver.builder(delegate)
83+
.executor(REFRESH_EXECUTOR)
84+
.build();
7385
}
7486

7587
@Override
@@ -85,9 +97,13 @@ public IdentityResult<S3ExpressIdentity> resolveIdentity(Context requestProperti
8597
IdentityResult<AwsCredentialsIdentity> baseResult = baseIdentityProvider.resolveIdentity(requestProperties);
8698
AwsCredentialsIdentity base = baseResult.identity();
8799
if (base == null) {
100+
String baseError = baseResult.error();
101+
Class<?> resolver = baseResult.resolver() == null ? getClass() : baseResult.resolver();
88102
return IdentityResult.ofError(
89-
getClass(),
90-
"Could not resolve base AWS credentials for S3 Express CreateSession");
103+
resolver,
104+
baseError == null
105+
? "Could not resolve base AWS credentials for S3 Express CreateSession"
106+
: "Could not resolve base AWS credentials for S3 Express CreateSession: " + baseError);
91107
}
92108

93109
var perBucket = cache.get(new S3ExpressIdentityKey(bucket, base));
@@ -101,8 +117,6 @@ public Class<S3ExpressIdentity> identityType() {
101117

102118
@Override
103119
public void invalidate() {
104-
// Per-bucket resolvers manage their own validity; punting on global invalidation for now.
105-
// If/when an auth-failure interceptor needs to invalidate, we can extend this to clear
106-
// the bucket entry for the call's specific (bucket, identity) key.
120+
cache.invalidateAll();
107121
}
108122
}

aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3VirtualHostStyleInterceptor.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ public <RequestT> RequestT modifyBeforeSigning(RequestHook<?, ?, RequestT> hook)
5656
if (!path.startsWith(prefix)) {
5757
return hook.request();
5858
}
59+
if (path.length() > prefix.length() && path.charAt(prefix.length()) != '/') {
60+
return hook.request();
61+
}
5962
String rest = path.substring(prefix.length());
6063
// "/<bucket>" → "/" (CreateSession-shaped), "/<bucket>/key" → "/key".
6164
if (rest.isEmpty()) {
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.aws.client.auth.scheme.s3express;
7+
8+
import static org.hamcrest.MatcherAssert.assertThat;
9+
import static org.hamcrest.Matchers.containsString;
10+
import static org.hamcrest.Matchers.equalTo;
11+
import static org.hamcrest.Matchers.is;
12+
13+
import java.util.concurrent.atomic.AtomicInteger;
14+
import org.junit.jupiter.api.Test;
15+
import software.amazon.smithy.java.auth.api.identity.IdentityResolver;
16+
import software.amazon.smithy.java.auth.api.identity.IdentityResult;
17+
import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity;
18+
import software.amazon.smithy.java.context.Context;
19+
20+
class S3ExpressIdentityProviderTest {
21+
22+
@Test
23+
void invalidateClearsCachedSessions() {
24+
var sessionCount = new AtomicInteger();
25+
var provider = new S3ExpressIdentityProvider(
26+
staticBaseIdentity(),
27+
(bucket, baseCredentials) -> S3ExpressIdentity.create(
28+
"SESSION-" + sessionCount.incrementAndGet(),
29+
"secret",
30+
"token",
31+
null));
32+
var context = contextWithBucket("bucket");
33+
34+
var first = provider.resolveIdentity(context).identity();
35+
var cached = provider.resolveIdentity(context).identity();
36+
provider.invalidate();
37+
var refreshed = provider.resolveIdentity(context).identity();
38+
39+
assertThat(cached, is(first));
40+
assertThat(refreshed.accessKeyId(), equalTo("SESSION-2"));
41+
assertThat(sessionCount.get(), equalTo(2));
42+
}
43+
44+
@Test
45+
void preservesBaseIdentityError() {
46+
IdentityResolver<AwsCredentialsIdentity> failingBase = new IdentityResolver<>() {
47+
@Override
48+
public IdentityResult<AwsCredentialsIdentity> resolveIdentity(Context requestProperties) {
49+
return IdentityResult.ofError(getClass(), "missing base creds");
50+
}
51+
52+
@Override
53+
public Class<AwsCredentialsIdentity> identityType() {
54+
return AwsCredentialsIdentity.class;
55+
}
56+
};
57+
var provider = new S3ExpressIdentityProvider(
58+
failingBase,
59+
(bucket, baseCredentials) -> {
60+
throw new AssertionError("CreateSession should not be called");
61+
});
62+
63+
var result = provider.resolveIdentity(contextWithBucket("bucket"));
64+
65+
assertThat(result.resolver(), equalTo(failingBase.getClass()));
66+
assertThat(result.error(), containsString("missing base creds"));
67+
}
68+
69+
private static IdentityResolver<AwsCredentialsIdentity> staticBaseIdentity() {
70+
return new IdentityResolver<>() {
71+
@Override
72+
public IdentityResult<AwsCredentialsIdentity> resolveIdentity(Context requestProperties) {
73+
return IdentityResult.of(AwsCredentialsIdentity.create("AKID", "secret"));
74+
}
75+
76+
@Override
77+
public Class<AwsCredentialsIdentity> identityType() {
78+
return AwsCredentialsIdentity.class;
79+
}
80+
};
81+
}
82+
83+
private static Context contextWithBucket(String bucket) {
84+
var context = Context.create();
85+
context.put(S3ExpressContext.BUCKET, bucket);
86+
return context;
87+
}
88+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.aws.client.auth.scheme.s3express;
7+
8+
import static org.hamcrest.MatcherAssert.assertThat;
9+
import static org.hamcrest.Matchers.equalTo;
10+
import static org.hamcrest.Matchers.sameInstance;
11+
12+
import java.net.URI;
13+
import java.util.List;
14+
import org.junit.jupiter.api.Test;
15+
import software.amazon.smithy.java.client.core.interceptors.RequestHook;
16+
import software.amazon.smithy.java.context.Context;
17+
import software.amazon.smithy.java.core.schema.ApiOperation;
18+
import software.amazon.smithy.java.core.schema.ApiService;
19+
import software.amazon.smithy.java.core.schema.Schema;
20+
import software.amazon.smithy.java.core.schema.SerializableStruct;
21+
import software.amazon.smithy.java.core.schema.ShapeBuilder;
22+
import software.amazon.smithy.java.core.serde.ShapeSerializer;
23+
import software.amazon.smithy.java.core.serde.TypeRegistry;
24+
import software.amazon.smithy.java.http.api.HttpRequest;
25+
import software.amazon.smithy.java.http.api.HttpVersion;
26+
import software.amazon.smithy.model.shapes.ShapeId;
27+
28+
class S3VirtualHostStyleInterceptorTest {
29+
30+
@Test
31+
void stripsBucketOnlyAtPathSegmentBoundary() {
32+
var request = request("/bucket/key");
33+
var result = intercept(request, "bucket");
34+
35+
assertThat(result.uri().getPath(), equalTo("/key"));
36+
}
37+
38+
@Test
39+
void doesNotStripBucketPrefixFromLongerPathSegment() {
40+
var request = request("/bucket-extra/key");
41+
var result = intercept(request, "bucket");
42+
43+
assertThat(result, sameInstance(request));
44+
}
45+
46+
@Test
47+
void stripsBucketOnlyPathToRoot() {
48+
var result = intercept(request("/bucket"), "bucket");
49+
50+
assertThat(result.uri().getPath(), equalTo("/"));
51+
}
52+
53+
private static HttpRequest intercept(HttpRequest request, String bucket) {
54+
var context = Context.create();
55+
context.put(S3ExpressContext.BUCKET, bucket);
56+
var hook = new RequestHook<>(operation(), context, new Input(), request);
57+
return S3VirtualHostStyleInterceptor.INSTANCE.modifyBeforeSigning(hook);
58+
}
59+
60+
private static HttpRequest request(String path) {
61+
return HttpRequest.create()
62+
.setMethod("GET")
63+
.setHttpVersion(HttpVersion.HTTP_1_1)
64+
.setUri(URI.create("https://example.com" + path))
65+
.toUnmodifiable();
66+
}
67+
68+
private static ApiOperation<Input, Input> operation() {
69+
return new ApiOperation<>() {
70+
@Override
71+
public ShapeBuilder<Input> inputBuilder() {
72+
return null;
73+
}
74+
75+
@Override
76+
public ShapeBuilder<Input> outputBuilder() {
77+
return null;
78+
}
79+
80+
@Override
81+
public Schema schema() {
82+
return null;
83+
}
84+
85+
@Override
86+
public Schema inputSchema() {
87+
return null;
88+
}
89+
90+
@Override
91+
public Schema outputSchema() {
92+
return null;
93+
}
94+
95+
@Override
96+
public TypeRegistry errorRegistry() {
97+
return null;
98+
}
99+
100+
@Override
101+
public List<ShapeId> effectiveAuthSchemes() {
102+
return List.of();
103+
}
104+
105+
@Override
106+
public List<Schema> errorSchemas() {
107+
return List.of();
108+
}
109+
110+
@Override
111+
public ApiService service() {
112+
return null;
113+
}
114+
};
115+
}
116+
117+
private static final class Input implements SerializableStruct {
118+
@Override
119+
public Schema schema() {
120+
return null;
121+
}
122+
123+
@Override
124+
public void serializeMembers(ShapeSerializer serializer) {}
125+
126+
@Override
127+
public <T> T getMemberValue(Schema member) {
128+
return null;
129+
}
130+
}
131+
}

0 commit comments

Comments
 (0)