Skip to content

Commit

Permalink
feat: Parse endpoints from auth token (#78)
Browse files Browse the repository at this point in the history
* feat: Parse endpoints from auth token

* add code comment

* pull out endpoint resolution to a helper

* update messages, clean up imports

* fix formatting
  • Loading branch information
gautamomento committed Oct 6, 2021
1 parent 9ac6ca7 commit 0a3b890
Show file tree
Hide file tree
Showing 7 changed files with 228 additions and 30 deletions.
1 change: 1 addition & 0 deletions build.gradle.kts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
ext["grpcVersion"] = "1.39.0"
ext["protobufVersion"] = "3.17.3"
ext["opentelemetryVersion"] = "1.4.1"
ext["jwtVersion"] = "0.11.2"

allprojects {
// These fields are used by the artifactoryPublish task
Expand Down
6 changes: 6 additions & 0 deletions momento-sdk/build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ plugins {
}

val opentelemetryVersion = rootProject.ext["opentelemetryVersion"]
val jwtVersion = rootProject.ext["jwtVersion"]

dependencies {
testImplementation("org.junit.jupiter:junit-jupiter:5.7.1")
Expand All @@ -16,6 +17,11 @@ dependencies {
implementation("io.opentelemetry:opentelemetry-exporter-otlp:$opentelemetryVersion")
implementation("io.grpc:grpc-netty:${rootProject.ext["grpcVersion"]}")

// For Auth token
implementation("io.jsonwebtoken:jjwt-api:$jwtVersion")
runtimeOnly("io.jsonwebtoken:jjwt-impl:$jwtVersion")
runtimeOnly("io.jsonwebtoken:jjwt-jackson:$jwtVersion")

// Internal Deps -------------------
implementation(project(":messages"))
}
Expand Down
48 changes: 48 additions & 0 deletions momento-sdk/src/intTest/java/momento/sdk/AuthTokenParserTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
package momento.sdk;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;

import momento.sdk.exceptions.ClientSdkException;
import org.junit.jupiter.api.Test;

final class AuthTokenParserTest {

// These secrets have botched up signature section, so should be okay to have them in source
// control.
private static final String TEST_AUTH_TOKEN_NO_ENDPOINT =
"eyJhbGciOiJIUzUxMiJ9.eyJzdWIiOiJpbnRlZ3JhdGlvbiJ9.ZOgkTs";
private static final String TEST_AUTH_TOKEN_ENDPOINT =
"eyJhbGciOiJIUzUxMiJ9.eyJzdWIiOiJzcXVpcnJlbCIsImNwIjoiY29udHJvbCBwbGFuZSBlbmRwb2ludCIsImMiOiJkYXRhIHBsYW5lIGVuZHBvaW50In0.zsTsEXFawetTCZI";

@Test
public void shouldParseAuthTokenWithNoEndpoints() {
AuthTokenParser.Claims claims = AuthTokenParser.parse(TEST_AUTH_TOKEN_NO_ENDPOINT);
assertFalse(claims.cacheEndpoint().isPresent());
assertFalse(claims.controlEndpoint().isPresent());
}

@Test
public void shouldParseAuthTokenWithEndpoints() {
AuthTokenParser.Claims claims = AuthTokenParser.parse(TEST_AUTH_TOKEN_ENDPOINT);
assertEquals("control plane endpoint", claims.controlEndpoint().get());
assertEquals("data plane endpoint", claims.cacheEndpoint().get());
}

@Test
public void throwExceptionWhenAuthTokenEmptyOrNull() {
assertThrows(ClientSdkException.class, () -> AuthTokenParser.parse(null));
assertThrows(ClientSdkException.class, () -> AuthTokenParser.parse(" "));
}

@Test
public void throwExceptionForInvalidClaimsToken() {
assertThrows(ClientSdkException.class, () -> AuthTokenParser.parse("abcd.effh.jdjjdjdj"));
}

@Test
public void throwExceptionForMalformedToken() {
assertThrows(ClientSdkException.class, () -> AuthTokenParser.parse("abcd"));
}
}
18 changes: 18 additions & 0 deletions momento-sdk/src/intTest/java/momento/sdk/MomentoTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import momento.sdk.exceptions.CacheAlreadyExistsException;
import momento.sdk.exceptions.ClientSdkException;
import momento.sdk.exceptions.InvalidArgumentException;
import momento.sdk.messages.CacheGetResponse;
import momento.sdk.messages.CacheSetResponse;
Expand All @@ -20,6 +21,11 @@ final class MomentoTest {
private String authToken;
private String cacheName;

// These secrets have botched up signature section, so should be okay to have them in source
// control.
private static final String TEST_AUTH_TOKEN_NO_ENDPOINT =
"eyJhbGciOiJIUzUxMiJ9.eyJzdWIiOiJpbnRlZ3JhdGlvbiJ9.ZOgkTs";

@BeforeAll
static void beforeAll() {
if (System.getenv("TEST_AUTH_TOKEN") == null) {
Expand Down Expand Up @@ -60,6 +66,18 @@ void testHappyPath() {
assertEquals("bar", rsp.asStringUtf8().get());
}

@Test
void authTokenWithNoEndpointAndNoEndpointOverride_throwsException() {
assertThrows(
ClientSdkException.class,
() -> Momento.builder().authToken(TEST_AUTH_TOKEN_NO_ENDPOINT).build());
}

@Test
void missingAuthToken_throwsException() {
assertThrows(ClientSdkException.class, () -> Momento.builder().build());
}

@Test
void testInvalidCacheName() {
Momento momento =
Expand Down
62 changes: 62 additions & 0 deletions momento-sdk/src/main/java/momento/sdk/AuthTokenParser.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
package momento.sdk;

import io.jsonwebtoken.Header;
import io.jsonwebtoken.Jwt;
import io.jsonwebtoken.Jwts;
import java.util.Optional;
import momento.sdk.exceptions.ClientSdkException;

final class AuthTokenParser {

private AuthTokenParser() {}

public static Claims parse(String authToken) {
try {
ensurePresent(authToken);
Jwt<Header, io.jsonwebtoken.Claims> parsed =
Jwts.parserBuilder().build().parseClaimsJwt(tokenWithOnlyHeaderAndClaims(authToken));
return new Claims(parsed.getBody());
} catch (Exception e) {
throw new ClientSdkException("Failed to parse Auth Token", e);
}
}

private static void ensurePresent(String authToken) {
if (authToken == null || authToken.isEmpty()) {
throw new IllegalArgumentException("Malformed Auth Token.");
}
}

// https://github.com/jwtk/jjwt/issues/280
private static String tokenWithOnlyHeaderAndClaims(String authToken) {
String[] splitToken = authToken.split("\\.");
if (splitToken == null || splitToken.length < 2) {
throw new IllegalArgumentException("Malformed Auth Token");
}
return splitToken[0] + "." + splitToken[1] + ".";
}

static class Claims {

private static final String CONTROL_ENDPOINT_CLAIM_NAME = "cp";
private static final String CACHE_ENDPOINT_CLAIM_NAME = "c";

private Optional<String> controlEndpoint;
private Optional<String> cacheEndpoint;

private Claims(io.jsonwebtoken.Claims claims) {
controlEndpoint =
Optional.ofNullable((String) claims.getOrDefault(CONTROL_ENDPOINT_CLAIM_NAME, null));
cacheEndpoint =
Optional.ofNullable((String) claims.getOrDefault(CACHE_ENDPOINT_CLAIM_NAME, null));
}

public Optional<String> controlEndpoint() {
return controlEndpoint;
}

public Optional<String> cacheEndpoint() {
return cacheEndpoint;
}
}
}
48 changes: 18 additions & 30 deletions momento-sdk/src/main/java/momento/sdk/Momento.java
Original file line number Diff line number Diff line change
Expand Up @@ -10,34 +10,37 @@
import java.io.Closeable;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import momento.sdk.exceptions.CacheAlreadyExistsException;
import momento.sdk.exceptions.CacheServiceExceptionMapper;
import momento.sdk.exceptions.ClientSdkException;
import momento.sdk.messages.CreateCacheResponse;

public final class Momento implements Closeable {

private static final String CONTROL_ENDPOINT_PREFIX = "control.";
private static final String CACHE_ENDPOINT_PREFIX = "cache.";

private final String authToken;
private final String hostedZone;
private final ScsControlBlockingStub blockingStub;
private final ManagedChannel channel;
private final MomentoEndpointsResolver.MomentoEndpoints momentoEndpoints;

private Momento(String authToken, Optional<String> hostedZoneOverride) {

private Momento(String authToken, String hostedZone) {
this.authToken = authToken;
this.hostedZone = hostedZone;
this.momentoEndpoints = MomentoEndpointsResolver.resolve(authToken, hostedZoneOverride);
this.channel = setupConnection(momentoEndpoints, authToken);
this.blockingStub = ScsControlGrpc.newBlockingStub(channel);
}

private static ManagedChannel setupConnection(
MomentoEndpointsResolver.MomentoEndpoints momentoEndpoints, String authToken) {
NettyChannelBuilder channelBuilder =
NettyChannelBuilder.forAddress(CONTROL_ENDPOINT_PREFIX + hostedZone, 443);
NettyChannelBuilder.forAddress(momentoEndpoints.controlEndpoint(), 443);
channelBuilder.useTransportSecurity();
channelBuilder.disableRetry();
List<ClientInterceptor> clientInterceptors = new ArrayList<>();
clientInterceptors.add(new AuthInterceptor(authToken));
channelBuilder.intercept(clientInterceptors);
ManagedChannel channel = channelBuilder.build();
this.blockingStub = ScsControlGrpc.newBlockingStub(channel);
this.channel = channel;
return channelBuilder.build();
}

/**
Expand Down Expand Up @@ -70,7 +73,7 @@ public CreateCacheResponse createCache(String cacheName) {

public Cache getCache(String cacheName) {
checkCacheNameValid(cacheName);
return makeCacheClient(authToken, cacheName, hostedZone);
return makeCacheClient(authToken, cacheName, momentoEndpoints.cacheEndpoint());
}

private CreateCacheRequest buildCreateCacheRequest(String cacheName) {
Expand All @@ -84,7 +87,7 @@ private static void checkCacheNameValid(String cacheName) {
}

private static Cache makeCacheClient(String authToken, String cacheName, String endpoint) {
return new Cache(authToken, cacheName, CACHE_ENDPOINT_PREFIX + endpoint);
return new Cache(authToken, cacheName, endpoint);
}

public void close() {
Expand All @@ -95,14 +98,9 @@ public static MomentoBuilder builder() {
return new MomentoBuilder();
}

// TODO: ParseJWT to determine the authToken
private static String extractEndpoint(String authToken) {
return null;
}

public static class MomentoBuilder {
private String authToken;
private String endpointOverride;
private Optional<String> endpointOverride = Optional.empty();

public MomentoBuilder authToken(String authToken) {
this.authToken = authToken;
Expand All @@ -119,25 +117,15 @@ public MomentoBuilder authToken(String authToken) {
// which the requests
// will be made.
public MomentoBuilder endpointOverride(String endpointOverride) {
this.endpointOverride = endpointOverride;
this.endpointOverride = Optional.ofNullable(endpointOverride);
return this;
}

public Momento build() {
if (authToken == null || authToken.isEmpty()) {
throw new ClientSdkException("Auth Token is required");
}
// Endpoint must be either available in the authToken or must be provided via
// endPointOverride.
String endpoint =
endpointOverride != null && !endpointOverride.isEmpty()
? endpointOverride
: extractEndpoint(authToken);

if (endpoint == null) {
throw new ClientSdkException("Endpoint for cache service is a required parameter.");
}
return new Momento(authToken, endpoint);
return new Momento(authToken, endpointOverride);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
package momento.sdk;

import java.util.Optional;
import momento.sdk.exceptions.ClientSdkException;

final class MomentoEndpointsResolver {

private static final String CONTROL_ENDPOINT_PREFIX = "control.";
private static final String CACHE_ENDPOINT_PREFIX = "cache.";

public static MomentoEndpoints resolve(String authToken, Optional<String> hostedZone) {
AuthTokenParser.Claims claims = AuthTokenParser.parse(authToken);
String controlEndpoint = getControlEndpoint(claims, hostedZone);
String cacheEndpoint = getCacheEndpoint(claims, hostedZone);
return new MomentoEndpoints(controlEndpoint, cacheEndpoint);
}

private static String getControlEndpoint(
AuthTokenParser.Claims claims, Optional<String> hostedZone) {
return controlEndpointFromHostedZone(hostedZone)
.orElseGet(
() ->
claims
.controlEndpoint()
.orElseThrow(
() ->
new ClientSdkException(
"Failed to determine control endpoint from the auth token or an override")));
}

private static String getCacheEndpoint(
AuthTokenParser.Claims claims, Optional<String> hostedZone) {
return cacheEndpointFromHostedZone(hostedZone)
.orElseGet(
() ->
claims
.cacheEndpoint()
.orElseThrow(
() ->
new ClientSdkException(
"Failed to determine cache endpoint from the auth token or an override")));
}

private static Optional<String> controlEndpointFromHostedZone(Optional<String> hostedZone) {
if (hostedZone.isPresent()) {
return Optional.of(CONTROL_ENDPOINT_PREFIX + hostedZone.get());
}
return Optional.empty();
}

private static Optional<String> cacheEndpointFromHostedZone(Optional<String> hostedZone) {
if (hostedZone.isPresent()) {
return Optional.of(CACHE_ENDPOINT_PREFIX + hostedZone.get());
}
return Optional.empty();
}

static class MomentoEndpoints {
private final String controlEndpoint;
private final String cacheEndpoint;

private MomentoEndpoints(String controlEndpoint, String cacheEndpoint) {
this.cacheEndpoint = cacheEndpoint;
this.controlEndpoint = controlEndpoint;
}

public String controlEndpoint() {
return this.controlEndpoint;
}

public String cacheEndpoint() {
return this.cacheEndpoint;
}
}
}

0 comments on commit 0a3b890

Please sign in to comment.