Skip to content

Commit

Permalink
migrate zms server use aws sdk v2
Browse files Browse the repository at this point in the history 
Signed-off-by: Henry Avetisyan <hga@yahooinc.com>
  • Loading branch information
@havetisyan
havetisyan committed Aug 8, 2024
1 parent b3f8680 commit 3b2d0f2
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 74 deletions.
15 changes: 4 additions & 11 deletions servers/zms/pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,6 @@

<dependencyManagement>
<dependencies>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-bom</artifactId>
<version>${aws.version}</version>
<type>pom</type>
<scope>import</scope>
</dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>bom</artifactId>
Expand Down Expand Up @@ -82,12 +75,12 @@
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-sts</artifactId>
<groupId>software.amazon.awssdk</groupId>
<artifactId>sts</artifactId>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-rds</artifactId>
<groupId>software.amazon.awssdk</groupId>
<artifactId>rds</artifactId>
</dependency>
<dependency>
<groupId>org.glassfish.jersey.inject</groupId>
Expand Down
70 changes: 34 additions & 36 deletions servers/zms/src/main/java/com/yahoo/athenz/zms/store/impl/AWSObjectStoreFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,16 @@
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import com.yahoo.athenz.common.server.util.Utils;
import org.eclipse.jetty.util.StringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.amazonaws.auth.InstanceProfileCredentialsProvider;
import com.amazonaws.services.rds.auth.GetIamAuthTokenRequest;
import com.amazonaws.services.rds.auth.RdsIamAuthTokenGenerator;
import com.amazonaws.util.EC2MetadataUtils;
import software.amazon.awssdk.auth.credentials.ProfileCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.rds.RdsClient;
import software.amazon.awssdk.services.rds.RdsUtilities;
import software.amazon.awssdk.services.rds.model.GenerateAuthenticationTokenRequest;

import com.yahoo.athenz.auth.PrivateKeyStore;
import com.yahoo.athenz.common.server.db.DataSourceFactory;
Expand Down Expand Up @@ -100,6 +103,10 @@ public ObjectStore create(PrivateKeyStore keyStore) {
return new JDBCObjectStore(dataPrimarySource, dataReplicaSource);
}

public void stop() {
scheduledThreadPool.shutdownNow();
}

void setConnectionProperties(Properties mysqlProperties, final String token) {
mysqlProperties.setProperty(ZMSConsts.DB_PROP_VERIFY_SERVER_CERT,
System.getProperty(ZMSConsts.ZMS_PROP_JDBC_VERIFY_SERVER_CERT, "true"));
Expand All @@ -111,41 +118,34 @@ void setConnectionProperties(Properties mysqlProperties, final String token) {
mysqlProperties.setProperty(ZMSConsts.DB_PROP_PASSWORD, token);
}

InstanceProfileCredentialsProvider getNewInstanceCredentialsProvider() {
return new InstanceProfileCredentialsProvider(true);
Region getRegion() {
return Utils.getAwsRegion(Region.US_EAST_1);
}

String getRegion() {
return EC2MetadataUtils.getEC2InstanceRegion();
}
String getAuthToken(String hostname, int port, String rdsUser) {

String getGeneratorAuthToken(RdsIamAuthTokenGenerator generator, final String hostname,
int port, final String rdsUser) {
return generator.getAuthToken(GetIamAuthTokenRequest.builder()
.hostname(hostname).port(port).userName(rdsUser)
.build());
}
String authToken = null;
try (RdsClient rdsClient = RdsClient.builder().region(getRegion())
.credentialsProvider(ProfileCredentialsProvider.create()).build()) {

String getAuthToken(String hostname, int port, String rdsUser) {
RdsUtilities utilities = rdsClient.utilities();

InstanceProfileCredentialsProvider awsCredProvider = getNewInstanceCredentialsProvider();

if (LOG.isDebugEnabled()) {
LOG.debug("getAuthToken: Access key id: {}", awsCredProvider.getCredentials().getAWSAccessKeyId());
}
GenerateAuthenticationTokenRequest tokenRequest = GenerateAuthenticationTokenRequest.builder()
.credentialsProvider(ProfileCredentialsProvider.create())
.username(rdsUser)
.port(port)
.hostname(hostname)
.build();

RdsIamAuthTokenGenerator generator = RdsIamAuthTokenGenerator.builder()
.credentials(awsCredProvider)
.region(getRegion())
.build();

if (LOG.isDebugEnabled()) {
LOG.debug("Instance {} Port {} User {} Region: {}", hostname, port, rdsUser, getRegion());
authToken = utilities.generateAuthenticationToken(tokenRequest);

} catch (Exception ex) {
LOG.error("getAuthToken: unable to generate auth token", ex);
}
return getGeneratorAuthToken(generator, hostname, port, rdsUser);

return authToken;
}

void updateCredentials(String hostname, Properties mysqlProperties) {

// if we have no hostname specified then we have nothing to do
Expand All @@ -155,12 +155,10 @@ void updateCredentials(String hostname, Properties mysqlProperties) {
}

// obtain iam role credentials and update the properties object
try {
final String rdsToken = getAuthToken(hostname, rdsPort, rdsUser);

final String rdsToken = getAuthToken(hostname, rdsPort, rdsUser);
if (!StringUtil.isEmpty(rdsToken)) {
mysqlProperties.setProperty(ZMSConsts.DB_PROP_PASSWORD, rdsToken);
} catch (Exception ex) {
LOG.error("CredentialsUpdater: unable to update auth token", ex);
}
}

Expand Down
47 changes: 20 additions & 27 deletions servers/zms/src/test/java/com/yahoo/athenz/zms/store/impl/AWSObjectStoreFactoryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,8 @@
*/
package com.yahoo.athenz.zms.store.impl;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.InstanceProfileCredentialsProvider;
import com.amazonaws.services.rds.auth.RdsIamAuthTokenGenerator;
import com.yahoo.athenz.zms.ZMSConsts;
import com.yahoo.athenz.zms.store.ObjectStore;
import org.mockito.Mockito;
import org.testng.annotations.Test;

import static org.testng.Assert.assertNotNull;
Expand All @@ -30,22 +26,7 @@ public class AWSObjectStoreFactoryTest {
static class TestAWSObjectStoreFactory extends AWSObjectStoreFactory {

@Override
InstanceProfileCredentialsProvider getNewInstanceCredentialsProvider() {
InstanceProfileCredentialsProvider provider = Mockito.mock(InstanceProfileCredentialsProvider.class);
AWSCredentials awsCredentials = Mockito.mock(AWSCredentials.class);
Mockito.when(provider.getCredentials()).thenReturn(awsCredentials);
Mockito.when(awsCredentials.getAWSAccessKeyId()).thenReturn("id");
return provider;
}

@Override
String getRegion() {
return "us-west-2";
}

@Override
String getGeneratorAuthToken(RdsIamAuthTokenGenerator generator, final String hostname,
int port, final String rdsUser) {
String getAuthToken(final String hostname, int port, final String rdsUser) {
if (rdsUser.equals("rds-user")) {
return "token";
}
Expand All @@ -71,6 +52,11 @@ public void testCreate() {
} catch (InterruptedException ignored) {
}
assertNotNull(store);
factory.stop();

System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_PRIMARY_INSTANCE);
System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_USER);
System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_CREDS_REFRESH_TIME);
}

@Test
Expand All @@ -84,20 +70,21 @@ public void testOriginalMethods() {

AWSObjectStoreFactory factory = new AWSObjectStoreFactory();

try {
factory.getNewInstanceCredentialsProvider();
} catch (Exception ignored) {
}

try {
factory.getRegion();
} catch (Exception ignored) {
}

try {
factory.getGeneratorAuthToken(null, "localhost", 40888, "rdsUser");
factory.getAuthToken("host", 3306, "user");
} catch (Exception ignored) {
}

factory.stop();

System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_PRIMARY_INSTANCE);
System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_USER);
System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_CREDS_REFRESH_TIME);
}

@Test
Expand All @@ -106,7 +93,7 @@ public void testCreateWithReplica() {
System.setProperty(ZMSConsts.ZMS_PROP_AWS_RDS_PRIMARY_INSTANCE, "instance");
System.setProperty(ZMSConsts.ZMS_PROP_AWS_RDS_REPLICA_INSTANCE, "replica");
System.setProperty(ZMSConsts.ZMS_PROP_AWS_RDS_USER, "rds-user");
System.setProperty(ZMSConsts.ZMS_PROP_AWS_RDS_CREDS_REFRESH_TIME, "30000");
System.setProperty(ZMSConsts.ZMS_PROP_AWS_RDS_CREDS_REFRESH_TIME, "1");

AWSObjectStoreFactory factory = new TestAWSObjectStoreFactory();
ObjectStore store = factory.create(null);
Expand All @@ -117,5 +104,11 @@ public void testCreateWithReplica() {
} catch (InterruptedException ignored) {
}
assertNotNull(store);
factory.stop();

System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_PRIMARY_INSTANCE);
System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_REPLICA_INSTANCE);
System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_USER);
System.clearProperty(ZMSConsts.ZMS_PROP_AWS_RDS_CREDS_REFRESH_TIME);
}
}

0 comments on commit 3b2d0f2

Please sign in to comment.