diff --git a/libs/java/dynamodb_client_factory/pom.xml b/libs/java/dynamodb_client_factory/pom.xml index fbbed834cc1..95d21f57368 100644 --- a/libs/java/dynamodb_client_factory/pom.xml +++ b/libs/java/dynamodb_client_factory/pom.xml @@ -31,27 +31,9 @@ athenz-dynamodb-client-factory Generate dynamodb client based on AWS temporary credentials jar - - - - com.amazonaws - aws-java-sdk-bom - ${aws.version} - pom - import - - - software.amazon.awssdk - bom - ${aws2.version} - pom - import - - - - 0.9230 + 0.9417 @@ -60,13 +42,10 @@ athenz-zts-java-client ${project.parent.version} - - com.amazonaws - aws-java-sdk-dynamodb - software.amazon.awssdk dynamodb + ${aws2.version} com.google.guava diff --git a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentials.java b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentials.java index 54068de8e80..79deec74eba 100644 --- a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentials.java +++ b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentials.java @@ -18,30 +18,43 @@ package com.yahoo.athenz.db.dynamodb; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.yahoo.athenz.zts.AWSCredentialsProviderImpl; +import com.yahoo.athenz.zts.AWSCredentialsProviderImplV2; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; public class DynamoDBClientAndCredentials { - private final AmazonDynamoDB amazonDynamoDB; - private final DynamoDbAsyncClient amazonDynamoAsyncDB; - private final AWSCredentialsProviderImpl awsCredentialsProvider; - public DynamoDBClientAndCredentials(AmazonDynamoDB amazonDynamoDB, DynamoDbAsyncClient amazonDynamoAsyncDB, AWSCredentialsProviderImpl awsCredentialsProvider) { - this.amazonDynamoDB = amazonDynamoDB; - this.amazonDynamoAsyncDB = amazonDynamoAsyncDB; + private final DynamoDbClient dynamoDbClient; + private final DynamoDbAsyncClient dynamoDbAsyncClient; + private final AWSCredentialsProviderImplV2 awsCredentialsProvider; + + public DynamoDBClientAndCredentials(DynamoDbClient dynamoDbClient, DynamoDbAsyncClient dynamoDbAsyncClient, + AWSCredentialsProviderImplV2 awsCredentialsProvider) { + this.dynamoDbClient = dynamoDbClient; + this.dynamoDbAsyncClient = dynamoDbAsyncClient; this.awsCredentialsProvider = awsCredentialsProvider; } - public AmazonDynamoDB getAmazonDynamoDB() { - return amazonDynamoDB; + public DynamoDbClient getDynamoDbClient() { + return dynamoDbClient; } - public DynamoDbAsyncClient getAmazonDynamoAsyncDB() { - return amazonDynamoAsyncDB; + public DynamoDbAsyncClient getDynamoDbAsyncClient() { + return dynamoDbAsyncClient; } - public AWSCredentialsProviderImpl getAwsCredentialsProvider() { - return awsCredentialsProvider; + public void close() { + if (dynamoDbClient != null) { + dynamoDbClient.close(); + } + if (dynamoDbAsyncClient != null) { + dynamoDbAsyncClient.close(); + } + if (awsCredentialsProvider != null) { + try { + awsCredentialsProvider.close(); + } catch (Exception ignored) { + } + } } } diff --git a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcher.java b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcher.java index 81d9810569d..ff53590caee 100644 --- a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcher.java +++ b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcher.java @@ -23,12 +23,14 @@ public interface DynamoDBClientFetcher { /** - * Returns a DynamoDBClient and the AWS credential provider used for authentication. - * The credentialProvider should be closed after DynamoDBClient is no longer needed. + * Returns a DynamoDBClient wrapper object that includes both regular + * and async clients along with the AWS credentials provider. + * The clients should be closed after DynamoDBClient is no + * longer needed which would close the associated AWS credential provider. * (GC might not run for a long period of time) * @param ztsClientNotificationSender notification sender object * @param dynamoDBClientSettings contains private key store and client settings - * @return DynamoDBClientAndCredentials which contains both a DynamoDB client and the credentialProvider used + * @return DynamoDBClientAndCredentials which contains both DynamoDB clients and the credentialProvider used */ DynamoDBClientAndCredentials getDynamoDBClient(ZTSClientNotificationSender ztsClientNotificationSender, DynamoDBClientSettings dynamoDBClientSettings); diff --git a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImpl.java b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImpl.java index 41009c7e36b..44593bdb851 100644 --- a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImpl.java +++ b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImpl.java @@ -18,33 +18,27 @@ package com.yahoo.athenz.db.dynamodb; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder; -import com.amazonaws.util.EC2MetadataUtils; import com.oath.auth.KeyRefresher; import com.oath.auth.Utils; -import com.yahoo.athenz.zts.AWSCredentialsProviderImpl; import com.yahoo.athenz.zts.AWSCredentialsProviderImplV2; import com.yahoo.athenz.zts.ZTSClientNotificationSender; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.utils.StringUtils; import javax.net.ssl.SSLContext; public class DynamoDBClientFetcherImpl implements DynamoDBClientFetcher { private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBClientFetcherImpl.class); - private String defaultAwsRegion; public DynamoDBClientFetcherImpl() { } - public DynamoDBClientFetcherImpl(String defaultAwsRegion) { - this.defaultAwsRegion = defaultAwsRegion; - } - @Override public DynamoDBClientAndCredentials getDynamoDBClient(ZTSClientNotificationSender ztsClientNotificationSender, DynamoDBClientSettings dynamoDBClientSettings) { @@ -58,28 +52,29 @@ public DynamoDBClientAndCredentials getDynamoDBClient(ZTSClientNotificationSende return getAuthenticatedDynamoDBClient(dynamoDBClientSettings, ztsClientNotificationSender); } else { LOGGER.info("DynamoDB client will use existing AWS authentication"); - String region = getAWSRegion(dynamoDBClientSettings.getRegion()); - AmazonDynamoDB client = AmazonDynamoDBClientBuilder - .standard() - .withRegion(region) - .build(); - DynamoDbAsyncClient asyncClient = DynamoDbAsyncClient.builder() - .region(Region.of(region)) - .build(); + DynamoDbClient dynamoDbClient = null; + DynamoDbAsyncClient dynamoDbAsyncClient = null; + if (dynamoDBClientSettings.isAsyncClient()) { + dynamoDbAsyncClient = DynamoDbAsyncClient.builder() + .region(getAWSRegion(dynamoDBClientSettings.getRegion())) + .build(); + } else { + dynamoDbClient = DynamoDbClient.builder() + .region(getAWSRegion(dynamoDBClientSettings.getRegion())) + .build(); + } - return new DynamoDBClientAndCredentials(client, asyncClient, null); + return new DynamoDBClientAndCredentials(dynamoDbClient, dynamoDbAsyncClient, null); } } - String getAWSRegion(final String settingRegion) { - if (StringUtils.isEmpty(settingRegion)) { - if (defaultAwsRegion == null) { - defaultAwsRegion = EC2MetadataUtils.getEC2InstanceRegion(); - } - return defaultAwsRegion; + Region getAWSRegion(final String settingRegion) { + if (!StringUtils.isEmpty(settingRegion)) { + return Region.of(settingRegion); } else { - return settingRegion; + DefaultAwsRegionProviderChain regionProvider = DefaultAwsRegionProviderChain.builder().build(); + return regionProvider.getRegion(); } } @@ -102,10 +97,10 @@ private DynamoDBClientAndCredentials getAuthenticatedDynamoDBClient(DynamoDBClie LOGGER.error("Failed to get AWS Temporary credentials", ex); } - AWSCredentialsProviderImpl credentialsProvider = null; - String region = dynamoDBClientSettings.getRegion(); + AWSCredentialsProviderImplV2 credentialsProvider = null; + try { - credentialsProvider = new AWSCredentialsProviderImpl( + credentialsProvider = new AWSCredentialsProviderImplV2( dynamoDBClientSettings.getZtsURL(), sslContext, dynamoDBClientSettings.getDomainName(), @@ -114,35 +109,24 @@ private DynamoDBClientAndCredentials getAuthenticatedDynamoDBClient(DynamoDBClie dynamoDBClientSettings.getMinExpiryTime(), dynamoDBClientSettings.getMaxExpiryTime(), ztsClientNotificationSender); - } catch (Exception ex) { - LOGGER.error("Failed to generate AmazonDynamoDB client", ex); + LOGGER.error("Failed to generate DynamoDbClient client", ex); } - AmazonDynamoDB client = AmazonDynamoDBClientBuilder.standard() - .withCredentials(credentialsProvider) - .withRegion(region) - .build(); - - AWSCredentialsProviderImplV2 credentialsProviderV2 = null; - try { - credentialsProviderV2 = new AWSCredentialsProviderImplV2( - dynamoDBClientSettings.getZtsURL(), - sslContext, - dynamoDBClientSettings.getDomainName(), - dynamoDBClientSettings.getRoleName(), - dynamoDBClientSettings.getExternalId(), - dynamoDBClientSettings.getMinExpiryTime(), - dynamoDBClientSettings.getMaxExpiryTime(), - ztsClientNotificationSender); - } catch (Exception ex) { - LOGGER.error("Failed to generate DynamoDbAsyncClient client", ex); + DynamoDbClient dynamoDbClient = null; + DynamoDbAsyncClient dynamoDbAsyncClient = null; + if (dynamoDBClientSettings.isAsyncClient()) { + dynamoDbAsyncClient = DynamoDbAsyncClient.builder() + .credentialsProvider(credentialsProvider) + .region(getAWSRegion(dynamoDBClientSettings.getRegion())) + .build(); + } else { + dynamoDbClient = DynamoDbClient.builder() + .credentialsProvider(credentialsProvider) + .region(getAWSRegion(dynamoDBClientSettings.getRegion())) + .build(); } - DynamoDbAsyncClient asyncClient = DynamoDbAsyncClient.builder() - .credentialsProvider(credentialsProviderV2) - .region(Region.of(region)) - .build(); - return new DynamoDBClientAndCredentials(client, asyncClient, credentialsProvider); + return new DynamoDBClientAndCredentials(dynamoDbClient, dynamoDbAsyncClient, credentialsProvider); } } diff --git a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettings.java b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettings.java index 3f0268ddb0b..1cf637fa982 100644 --- a/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettings.java +++ b/libs/java/dynamodb_client_factory/src/main/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettings.java @@ -34,22 +34,14 @@ public class DynamoDBClientSettings { private final String externalId; private final Integer minExpiryTime; private final Integer maxExpiryTime; - private final String keygroupName; - - public DynamoDBClientSettings(String certPath, - String domainName, - String roleName, - String trustStore, - String trustStorePassword, - String ztsURL, - String region, - String keyPath, - String appName, - PrivateKeyStore keyStore, - String externalId, - Integer minExpiryTime, - Integer maxExpiryTime, - String keygroupName) { + private final String keyGroupName; + private final boolean isAsyncClient; + + public DynamoDBClientSettings(String certPath, String domainName, String roleName, String trustStore, + String trustStorePassword, String ztsURL, String region, String keyPath, String appName, + PrivateKeyStore keyStore, String externalId, Integer minExpiryTime, Integer maxExpiryTime, + String keyGroupName, boolean isAsyncClient) { + this.certPath = certPath; this.domainName = domainName; this.roleName = roleName; @@ -63,7 +55,8 @@ public DynamoDBClientSettings(String certPath, this.externalId = externalId; this.minExpiryTime = minExpiryTime; this.maxExpiryTime = maxExpiryTime; - this.keygroupName = keygroupName; + this.keyGroupName = keyGroupName; + this.isAsyncClient = isAsyncClient; } public boolean areCredentialsProvided() { @@ -123,6 +116,10 @@ public char[] getTrustStorePasswordChars() { return null; } - return keyStore.getSecret(appName, keygroupName, trustStorePassword); + return keyStore.getSecret(appName, keyGroupName, trustStorePassword); + } + + public boolean isAsyncClient() { + return isAsyncClient; } } diff --git a/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentialsTest.java b/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentialsTest.java index f121af18388..a592f2b4793 100644 --- a/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentialsTest.java +++ b/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientAndCredentialsTest.java @@ -18,28 +18,61 @@ package com.yahoo.athenz.db.dynamodb; +import com.yahoo.athenz.zts.AWSCredentialsProviderImplV2; import org.mockito.Mockito; import org.testng.annotations.Test; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.yahoo.athenz.zts.AWSCredentialsProviderImpl; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; + +import java.io.IOException; import static org.testng.Assert.assertEquals; public class DynamoDBClientAndCredentialsTest { @Test - public void testDynamoDBClientAndCredentials() { + public void testDynamoDBClientAndCredentials() throws IOException { + + DynamoDbClient dynamoDbClient = Mockito.mock(DynamoDbClient.class); + Mockito.doNothing().when(dynamoDbClient).close(); + DynamoDbAsyncClient dynamoDbAsyncClient = Mockito.mock(DynamoDbAsyncClient.class); + Mockito.doNothing().when(dynamoDbAsyncClient).close(); + AWSCredentialsProviderImplV2 credentialsProvider = Mockito.mock(AWSCredentialsProviderImplV2.class); + Mockito.doNothing().when(credentialsProvider).close(); + DynamoDBClientAndCredentials dynamoDBClientAndCredentials = + new DynamoDBClientAndCredentials(dynamoDbClient, dynamoDbAsyncClient, credentialsProvider); + assertEquals(dynamoDbClient, dynamoDBClientAndCredentials.getDynamoDbClient()); + assertEquals(dynamoDbAsyncClient, dynamoDBClientAndCredentials.getDynamoDbAsyncClient()); + dynamoDBClientAndCredentials.close(); + } - AmazonDynamoDB amazonDynamoDB = Mockito.mock(AmazonDynamoDB.class); - DynamoDbAsyncClient amazonDynamoAsyncDB = Mockito.mock(DynamoDbAsyncClient.class); - AWSCredentialsProviderImpl awsCredentialsProvider = Mockito.mock(AWSCredentialsProviderImpl.class); + @Test + public void testDynamoDBClientAndCredentialsNullProvider() { - DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials( - amazonDynamoDB, amazonDynamoAsyncDB, awsCredentialsProvider); + DynamoDbClient dynamoDbClient = Mockito.mock(DynamoDbClient.class); + Mockito.doNothing().when(dynamoDbClient).close(); + DynamoDbAsyncClient dynamoDbAsyncClient = Mockito.mock(DynamoDbAsyncClient.class); + Mockito.doNothing().when(dynamoDbAsyncClient).close(); + DynamoDBClientAndCredentials dynamoDBClientAndCredentials = + new DynamoDBClientAndCredentials(dynamoDbClient, dynamoDbAsyncClient, null); + assertEquals(dynamoDbClient, dynamoDBClientAndCredentials.getDynamoDbClient()); + assertEquals(dynamoDbAsyncClient, dynamoDBClientAndCredentials.getDynamoDbAsyncClient()); + dynamoDBClientAndCredentials.close(); + } + + @Test + public void testDynamoDBClientAndCredentialsException() throws IOException { - assertEquals(amazonDynamoDB, dynamoDBClientAndCredentials.getAmazonDynamoDB()); - assertEquals(amazonDynamoAsyncDB, dynamoDBClientAndCredentials.getAmazonDynamoAsyncDB()); - assertEquals(awsCredentialsProvider, dynamoDBClientAndCredentials.getAwsCredentialsProvider()); + DynamoDbClient dynamoDbClient = Mockito.mock(DynamoDbClient.class); + Mockito.doNothing().when(dynamoDbClient).close(); + DynamoDbAsyncClient dynamoDbAsyncClient = Mockito.mock(DynamoDbAsyncClient.class); + Mockito.doNothing().when(dynamoDbAsyncClient).close(); + AWSCredentialsProviderImplV2 credentialsProvider = Mockito.mock(AWSCredentialsProviderImplV2.class); + Mockito.doThrow(new IllegalArgumentException()).when(credentialsProvider).close(); + DynamoDBClientAndCredentials dynamoDBClientAndCredentials = + new DynamoDBClientAndCredentials(dynamoDbClient, dynamoDbAsyncClient, credentialsProvider); + assertEquals(dynamoDbClient, dynamoDBClientAndCredentials.getDynamoDbClient()); + assertEquals(dynamoDbAsyncClient, dynamoDBClientAndCredentials.getDynamoDbAsyncClient()); + dynamoDBClientAndCredentials.close(); } } diff --git a/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImplTest.java b/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImplTest.java index 0b9bc8bf921..c75db45ac38 100644 --- a/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImplTest.java +++ b/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientFetcherImplTest.java @@ -18,110 +18,124 @@ package com.yahoo.athenz.db.dynamodb; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.google.common.io.Resources; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.zts.ZTSClientNotificationSender; import org.mockito.Mockito; import org.testng.annotations.Test; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; import static org.testng.AssertJUnit.assertNotNull; public class DynamoDBClientFetcherImplTest { - private static final String ZTS_PROP_DYNAMODB_KEY_PATH = "athenz.zts.dynamodb_key_path"; - private static final String ZTS_PROP_DYNAMODB_CERT_PATH = "athenz.zts.dynamodb_cert_path"; - private static final String ZTS_PROP_DYNAMODB_DOMAIN = "athenz.zts.dynamodb_aws_domain"; - private static final String ZTS_PROP_DYNAMODB_ROLE = "athenz.zts.dynamodb_aws_role"; - private static final String ZTS_PROP_DYNAMODB_TRUSTSTORE = "athenz.zts.dynamodb_trust_store_path"; - private static final String ZTS_PROP_DYNAMODB_TRUSTSTORE_PASSWORD = "athenz.zts.dynamodb_trust_store_password"; - private static final String ZTS_PROP_DYNAMODB_REGION = "athenz.zts.dynamodb_region"; - private static final String ZTS_PROP_DYNAMODB_ZTS_URL = "athenz.zts.dynamodb_zts_url"; - private static final String ZTS_PROP_DYNAMODB_MIN_EXPIRY_TIME = "athenz.zts.dynamodb_min_expiry_time"; - private static final String ZTS_PROP_DYNAMODB_MAX_EXPIRY_TIME = "athenz.zts.dynamodb_max_expiry_time"; - @Test - public void testGetClientWitSpecifiedRegion() { - System.setProperty(ZTS_PROP_DYNAMODB_REGION, "test.region"); - DynamoDBClientFetcher dynamoDBClientFetcher = DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); - PrivateKeyStore keyStore = Mockito.mock(PrivateKeyStore.class); - ZTSClientNotificationSender ztsClientNotificationSender = Mockito.mock(ZTSClientNotificationSender.class); - DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(null, null, null, null, null, - null, "test.region", null, null, keyStore, null, null, null, null); - AmazonDynamoDB dynamoDBClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, - dynamoDBClientSettings).getAmazonDynamoDB(); - assertNotNull(dynamoDBClient); - System.clearProperty(ZTS_PROP_DYNAMODB_REGION); + public void testDynamoDBClientFetcherFactory() { + DynamoDBClientFetcherFactory dynamoDBClientFetcherFactory = new DynamoDBClientFetcherFactory(); + assertNotNull(dynamoDBClientFetcherFactory); + assertNotNull(DynamoDBClientFetcherFactory.getDynamoDBClientFetcher()); } @Test - public void testGetClientWithDefaultRegion() { - DynamoDBClientFetcher dynamoDBClientFetcher = new DynamoDBClientFetcherImpl("testRegion"); + public void testGetClientWitSpecifiedRegion() { + DynamoDBClientFetcher dynamoDBClientFetcher = DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); PrivateKeyStore keyStore = Mockito.mock(PrivateKeyStore.class); ZTSClientNotificationSender ztsClientNotificationSender = Mockito.mock(ZTSClientNotificationSender.class); DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(null, null, null, null, null, - null, "testRegion", null, null, keyStore, null, null, null, null); - AmazonDynamoDB dynamoDBClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, - dynamoDBClientSettings).getAmazonDynamoDB(); - assertNotNull(dynamoDBClient); + null, "test.region", null, null, keyStore, null, null, null, null, false); + DynamoDbClient dynamoDbClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbClient(); + assertNotNull(dynamoDbClient); } @Test - public void testGetAuthenticatedClient() { + public void testGetAuthenticatedClientWithCreds() { String certPath = Resources.getResource("gdpr.aws.core.cert.pem").getPath();//public String keyPath = Resources.getResource("gdpr.aws.core.key.pem").getPath();//private - System.setProperty(ZTS_PROP_DYNAMODB_KEY_PATH, keyPath); - System.setProperty(ZTS_PROP_DYNAMODB_CERT_PATH, certPath); - System.setProperty(ZTS_PROP_DYNAMODB_DOMAIN, "test.domain"); - System.setProperty(ZTS_PROP_DYNAMODB_REGION, "test.region"); - System.setProperty(ZTS_PROP_DYNAMODB_ROLE, "test.role"); - System.setProperty(ZTS_PROP_DYNAMODB_TRUSTSTORE, "test.truststore"); - System.setProperty(ZTS_PROP_DYNAMODB_TRUSTSTORE_PASSWORD, "test.truststore.password"); - System.setProperty(ZTS_PROP_DYNAMODB_ZTS_URL, "https://dev.zts.athenzcompany.com:4443/zts/v1"); - DynamoDBClientFetcherImpl dynamoDBClientFetcher = new DynamoDBClientFetcherImpl(); PrivateKeyStore keyStore = Mockito.mock(PrivateKeyStore.class); when(keyStore.getSecret(Mockito.eq(""), Mockito.eq(null), Mockito.eq("test.truststore.password"))) .thenReturn("mockPassword".toCharArray()); ZTSClientNotificationSender ztsClientNotificationSender = Mockito.mock(ZTSClientNotificationSender.class); + // first we test the regular client + DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(certPath, "test.domain", "test.role", "test.truststore", "test.truststore.password", "https://dev.zts.athenzcompany.com:4443/zts/v1", "test.region", keyPath, null, - keyStore, null, null, null, null); - AmazonDynamoDB dynamoDBClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, - dynamoDBClientSettings).getAmazonDynamoDB(); - assertNotNull(dynamoDBClient); - - System.clearProperty(ZTS_PROP_DYNAMODB_KEY_PATH); - System.clearProperty(ZTS_PROP_DYNAMODB_CERT_PATH); - System.clearProperty(ZTS_PROP_DYNAMODB_DOMAIN); - System.clearProperty(ZTS_PROP_DYNAMODB_REGION); - System.clearProperty(ZTS_PROP_DYNAMODB_ROLE); - System.clearProperty(ZTS_PROP_DYNAMODB_TRUSTSTORE); - System.clearProperty(ZTS_PROP_DYNAMODB_TRUSTSTORE_PASSWORD); - System.clearProperty(ZTS_PROP_DYNAMODB_ZTS_URL); + keyStore, null, null, null, null, false); + DynamoDbClient dynamoDbClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbClient(); + assertNotNull(dynamoDbClient); + DynamoDbAsyncClient dynamoDbAsyncClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbAsyncClient(); + assertNull(dynamoDbAsyncClient); + + // now the async client + + dynamoDBClientSettings = new DynamoDBClientSettings(certPath, "test.domain", + "test.role", "test.truststore", "test.truststore.password", + "https://dev.zts.athenzcompany.com:4443/zts/v1", "test.region", keyPath, null, + keyStore, null, null, null, null, true); + dynamoDbClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbClient(); + assertNull(dynamoDbClient); + dynamoDbAsyncClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbAsyncClient(); + assertNotNull(dynamoDbAsyncClient); } @Test - public void testGetAWSRegion() { + public void testGetAuthenticatedClientWithoutCreds() { DynamoDBClientFetcherImpl dynamoDBClientFetcher = new DynamoDBClientFetcherImpl(); - assertEquals(dynamoDBClientFetcher.getAWSRegion("us-west-2"), "us-west-2"); + ZTSClientNotificationSender ztsClientNotificationSender = Mockito.mock(ZTSClientNotificationSender.class); - dynamoDBClientFetcher = new DynamoDBClientFetcherImpl("us-east-1"); - assertEquals(dynamoDBClientFetcher.getAWSRegion("us-west-2"), "us-west-2"); - assertEquals(dynamoDBClientFetcher.getAWSRegion(""), "us-east-1"); - assertEquals(dynamoDBClientFetcher.getAWSRegion(null), "us-east-1"); + // first we test the regular client + + DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(null, "test.domain", + "test.role", null, null, "https://dev.zts.athenzcompany.com:4443/zts/v1", "test.region", + null, null, null, null, null, null, null, false); + DynamoDbClient dynamoDbClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbClient(); + assertNotNull(dynamoDbClient); + DynamoDbAsyncClient dynamoDbAsyncClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbAsyncClient(); + assertNull(dynamoDbAsyncClient); + + // now the async client + + dynamoDBClientSettings = new DynamoDBClientSettings(null, "test.domain", + "test.role", null, null, "https://dev.zts.athenzcompany.com:4443/zts/v1", "test.region", + null, null, null, null, null, null, null, true); + dynamoDbClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbClient(); + assertNull(dynamoDbClient); + dynamoDbAsyncClient = dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + dynamoDBClientSettings).getDynamoDbAsyncClient(); + assertNotNull(dynamoDbAsyncClient); + } + + @Test + public void testGetAWSRegion() { + + DynamoDBClientFetcherImpl dynamoDBClientFetcher = new DynamoDBClientFetcherImpl(); + assertEquals(dynamoDBClientFetcher.getAWSRegion("us-west-2"), Region.of("us-west-2")); // if this test is running in aws, then we'll get a valid region // value. if running on-prem, then we'll get an exception when ec2 meta // api is called. so we'll get a null value back dynamoDBClientFetcher = new DynamoDBClientFetcherImpl(); - dynamoDBClientFetcher.getAWSRegion(null); + try { + dynamoDBClientFetcher.getAWSRegion(null); + } catch (Exception ignored) { + } } } diff --git a/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettingsTest.java b/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettingsTest.java index 2891f392557..5112d92f5e4 100644 --- a/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettingsTest.java +++ b/libs/java/dynamodb_client_factory/src/test/java/com/yahoo/athenz/db/dynamodb/DynamoDBClientSettingsTest.java @@ -30,21 +30,8 @@ public class DynamoDBClientSettingsTest { @Test public void credentialsNotProvided() { PrivateKeyStore keyStore = Mockito.mock(PrivateKeyStore.class); - DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(null, - null, - null, - null, - null, - null, - null, - null, - null, - keyStore, - null, - null, - null, - null); - + DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(null, null, null, null, + null, null, null, null, null, keyStore, null, null, null, null, false); assertFalse(dynamoDBClientSettings.areCredentialsProvided()); } @@ -64,7 +51,9 @@ public void testCredentialsProvided() { when(keyStore.getSecret(Mockito.eq("test.appname"), Mockito.eq(null), Mockito.eq("test.truststore.password"))) .thenReturn("decryptedPassword".toCharArray()); - DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(certPath, domain, role, trustStore, trustStorePassword, ztsUrl, region, keyPath, appName, keyStore, null, null, null, null); + DynamoDBClientSettings dynamoDBClientSettings = new DynamoDBClientSettings(certPath, domain, role, + trustStore, trustStorePassword, ztsUrl, region, keyPath, appName, keyStore, null, null, + null, null, false); assertTrue(dynamoDBClientSettings.areCredentialsProvided()); assertEquals("test.keypath", dynamoDBClientSettings.getKeyPath()); @@ -77,7 +66,8 @@ public void testCredentialsProvided() { assertEquals("test.ztsurl", dynamoDBClientSettings.getZtsURL()); // Now verify that when keyStore isn't provided, trustStorePassword will be null - dynamoDBClientSettings = new DynamoDBClientSettings(certPath, domain, role, trustStore, trustStorePassword, ztsUrl, region, keyPath, appName, null, null, null, null, null); + dynamoDBClientSettings = new DynamoDBClientSettings(certPath, domain, role, trustStore, + trustStorePassword, ztsUrl, region, keyPath, appName, null, null, null, null, null, false); assertNull(dynamoDBClientSettings.getTrustStorePasswordChars()); } } \ No newline at end of file diff --git a/servers/zms/src/main/java/com/yahoo/athenz/zms/ZMSImpl.java b/servers/zms/src/main/java/com/yahoo/athenz/zms/ZMSImpl.java index a5f8d114606..9c3dc7ba7ea 100644 --- a/servers/zms/src/main/java/com/yahoo/athenz/zms/ZMSImpl.java +++ b/servers/zms/src/main/java/com/yahoo/athenz/zms/ZMSImpl.java @@ -83,7 +83,6 @@ import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; -import java.security.PrivateKey; import java.text.ParseException; import java.text.SimpleDateFormat; import java.util.*; @@ -1102,9 +1101,7 @@ void loadPrivateKeyStore() { // when signing policy files if (privateECKey == null && privateRSAKey == null) { - StringBuilder privKeyId = new StringBuilder(256); - PrivateKey pkey = keyStore.getPrivateKey(ZMSConsts.ZMS_SERVICE, serverHostName, privKeyId); - privateKey = new ServerPrivateKey(pkey, privKeyId.toString()); + privateKey = keyStore.getPrivateKey(ZMSConsts.ZMS_SERVICE, serverHostName, serverRegion, null); } else { privateKey = Objects.requireNonNullElseGet(privateECKey, () -> privateRSAKey); } diff --git a/servers/zts/pom.xml b/servers/zts/pom.xml index 10e94752c03..a1949590bc7 100644 --- a/servers/zts/pom.xml +++ b/servers/zts/pom.xml @@ -32,18 +32,6 @@ 0.9958 - - - - com.amazonaws - aws-java-sdk-bom - ${aws.version} - pom - import - - - - org.slf4j @@ -148,20 +136,24 @@ ${jersey.version} - com.amazonaws - aws-java-sdk-s3 + software.amazon.awssdk + s3 + ${aws2.version} - com.amazonaws - aws-java-sdk-sts + software.amazon.awssdk + sts + ${aws2.version} - com.amazonaws - aws-java-sdk-rds + software.amazon.awssdk + rds + ${aws2.version} - com.amazonaws - aws-java-sdk-dynamodb + software.amazon.awssdk + dynamodb + ${aws2.version} org.eclipse.jetty diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/ZTSConsts.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/ZTSConsts.java index 21a3fd0d0dd..bd1e70bd192 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/ZTSConsts.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/ZTSConsts.java @@ -186,8 +186,8 @@ public final class ZTSConsts { public static final String ZTS_PROP_AWS_RDS_IAM_ROLE = "athenz.zts.aws_rds_iam_role"; public static final String ZTS_PROP_AWS_RDS_ENGINE = "athenz.zts.aws_rds_engine"; public static final String ZTS_PROP_AWS_RDS_DATABASE = "athenz.zts.aws_rds_database"; - public static final String ZTS_PROP_AWS_RDS_MASTER_INSTANCE = "athenz.zts.aws_rds_master_instance"; - public static final String ZTS_PROP_AWS_RDS_MASTER_PORT = "athenz.zts.aws_rds_master_port"; + public static final String ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE = "athenz.zts.aws_rds_master_instance"; + public static final String ZTS_PROP_AWS_RDS_PRIMARY_PORT = "athenz.zts.aws_rds_master_port"; public static final String ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME = "athenz.zts.aws_rds_creds_refresh_time"; public static final String ZTS_SERVICE = "zts"; diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactory.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactory.java index 4cb7d2fabd8..192169d1445 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactory.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactory.java @@ -22,18 +22,19 @@ import com.yahoo.athenz.common.server.cert.CertRecordStore; import com.yahoo.athenz.common.server.cert.CertRecordStoreFactory; +import com.yahoo.athenz.common.server.util.Utils; 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 com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.db.DataSourceFactory; import com.yahoo.athenz.common.server.db.PoolableDataSource; import com.yahoo.athenz.zts.ZTSConsts; +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; public class AWSCertRecordStoreFactory implements CertRecordStoreFactory { @@ -41,23 +42,21 @@ public class AWSCertRecordStoreFactory implements CertRecordStoreFactory { private static Properties mysqlConnectionProperties = new Properties(); private static String rdsUser = null; - private static String rdsIamRole = null; - private static String rdsMaster = null; + private static String rdsPrimary = null; private int rdsPort = 3306; @Override public CertRecordStore create(PrivateKeyStore keyStore) { rdsUser = System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER); - rdsIamRole = System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_IAM_ROLE); - rdsMaster = System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_INSTANCE); - rdsPort = Integer.parseInt(System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_PORT, "3306")); + rdsPrimary = System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE); + rdsPort = Integer.parseInt(System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_PORT, "3306")); final String rdsEngine = System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_ENGINE, "mysql"); final String rdsDatabase = System.getProperty(ZTSConsts.ZTS_PROP_AWS_RDS_DATABASE, "zts_store"); - final String jdbcStore = String.format("jdbc:%s://%s:%d/%s", rdsEngine, rdsMaster, rdsPort, rdsDatabase); - String rdsToken = getAuthToken(rdsMaster, rdsPort, rdsUser, rdsIamRole); + final String jdbcStore = String.format("jdbc:%s://%s:%d/%s", rdsEngine, rdsPrimary, rdsPort, rdsDatabase); + String rdsToken = getAuthToken(rdsPrimary, rdsPort, rdsUser); if (LOGGER.isDebugEnabled()) { LOGGER.debug("Connecting to {} with auth token {}", jdbcStore, rdsToken); @@ -81,30 +80,28 @@ public CertRecordStore create(PrivateKeyStore keyStore) { return new JDBCCertRecordStore(dataSource); } - String getInstanceRegion() { - return EC2MetadataUtils.getEC2InstanceRegion(); - } + String getAuthToken(String hostname, int port, String rdsUser) { - RdsIamAuthTokenGenerator getTokenGenerator(InstanceProfileCredentialsProvider awsCredProvider) { - return RdsIamAuthTokenGenerator.builder() - .credentials(awsCredProvider) - .region(getInstanceRegion()) - .build(); - } + String authToken = null; + try (RdsClient rdsClient = RdsClient.builder().region(Utils.getAwsRegion((Region.US_EAST_1))) + .credentialsProvider(ProfileCredentialsProvider.create()).build()) { - String getAuthToken(String hostname, int port, String rdsUser, String rdsIamRole) { + RdsUtilities utilities = rdsClient.utilities(); - InstanceProfileCredentialsProvider awsCredProvider = new InstanceProfileCredentialsProvider(true); - RdsIamAuthTokenGenerator generator = getTokenGenerator(awsCredProvider); - - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("Instance {} Port {} User {} Region: {} Role: {}", hostname, port, rdsUser, - getInstanceRegion(), rdsIamRole); + GenerateAuthenticationTokenRequest tokenRequest = GenerateAuthenticationTokenRequest.builder() + .credentialsProvider(ProfileCredentialsProvider.create()) + .username(rdsUser) + .port(port) + .hostname(hostname) + .build(); + + authToken = utilities.generateAuthenticationToken(tokenRequest); + + } catch (Exception ex) { + LOGGER.error("getAuthToken: unable to generate auth token", ex); } - - return generator.getAuthToken(GetIamAuthTokenRequest.builder() - .hostname(hostname).port(port).userName(rdsUser) - .build()); + + return authToken; } class CredentialsUpdater implements Runnable { @@ -117,7 +114,7 @@ public void run() { } try { - final String rdsToken = getAuthToken(rdsMaster, rdsPort, rdsUser, rdsIamRole); + final String rdsToken = getAuthToken(rdsPrimary, rdsPort, rdsUser); mysqlConnectionProperties.setProperty(ZTSConsts.DB_PROP_PASSWORD, rdsToken); } catch (Throwable t) { diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStore.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStore.java index a8ab053d51b..f7d547e8dd9 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStore.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStore.java @@ -15,18 +15,16 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; import com.yahoo.athenz.auth.Principal; import com.yahoo.athenz.common.server.cert.CertRecordStore; import com.yahoo.athenz.common.server.cert.CertRecordStoreConnection; import com.yahoo.athenz.common.server.db.RolesProvider; import com.yahoo.athenz.common.server.notification.NotificationManager; import com.yahoo.athenz.common.utils.X509CertUtils; -import com.yahoo.athenz.zts.ResourceException; import com.yahoo.athenz.zts.notification.ZTSClientNotificationSenderImpl; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import java.security.cert.X509Certificate; @@ -38,11 +36,12 @@ public class DynamoDBCertRecordStore implements CertRecordStore { private final String tableName; private final String currentTimeIndexName; private final String hostIndexName; - private final DynamoDB dynamoDB; + private final DynamoDbClient dynamoDB; private final ZTSClientNotificationSenderImpl ztsClientNotificationSender; - public DynamoDBCertRecordStore(AmazonDynamoDB client, final String tableName, final String currentTimeIndexName, String hostIndexName, ZTSClientNotificationSenderImpl ztsClientNotificationSender) { - this.dynamoDB = new DynamoDB(client); + public DynamoDBCertRecordStore(DynamoDbClient client, final String tableName, final String currentTimeIndexName, + String hostIndexName, ZTSClientNotificationSenderImpl ztsClientNotificationSender) { + this.dynamoDB = client; this.tableName = tableName; this.currentTimeIndexName = currentTimeIndexName; this.hostIndexName = hostIndexName; @@ -51,12 +50,7 @@ public DynamoDBCertRecordStore(AmazonDynamoDB client, final String tableName, fi @Override public CertRecordStoreConnection getConnection() { - try { - return new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, currentTimeIndexName, hostIndexName); - } catch (Exception ex) { - LOGGER.error("getConnection: {}", ex.getMessage()); - throw new ResourceException(ResourceException.SERVICE_UNAVAILABLE, ex.getMessage()); - } + return new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, currentTimeIndexName, hostIndexName); } @Override @@ -69,12 +63,13 @@ public void clearConnections() { @Override public void log(final Principal principal, final String ip, final String provider, - final String instanceId, final X509Certificate x509Cert) { + final String instanceId, final X509Certificate x509Cert) { X509CertUtils.logCert(CERTLOGGER, principal, ip, provider, instanceId, x509Cert); } @Override - public boolean enableNotifications(NotificationManager notificationManager, RolesProvider rolesProvider, final String serverName) { + public boolean enableNotifications(NotificationManager notificationManager, RolesProvider rolesProvider, + final String serverName) { if (ztsClientNotificationSender != null) { return ztsClientNotificationSender.init(notificationManager, rolesProvider, serverName); } else { diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnection.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnection.java index eb29a4186a9..3a14d8b8723 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnection.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnection.java @@ -19,17 +19,12 @@ import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; -import com.amazonaws.services.dynamodbv2.document.*; -import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec; -import com.amazonaws.services.dynamodbv2.document.utils.ValueMap; -import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException; import com.yahoo.athenz.common.server.cert.CertRecordStoreConnection; import com.yahoo.athenz.common.server.cert.X509CertRecord; import com.yahoo.athenz.zts.ZTSConsts; - -import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; +import software.amazon.awssdk.services.dynamodb.model.*; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import com.yahoo.athenz.zts.utils.DynamoDBUtils; import com.yahoo.athenz.zts.utils.RetryDynamoDBCommand; @@ -61,7 +56,7 @@ public class DynamoDBCertRecordStoreConnection implements CertRecordStoreConnect private static final String KEY_REGISTER_TIME = "registerTime"; private static final String KEY_SVC_DATA_UPDATE_TIME = "svcDataUpdateTime"; - // the configuration setting is in hours so we'll automatically + // the configuration setting is in hours, so we'll automatically // convert into seconds since that's what dynamoDB needs // we need to expire records in 30 days private static final Long EXPIRY_HOURS = Long.parseLong( @@ -73,22 +68,24 @@ public class DynamoDBCertRecordStoreConnection implements CertRecordStoreConnect private static long expiryTime = 3660 * EXPIRY_HOURS; - private Table table; - private final Index currentTimeIndex; - private final Index hostNameIndex; + private final String tableName; + private final String currentTimeIndexName; + private final String hostNameIndexName; + private final DynamoDbClient dynamoDB; private final DynamoDBNotificationsHelper dynamoDBNotificationsHelper = new DynamoDBNotificationsHelper(); - private final RetryDynamoDBCommand getItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); - private final RetryDynamoDBCommand updateItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); - private final RetryDynamoDBCommand putItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); - private final RetryDynamoDBCommand deleteItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); - private final RetryDynamoDBCommand> itemCollectionRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); - - - public DynamoDBCertRecordStoreConnection(DynamoDB dynamoDB, final String tableName, String currentTimeIndexName, String hostIndexName) { - this.table = dynamoDB.getTable(tableName); - this.currentTimeIndex = table.getIndex(currentTimeIndexName); - this.hostNameIndex = table.getIndex(hostIndexName); + private final RetryDynamoDBCommand getItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + private final RetryDynamoDBCommand updateItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + private final RetryDynamoDBCommand putItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + private final RetryDynamoDBCommand deleteItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + private final RetryDynamoDBCommand itemCollectionRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + + public DynamoDBCertRecordStoreConnection(DynamoDbClient dynamoDB, final String tableName, + String currentTimeIndexName, String hostNameIndexName) { + this.tableName = tableName; + this.dynamoDB = dynamoDB; + this.currentTimeIndexName = currentTimeIndexName; + this.hostNameIndexName = hostNameIndexName; } @Override @@ -103,9 +100,19 @@ public void close() { public X509CertRecord getX509CertRecord(String provider, String instanceId, String service) { final String primaryKey = getPrimaryKey(provider, instanceId, service); + + HashMap keyToGet = new HashMap<>(); + keyToGet.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); + try { - Item item = getItemRetryDynamoDBCommand.run(() -> table.getItem(KEY_PRIMARY, primaryKey)); - if (item == null) { + GetItemResponse response = getItemRetryDynamoDBCommand.run(() -> dynamoDB.getItem(request)); + Map item = response.item(); + if (item == null || item.isEmpty()) { LOGGER.error("DynamoDB Get Error for {}: item not found", primaryKey); return null; } @@ -116,29 +123,23 @@ public X509CertRecord getX509CertRecord(String provider, String instanceId, Stri } } - private X509CertRecord itemToX509CertRecord(Item item) { - boolean clientCert; - try { - clientCert = item.getBoolean(KEY_CLIENT_CERT); - } catch (Exception ex) { - LOGGER.warn("clientCert for item doesn't exist. Will set it to false. Item: {}", item.toString()); - clientCert = false; - } + private X509CertRecord itemToX509CertRecord(Map item) { + X509CertRecord certRecord = new X509CertRecord(); - certRecord.setProvider(item.getString(KEY_PROVIDER)); - certRecord.setInstanceId(item.getString(KEY_INSTANCE_ID)); - certRecord.setService(item.getString(KEY_SERVICE)); - certRecord.setCurrentSerial(item.getString(KEY_CURRENT_SERIAL)); - certRecord.setCurrentIP(item.getString(KEY_CURRENT_IP)); + certRecord.setProvider(DynamoDBUtils.getString(item, KEY_PROVIDER)); + certRecord.setInstanceId(DynamoDBUtils.getString(item, KEY_INSTANCE_ID)); + certRecord.setService(DynamoDBUtils.getString(item, KEY_SERVICE)); + certRecord.setCurrentSerial(DynamoDBUtils.getString(item, KEY_CURRENT_SERIAL)); + certRecord.setCurrentIP(DynamoDBUtils.getString(item, KEY_CURRENT_IP)); certRecord.setCurrentTime(DynamoDBUtils.getDateFromItem(item, KEY_CURRENT_TIME)); - certRecord.setPrevSerial(item.getString(KEY_PREV_SERIAL)); - certRecord.setPrevIP(item.getString(KEY_PREV_IP)); + certRecord.setPrevSerial(DynamoDBUtils.getString(item, KEY_PREV_SERIAL)); + certRecord.setPrevIP(DynamoDBUtils.getString(item, KEY_PREV_IP)); certRecord.setPrevTime(DynamoDBUtils.getDateFromItem(item, KEY_PREV_TIME)); - certRecord.setClientCert(clientCert); + certRecord.setClientCert(DynamoDBUtils.getBoolean(item, KEY_CLIENT_CERT)); certRecord.setLastNotifiedTime(DynamoDBUtils.getDateFromItem(item, KEY_LAST_NOTIFIED_TIME)); - certRecord.setLastNotifiedServer(item.getString(KEY_LAST_NOTIFIED_SERVER)); + certRecord.setLastNotifiedServer(DynamoDBUtils.getString(item, KEY_LAST_NOTIFIED_SERVER)); certRecord.setExpiryTime(DynamoDBUtils.getDateFromItem(item, KEY_EXPIRY_TIME)); - certRecord.setHostName(item.getString(KEY_HOSTNAME)); + certRecord.setHostName(DynamoDBUtils.getString(item, KEY_HOSTNAME)); certRecord.setSvcDataUpdateTime(DynamoDBUtils.getDateFromItem(item, KEY_SVC_DATA_UPDATE_TIME)); return certRecord; } @@ -156,33 +157,41 @@ public boolean updateX509CertRecord(X509CertRecord certRecord) { certRecord.setSvcDataUpdateTime(new Date()); } - String hostName = certRecord.getHostName(); // Prevent inserting null values in hostName as the hostName-Index will not allow it + + String hostName = certRecord.getHostName(); if (StringUtil.isEmpty(hostName)) { hostName = primaryKey; } + HashMap itemKey = new HashMap<>(); + itemKey.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + try { - UpdateItemSpec updateItemSpec = new UpdateItemSpec() - .withPrimaryKey(KEY_PRIMARY, primaryKey) - .withAttributeUpdate( - new AttributeUpdate(KEY_INSTANCE_ID).put(certRecord.getInstanceId()), - new AttributeUpdate(KEY_PROVIDER).put(certRecord.getProvider()), - new AttributeUpdate(KEY_SERVICE).put(certRecord.getService()), - new AttributeUpdate(KEY_CURRENT_SERIAL).put(certRecord.getCurrentSerial()), - new AttributeUpdate(KEY_CURRENT_IP).put(certRecord.getCurrentIP()), - new AttributeUpdate(KEY_CURRENT_TIME).put(DynamoDBUtils.getLongFromDate(certRecord.getCurrentTime())), - new AttributeUpdate(KEY_CURRENT_DATE).put(DynamoDBUtils.getIso8601FromDate(certRecord.getCurrentTime())), - new AttributeUpdate(KEY_PREV_SERIAL).put(certRecord.getPrevSerial()), - new AttributeUpdate(KEY_PREV_IP).put(certRecord.getPrevIP()), - new AttributeUpdate(KEY_PREV_TIME).put(DynamoDBUtils.getLongFromDate(certRecord.getPrevTime())), - new AttributeUpdate(KEY_CLIENT_CERT).put(certRecord.getClientCert()), - new AttributeUpdate(KEY_TTL).put(certRecord.getCurrentTime().getTime() / 1000L + expiryTime), - new AttributeUpdate(KEY_SVC_DATA_UPDATE_TIME).put(DynamoDBUtils.getLongFromDate(certRecord.getSvcDataUpdateTime())), - new AttributeUpdate(KEY_EXPIRY_TIME).put(DynamoDBUtils.getLongFromDate(certRecord.getExpiryTime())), - new AttributeUpdate(KEY_HOSTNAME).put(hostName) - ); - updateItemRetryDynamoDBCommand.run(() -> table.updateItem(updateItemSpec)); + HashMap updatedValues = new HashMap<>(); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_INSTANCE_ID, certRecord.getInstanceId()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_PROVIDER, certRecord.getProvider()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_SERVICE, certRecord.getService()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_CURRENT_SERIAL, certRecord.getCurrentSerial()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_CURRENT_IP, certRecord.getCurrentIP()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_CURRENT_TIME, certRecord.getCurrentTime()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_CURRENT_DATE, DynamoDBUtils.getIso8601FromDate(certRecord.getCurrentTime())); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_PREV_SERIAL, certRecord.getPrevSerial()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_PREV_IP, certRecord.getPrevIP()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_PREV_TIME, certRecord.getPrevTime()); + DynamoDBUtils.updateItemBoolValue(updatedValues, KEY_CLIENT_CERT, certRecord.getClientCert()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_TTL, certRecord.getCurrentTime().getTime() / 1000L + expiryTime); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_SVC_DATA_UPDATE_TIME, certRecord.getSvcDataUpdateTime()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_EXPIRY_TIME, certRecord.getExpiryTime()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_HOSTNAME, hostName); + + UpdateItemRequest request = UpdateItemRequest.builder() + .tableName(tableName) + .key(itemKey) + .attributeUpdates(updatedValues) + .build(); + + updateItemRetryDynamoDBCommand.run(() -> dynamoDB.updateItem(request)); return true; } catch (Exception ex) { LOGGER.error("DynamoDB Update Error for {}: {}/{}", primaryKey, ex.getClass(), ex.getMessage()); @@ -195,31 +204,39 @@ public boolean insertX509CertRecord(X509CertRecord certRecord) { final String primaryKey = getPrimaryKey(certRecord.getProvider(), certRecord.getInstanceId(), certRecord.getService()); - String hostName = certRecord.getHostName(); + // Prevent inserting null values in hostName as the hostName-Index will not allow it + + String hostName = certRecord.getHostName(); if (StringUtil.isEmpty(hostName)) { hostName = primaryKey; } try { - Item item = new Item() - .withPrimaryKey(KEY_PRIMARY, primaryKey) - .withString(KEY_INSTANCE_ID, certRecord.getInstanceId()) - .withString(KEY_PROVIDER, certRecord.getProvider()) - .withString(KEY_SERVICE, certRecord.getService()) - .withString(KEY_CURRENT_SERIAL, certRecord.getCurrentSerial()) - .withString(KEY_CURRENT_IP, certRecord.getCurrentIP()) - .with(KEY_CURRENT_TIME, DynamoDBUtils.getLongFromDate(certRecord.getCurrentTime())) - .withString(KEY_CURRENT_DATE, DynamoDBUtils.getIso8601FromDate(certRecord.getCurrentTime())) - .withString(KEY_PREV_SERIAL, certRecord.getPrevSerial()) - .withString(KEY_PREV_IP, certRecord.getPrevIP()) - .with(KEY_PREV_TIME, DynamoDBUtils.getLongFromDate(certRecord.getPrevTime())) - .withBoolean(KEY_CLIENT_CERT, certRecord.getClientCert()) - .withLong(KEY_TTL, certRecord.getCurrentTime().getTime() / 1000L + expiryTime) - .with(KEY_EXPIRY_TIME, DynamoDBUtils.getLongFromDate(certRecord.getExpiryTime())) - .with(KEY_SVC_DATA_UPDATE_TIME, DynamoDBUtils.getLongFromDate(certRecord.getSvcDataUpdateTime())) - .withLong(KEY_REGISTER_TIME, System.currentTimeMillis()) - .with(KEY_HOSTNAME, hostName); - putItemRetryDynamoDBCommand.run(() -> table.putItem(item)); + HashMap itemValues = new HashMap<>(); + itemValues.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + itemValues.put(KEY_INSTANCE_ID, AttributeValue.fromS(certRecord.getInstanceId())); + itemValues.put(KEY_PROVIDER, AttributeValue.fromS(certRecord.getProvider())); + itemValues.put(KEY_SERVICE, AttributeValue.fromS(certRecord.getService())); + itemValues.put(KEY_CURRENT_SERIAL, AttributeValue.fromS(certRecord.getCurrentSerial())); + itemValues.put(KEY_CURRENT_IP, AttributeValue.fromS(certRecord.getCurrentIP())); + itemValues.put(KEY_CURRENT_TIME, AttributeValue.fromN(DynamoDBUtils.getNumberFromDate(certRecord.getCurrentTime()))); + itemValues.put(KEY_CURRENT_DATE, AttributeValue.fromS(DynamoDBUtils.getIso8601FromDate(certRecord.getCurrentTime()))); + itemValues.put(KEY_PREV_SERIAL, AttributeValue.fromS(certRecord.getPrevSerial())); + itemValues.put(KEY_PREV_IP, AttributeValue.fromS(certRecord.getPrevIP())); + itemValues.put(KEY_PREV_TIME, AttributeValue.fromN(DynamoDBUtils.getNumberFromDate(certRecord.getPrevTime()))); + itemValues.put(KEY_CLIENT_CERT, AttributeValue.fromBool(certRecord.getClientCert())); + itemValues.put(KEY_TTL, AttributeValue.fromN(Long.toString(certRecord.getCurrentTime().getTime() / 1000L + expiryTime))); + itemValues.put(KEY_EXPIRY_TIME, AttributeValue.fromN(DynamoDBUtils.getNumberFromDate(certRecord.getExpiryTime()))); + itemValues.put(KEY_SVC_DATA_UPDATE_TIME, AttributeValue.fromN(DynamoDBUtils.getNumberFromDate(certRecord.getSvcDataUpdateTime()))); + itemValues.put(KEY_REGISTER_TIME, AttributeValue.fromN(String.valueOf(System.currentTimeMillis()))); + itemValues.put(KEY_HOSTNAME, AttributeValue.fromS(hostName)); + + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(itemValues) + .build(); + + putItemRetryDynamoDBCommand.run(() -> dynamoDB.putItem(request)); return true; } catch (Exception ex) { LOGGER.error("DynamoDB Insert Error for {}: {}/{}", primaryKey, ex.getClass(), ex.getMessage()); @@ -231,10 +248,17 @@ public boolean insertX509CertRecord(X509CertRecord certRecord) { public boolean deleteX509CertRecord(String provider, String instanceId, String service) { final String primaryKey = getPrimaryKey(provider, instanceId, service); + + HashMap keyToGet = new HashMap<>(); + keyToGet.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + + DeleteItemRequest request = DeleteItemRequest.builder() + .tableName(tableName) + .key(keyToGet) + .build(); + try { - DeleteItemSpec deleteItemSpec = new DeleteItemSpec() - .withPrimaryKey(KEY_PRIMARY, primaryKey); - deleteItemRetryDynamoDBCommand.run(() -> table.deleteItem(deleteItemSpec)); + deleteItemRetryDynamoDBCommand.run(() -> dynamoDB.deleteItem(request)); return true; } catch (Exception ex) { LOGGER.error("DynamoDB Delete Error for {}: {}/{}", primaryKey, ex.getClass(), ex.getMessage()); @@ -246,7 +270,7 @@ public boolean deleteX509CertRecord(String provider, String instanceId, String s public int deleteExpiredX509CertRecords(int expiryTimeMins) { // with dynamo db there is no need to manually expunge expired - // record since we have the TTL option enabled for our table + // record since we have the TTL option enabled for our table, // and we just need to make sure the attribute is updated with // the epoch time + timeout seconds when it should retire @@ -254,10 +278,9 @@ public int deleteExpiredX509CertRecords(int expiryTimeMins) { } @Override - public List updateUnrefreshedCertificatesNotificationTimestamp(String lastNotifiedServer, - long lastNotifiedTime, - String provider) { - List items = getUnrefreshedCertsRecords(lastNotifiedTime, provider); + public List updateUnrefreshedCertificatesNotificationTimestamp(final String lastNotifiedServer, + long lastNotifiedTime, final String provider) { + List> items = getUnrefreshedCertsRecords(lastNotifiedTime, provider); return updateLastNotified(lastNotifiedServer, lastNotifiedTime, items); } @@ -265,20 +288,20 @@ private String getPrimaryKey(final String provider, final String instanceId, fin return provider + ":" + service + ":" + instanceId; } - private List updateLastNotified(String lastNotifiedServer, long lastNotifiedTime, List items) { + private List updateLastNotified(String lastNotifiedServer, long lastNotifiedTime, + List> items) { + long yesterday = lastNotifiedTime - TimeUnit.DAYS.toMillis(1); List updatedRecords = new ArrayList<>(); - for (Item item : items) { + for (Map item : items) { try { - Item updatedItem = dynamoDBNotificationsHelper.updateLastNotifiedItem(lastNotifiedServer, lastNotifiedTime, yesterday, item, KEY_PRIMARY, table); + Map updatedItem = dynamoDBNotificationsHelper.updateLastNotifiedItem( + lastNotifiedServer, lastNotifiedTime, yesterday, item, KEY_PRIMARY, tableName, dynamoDB); if (isRecordUpdatedWithNotificationTimeAndServer(lastNotifiedServer, lastNotifiedTime, updatedItem)) { - X509CertRecord x509CertRecord = itemToX509CertRecord(updatedItem); - updatedRecords.add(x509CertRecord); + updatedRecords.add(itemToX509CertRecord(updatedItem)); } - } catch (ConditionalCheckFailedException ex) { - // This error appears when the update didn't work because it was already updated by another server. We can ignore it. } catch (Exception ex) { LOGGER.error("DynamoDB updateLastNotified failed for item: {}, error: {}", item.toString(), ex.getMessage()); } @@ -287,79 +310,93 @@ private List updateLastNotified(String lastNotifiedServer, long return updatedRecords; } - private boolean isRecordUpdatedWithNotificationTimeAndServer(String lastNotifiedServer, long lastNotifiedTime, Item updatedItem) { + private boolean isRecordUpdatedWithNotificationTimeAndServer(String lastNotifiedServer, long lastNotifiedTime, + Map updatedItem) { + return updatedItem != null && - updatedItem.getLong(KEY_LAST_NOTIFIED_TIME) == lastNotifiedTime && - updatedItem.getString(KEY_LAST_NOTIFIED_SERVER) != null && - updatedItem.getString(KEY_LAST_NOTIFIED_SERVER).equals(lastNotifiedServer); + DynamoDBUtils.getLong(updatedItem, KEY_LAST_NOTIFIED_TIME) == lastNotifiedTime && + lastNotifiedServer.equals(DynamoDBUtils.getString(updatedItem, KEY_LAST_NOTIFIED_SERVER)); } - private List getUnrefreshedCertsRecords(long lastNotifiedTime, String provider) { + private List> getUnrefreshedCertsRecords(long lastNotifiedTime, String provider) { + long yesterday = lastNotifiedTime - TimeUnit.DAYS.toMillis(1); long unrefreshedCertsRangeBegin = lastNotifiedTime - TimeUnit.HOURS.toMillis(EXPIRY_HOURS); long unrefreshedCertsRangeEnd = lastNotifiedTime - TimeUnit.HOURS.toMillis(EXPIRY_HOURS_GRACE); - List items = new ArrayList<>(); + List> items = new ArrayList<>(); List unrefreshedCertDates = DynamoDBUtils.getISODatesByRange(unrefreshedCertsRangeBegin, unrefreshedCertsRangeEnd); for (String unrefreshedCertDate : unrefreshedCertDates) { items.addAll(getUnrefreshedCertRecordsByDate(provider, yesterday, unrefreshedCertDate)); } - // Filter outdated records from before re-bootstrapping (another record exist with a new uuid) + // Filter outdated records from before re-bootstrapping + // (another record exist with a new uuid) + items = items.stream() - .filter(item -> (mostUpdatedHostRecord(item))) + .filter(this::mostUpdatedHostRecord) .collect(Collectors.toList()); return items; } - private boolean mostUpdatedHostRecord(Item recordToCheck) { + private boolean mostUpdatedHostRecord(Map recordToCheck) { try { - // Query all records with the same hostName / provider / service as recordToCheck - QuerySpec spec = new QuerySpec() - .withKeyConditionExpression("hostName = :v_host_name") - .withFilterExpression("attribute_exists(provider) AND provider = :v_provider AND attribute_exists(service) AND service = :v_service") - .withValueMap(new ValueMap() - .withString(":v_host_name", recordToCheck.getString(KEY_HOSTNAME)) - .withString(":v_provider", recordToCheck.getString(KEY_PROVIDER)) - .withString(":v_service", recordToCheck.getString(KEY_SERVICE)) - ); - - ItemCollection outcome = itemCollectionRetryDynamoDBCommand.run(() -> hostNameIndex.query(spec)); - List allRecordsWithHost = new ArrayList<>(); - for (Item item : outcome) { - allRecordsWithHost.add(item); - } + // Set up mapping of the partition name with the value. + + HashMap attrValues = new HashMap<>(); + attrValues.put(":v_host_name", AttributeValue.fromS(DynamoDBUtils.getString(recordToCheck, KEY_HOSTNAME))); + attrValues.put(":v_provider", AttributeValue.fromS(DynamoDBUtils.getString(recordToCheck, KEY_PROVIDER))); + attrValues.put(":v_service", AttributeValue.fromS(DynamoDBUtils.getString(recordToCheck, KEY_SERVICE))); + + QueryRequest request = QueryRequest.builder() + .tableName(tableName) + .indexName(hostNameIndexName) + .keyConditionExpression("hostName = :v_host_name") + .filterExpression("attribute_exists(provider) AND provider = :v_provider AND attribute_exists(service) AND service = :v_service") + .expressionAttributeValues(attrValues) + .build(); + + QueryResponse response = itemCollectionRetryDynamoDBCommand.run(() -> dynamoDB.query(request)); + List> allRecordsWithHost = new ArrayList<>(response.items()); // Verify recordToCheck is the most updated record with this hostName - return dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(recordToCheck, allRecordsWithHost, KEY_CURRENT_TIME, KEY_PRIMARY); + return dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(recordToCheck, allRecordsWithHost, + KEY_CURRENT_TIME, KEY_PRIMARY); } catch (Exception ex) { LOGGER.error("DynamoDB mostUpdatedHostRecord failed for item: {}, error: {}", recordToCheck.toString(), ex.getMessage()); return false; } } - private List getUnrefreshedCertRecordsByDate(String provider, long yesterday, String unrefreshedCertDate) { + private List> getUnrefreshedCertRecordsByDate(final String provider, long yesterday, + final String unrefreshedCertDate) { + + List> items = new ArrayList<>(); try { - QuerySpec spec = new QuerySpec() - .withKeyConditionExpression("currentDate = :v_current_date") - .withFilterExpression("provider = :v_provider AND attribute_exists(hostName) AND (attribute_not_exists(lastNotifiedTime) OR lastNotifiedTime < :v_last_notified)") - .withValueMap(new ValueMap() - .withString(":v_current_date", unrefreshedCertDate) - .withNumber(":v_last_notified", yesterday) - .withString(":v_provider", provider)); - - ItemCollection outcome = itemCollectionRetryDynamoDBCommand.run(() -> currentTimeIndex.query(spec)); - List items = new ArrayList<>(); - for (Item item : outcome) { - items.add(item); - } - return items; + // Set up mapping of the partition name with the value. + + HashMap attrValues = new HashMap<>(); + attrValues.put(":v_current_date", AttributeValue.fromS(unrefreshedCertDate)); + attrValues.put(":v_last_notified", AttributeValue.fromN(String.valueOf(yesterday))); + attrValues.put(":v_provider", AttributeValue.fromS(provider)); + + QueryRequest request = QueryRequest.builder() + .tableName(tableName) + .indexName(currentTimeIndexName) + .keyConditionExpression("currentDate = :v_current_date") + .filterExpression("provider = :v_provider AND attribute_exists(hostName) AND (attribute_not_exists(lastNotifiedTime) OR lastNotifiedTime < :v_last_notified)") + .expressionAttributeValues(attrValues) + .build(); + + QueryResponse response = itemCollectionRetryDynamoDBCommand.run(() -> dynamoDB.query(request)); + items.addAll(response.items()); } catch (Exception ex) { - LOGGER.error("DynamoDB getUnrefreshedCertRecordsByDate failed for provider: {}, date: {} error: {}", provider, unrefreshedCertDate, ex.getMessage()); + LOGGER.error("DynamoDB getUnrefreshedCertRecordsByDate failed for provider: {}, date: {} error: {}", + provider, unrefreshedCertDate, ex.getMessage()); } - return new ArrayList<>(); + return items; } } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactory.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactory.java index 1067beedcbe..9fdecd87b41 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactory.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactory.java @@ -15,7 +15,6 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.cert.CertRecordStore; import com.yahoo.athenz.common.server.cert.CertRecordStoreFactory; @@ -24,42 +23,47 @@ import com.yahoo.athenz.zts.ResourceException; import com.yahoo.athenz.zts.ZTSConsts; import com.yahoo.athenz.zts.notification.ZTSClientNotificationSenderImpl; +import org.eclipse.jetty.util.StringUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; public class DynamoDBCertRecordStoreFactory implements CertRecordStoreFactory { + private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBCertRecordStoreFactory.class); @Override public CertRecordStore create(PrivateKeyStore keyStore) { final String tableName = System.getProperty(ZTSConsts.ZTS_PROP_CERT_DYNAMODB_TABLE_NAME); - if (tableName == null || tableName.isEmpty()) { + if (StringUtil.isEmpty(tableName)) { LOGGER.error("Cert Store DynamoDB table name not specified"); throw new ResourceException(ResourceException.SERVICE_UNAVAILABLE, "DynamoDB table name not specified"); } final String currentTimeIndexName = System.getProperty(ZTSConsts.ZTS_PROP_CERT_DYNAMODB_INDEX_CURRENT_TIME_NAME); - if (currentTimeIndexName == null || currentTimeIndexName.isEmpty()) { + if (StringUtil.isEmpty(currentTimeIndexName)) { LOGGER.error("Cert Store DynamoDB index current-time not specified"); throw new ResourceException(ResourceException.SERVICE_UNAVAILABLE, "DynamoDB index current-time not specified"); } final String hostNameIndex = System.getProperty(ZTSConsts.ZTS_PROP_CERT_DYNAMODB_INDEX_HOST_NAME); - if (hostNameIndex == null || hostNameIndex.isEmpty()) { + if (StringUtil.isEmpty(hostNameIndex)) { LOGGER.error("Cert Store DynamoDB index host-name not specified"); throw new ResourceException(ResourceException.SERVICE_UNAVAILABLE, "DynamoDB index host-name not specified"); } - ZTSClientNotificationSenderImpl ztsClientNotificationSender = new ZTSClientNotificationSenderImpl(); - AmazonDynamoDB client = getDynamoDBClient(ztsClientNotificationSender, keyStore); + DynamoDbClient client = getDynamoDBClient(ztsClientNotificationSender, keyStore); return new DynamoDBCertRecordStore(client, tableName, currentTimeIndexName, hostNameIndex, ztsClientNotificationSender); } - AmazonDynamoDB getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { + DynamoDbClient getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, + PrivateKeyStore keyStore) { + DynamoDBClientFetcher dynamoDBClientFetcher = DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); ZTSDynamoDBClientSettingsFactory ztsDynamoDBClientSettingsFactory = new ZTSDynamoDBClientSettingsFactory(keyStore); - return dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings()).getAmazonDynamoDB(); + return dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings(false)).getDynamoDbClient(); } } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelper.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelper.java index 218f4eb45fc..67a9ca1aa7a 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelper.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelper.java @@ -16,47 +16,61 @@ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.document.Item; -import com.amazonaws.services.dynamodbv2.document.Table; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; -import com.amazonaws.services.dynamodbv2.document.utils.ValueMap; -import com.amazonaws.services.dynamodbv2.model.ReturnValue; import com.yahoo.athenz.zts.utils.RetryDynamoDBCommand; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.ReturnValue; +import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.TimeoutException; public class DynamoDBNotificationsHelper { - public boolean isMostUpdatedRecordBasedOnAttribute(Item recordToCheck, List allRecordsWithHost, String dateKey, String primaryKey) { - Item mostUpdatedHostRecord = allRecordsWithHost.stream() + public boolean isMostUpdatedRecordBasedOnAttribute(Map recordToCheck, + List> allRecordsWithHost, String dateKey, String primaryKey) { + Map mostUpdatedHostRecord = allRecordsWithHost.stream() .reduce((record1, record2) -> { - if (record1.isNull(dateKey) || record1.get(dateKey) == null) { + if (record1.get(dateKey) == null) { return record2; } - if (record2.isNull(dateKey) || record2.get(dateKey) == null) { + if (record2.get(dateKey) == null) { return record1; } - return record1.getLong(dateKey) > record2.getLong(dateKey) ? record1 : record2; + return Long.parseLong(record1.get(dateKey).n()) > Long.parseLong(record2.get(dateKey).n()) ? record1 : record2; }).get(); return recordToCheck.get(primaryKey).equals(mostUpdatedHostRecord.get(primaryKey)); } - public Item updateLastNotifiedItem(String lastNotifiedServer, long lastNotifiedTime, long yesterday, Item item, String primaryKey, Table table) throws InterruptedException, TimeoutException { - RetryDynamoDBCommand retryDynamoDBCommand = new RetryDynamoDBCommand<>(); + public Map updateLastNotifiedItem(final String lastNotifiedServer, long lastNotifiedTime, + long yesterday, Map item, final String primaryKey, final String tableName, + DynamoDbClient dbClient) throws InterruptedException, TimeoutException { + + RetryDynamoDBCommand> retryDynamoDBCommand = new RetryDynamoDBCommand<>(); return retryDynamoDBCommand.run(() -> { + + HashMap itemKey = new HashMap<>(); + itemKey.put(primaryKey, item.get(primaryKey)); + + HashMap attrValues = new HashMap<>(); + attrValues.put(":lastNotifiedTimeVal", AttributeValue.fromN(String.valueOf(lastNotifiedTime))); + attrValues.put(":v_yesterday", AttributeValue.fromN(String.valueOf(yesterday))); + attrValues.put(":lastNotifiedServerVal", AttributeValue.fromS(lastNotifiedServer)); + // For each item, update lastNotifiedTime and lastNotifiedServer (unless they were already updated) - UpdateItemSpec updateItemSpec = new UpdateItemSpec().withPrimaryKey(primaryKey, item.getString(primaryKey)) - .withReturnValues(ReturnValue.ALL_NEW) - .withUpdateExpression("set lastNotifiedTime = :lastNotifiedTimeVal, lastNotifiedServer = :lastNotifiedServerVal") - .withConditionExpression("attribute_not_exists(lastNotifiedTime) OR lastNotifiedTime < :v_yesterday") - .withValueMap(new ValueMap() - .with(":lastNotifiedTimeVal", lastNotifiedTime) - .withNumber(":v_yesterday", yesterday) - .withString(":lastNotifiedServerVal", lastNotifiedServer)); - - return table.updateItem(updateItemSpec).getItem(); + UpdateItemRequest request = UpdateItemRequest.builder() + .tableName(tableName) + .key(itemKey) + .returnValues(ReturnValue.ALL_NEW) + .updateExpression("set lastNotifiedTime = :lastNotifiedTimeVal, lastNotifiedServer = :lastNotifiedServerVal") + .conditionExpression("attribute_not_exists(lastNotifiedTime) OR lastNotifiedTime < :v_yesterday") + .expressionAttributeValues(attrValues) + .build(); + + return dbClient.updateItem(request).attributes(); }); } } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStore.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStore.java index b1381693adb..8c8c3d66f9f 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStore.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStore.java @@ -15,42 +15,36 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; import com.yahoo.athenz.auth.Principal; import com.yahoo.athenz.common.server.db.RolesProvider; import com.yahoo.athenz.common.server.notification.NotificationManager; import com.yahoo.athenz.common.server.ssh.SSHRecordStore; import com.yahoo.athenz.common.server.ssh.SSHRecordStoreConnection; import com.yahoo.athenz.common.utils.X509CertUtils; -import com.yahoo.athenz.zts.ResourceException; import com.yahoo.athenz.zts.notification.ZTSClientNotificationSenderImpl; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; public class DynamoDBSSHRecordStore implements SSHRecordStore { private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBSSHRecordStore.class); private static final Logger SSHLOGGER = LoggerFactory.getLogger("SSHCertLogger"); - private final DynamoDB dynamoDB; + private final DynamoDbClient dynamoDB; private final String tableName; private final ZTSClientNotificationSenderImpl ztsClientNotificationSender; - public DynamoDBSSHRecordStore(AmazonDynamoDB client, final String tableName, ZTSClientNotificationSenderImpl ztsClientNotificationSender) { - this.dynamoDB = new DynamoDB(client); + public DynamoDBSSHRecordStore(DynamoDbClient client, final String tableName, + ZTSClientNotificationSenderImpl ztsClientNotificationSender) { + this.dynamoDB = client; this.tableName = tableName; this.ztsClientNotificationSender = ztsClientNotificationSender; } @Override public SSHRecordStoreConnection getConnection() { - try { - return new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); - } catch (Exception ex) { - LOGGER.error("getConnection: {}", ex.getMessage()); - throw new ResourceException(ResourceException.SERVICE_UNAVAILABLE, ex.getMessage()); - } + return new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); } @Override @@ -63,12 +57,13 @@ public void clearConnections() { @Override public void log(final Principal principal, final String ip, final String service, - final String instanceId) { + final String instanceId) { X509CertUtils.logSSH(SSHLOGGER, principal, ip, service, instanceId); } @Override - public boolean enableNotifications(NotificationManager notificationManager, RolesProvider rolesProvider, final String serverName) { + public boolean enableNotifications(NotificationManager notificationManager, RolesProvider rolesProvider, + final String serverName) { if (ztsClientNotificationSender != null) { return ztsClientNotificationSender.init(notificationManager, rolesProvider, serverName); } else { diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnection.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnection.java index 9d053d4e030..bc01a73d0ed 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnection.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnection.java @@ -15,17 +15,17 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.document.AttributeUpdate; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; -import com.amazonaws.services.dynamodbv2.document.Item; -import com.amazonaws.services.dynamodbv2.document.Table; -import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; import com.yahoo.athenz.common.server.ssh.SSHCertRecord; import com.yahoo.athenz.common.server.ssh.SSHRecordStoreConnection; import com.yahoo.athenz.zts.ZTSConsts; +import com.yahoo.athenz.zts.utils.DynamoDBUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.*; + +import java.util.HashMap; +import java.util.Map; public class DynamoDBSSHRecordStoreConnection implements SSHRecordStoreConnection { @@ -45,10 +45,12 @@ public class DynamoDBSSHRecordStoreConnection implements SSHRecordStoreConnectio private static long expiryTime = 3660 * Long.parseLong( System.getProperty(ZTSConsts.ZTS_PROP_SSH_DYNAMODB_ITEM_TTL_HOURS, "720")); - private Table table; + private final DynamoDbClient dynamoDB; + private final String tableName; - public DynamoDBSSHRecordStoreConnection(DynamoDB dynamoDB, final String tableName) { - table = dynamoDB.getTable(tableName); + public DynamoDBSSHRecordStoreConnection(DynamoDbClient dynamoDB, final String tableName) { + this.dynamoDB = dynamoDB; + this.tableName = tableName; } @Override @@ -63,19 +65,31 @@ public void close() { public SSHCertRecord getSSHCertRecord(String instanceId, String service) { final String primaryKey = getPrimaryKey(instanceId, service); + + HashMap keyToGet = new HashMap<>(); + keyToGet.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); + try { - Item item = table.getItem(KEY_PRIMARY, primaryKey); - if (item == null) { + GetItemResponse response = dynamoDB.getItem(request); + Map item = response.item(); + if (item == null || item.isEmpty()) { LOGGER.error("DynamoDB Get Error for {}: item not found", primaryKey); return null; } + SSHCertRecord certRecord = new SSHCertRecord(); certRecord.setInstanceId(instanceId); certRecord.setService(service); - certRecord.setPrincipals(item.getString(KEY_PRINCIPALS)); - certRecord.setClientIP(item.getString(KEY_CLIENT_IP)); - certRecord.setPrivateIP(item.getString(KEY_PRIVATE_IP)); + certRecord.setPrincipals(DynamoDBUtils.getString(item, KEY_PRINCIPALS)); + certRecord.setClientIP(DynamoDBUtils.getString(item, KEY_CLIENT_IP)); + certRecord.setPrivateIP(DynamoDBUtils.getString(item, KEY_PRIVATE_IP)); return certRecord; + } catch (Exception ex) { LOGGER.error("DynamoDB Get Error for {}: {}/{}", primaryKey, ex.getClass(), ex.getMessage()); return null; @@ -87,18 +101,25 @@ public boolean updateSSHCertRecord(SSHCertRecord certRecord) { final String primaryKey = getPrimaryKey(certRecord.getInstanceId(), certRecord.getService()); + HashMap itemKey = new HashMap<>(); + itemKey.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + try { - UpdateItemSpec updateItemSpec = new UpdateItemSpec() - .withPrimaryKey(KEY_PRIMARY, primaryKey) - .withAttributeUpdate( - new AttributeUpdate(KEY_INSTANCE_ID).put(certRecord.getInstanceId()), - new AttributeUpdate(KEY_SERVICE).put(certRecord.getService()), - new AttributeUpdate(KEY_CLIENT_IP).put(certRecord.getClientIP()), - new AttributeUpdate(KEY_PRINCIPALS).put(certRecord.getPrincipals()), - new AttributeUpdate(KEY_PRIVATE_IP).put(certRecord.getPrivateIP()), - new AttributeUpdate(KEY_TTL).put(System.currentTimeMillis() / 1000L + expiryTime) - ); - table.updateItem(updateItemSpec); + HashMap updatedValues = new HashMap<>(); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_INSTANCE_ID, certRecord.getInstanceId()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_SERVICE, certRecord.getService()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_CLIENT_IP, certRecord.getClientIP()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_PRINCIPALS, certRecord.getPrincipals()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_PRIVATE_IP, certRecord.getPrivateIP()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_TTL, Long.toString(System.currentTimeMillis() / 1000L + expiryTime)); + + UpdateItemRequest request = UpdateItemRequest.builder() + .tableName(tableName) + .key(itemKey) + .attributeUpdates(updatedValues) + .build(); + + dynamoDB.updateItem(request); return true; } catch (Exception ex) { LOGGER.error("DynamoDB Update Error for {}: {}/{}", primaryKey, ex.getClass(), ex.getMessage()); @@ -111,15 +132,21 @@ public boolean insertSSHCertRecord(SSHCertRecord certRecord) { final String primaryKey = getPrimaryKey(certRecord.getInstanceId(), certRecord.getService()); try { - Item item = new Item() - .withPrimaryKey(KEY_PRIMARY, primaryKey) - .withString(KEY_INSTANCE_ID, certRecord.getInstanceId()) - .withString(KEY_SERVICE, certRecord.getService()) - .withString(KEY_CLIENT_IP, certRecord.getClientIP()) - .withString(KEY_PRINCIPALS, certRecord.getPrincipals()) - .withString(KEY_PRIVATE_IP, certRecord.getPrivateIP()) - .withLong(KEY_TTL, System.currentTimeMillis() / 1000L + expiryTime); - table.putItem(item); + HashMap itemValues = new HashMap<>(); + itemValues.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + itemValues.put(KEY_INSTANCE_ID, AttributeValue.fromS(certRecord.getInstanceId())); + itemValues.put(KEY_SERVICE, AttributeValue.fromS(certRecord.getService())); + itemValues.put(KEY_CLIENT_IP, AttributeValue.fromS(certRecord.getClientIP())); + itemValues.put(KEY_PRINCIPALS, AttributeValue.fromS(certRecord.getPrincipals())); + itemValues.put(KEY_PRIVATE_IP, AttributeValue.fromS(certRecord.getPrivateIP())); + itemValues.put(KEY_TTL, AttributeValue.fromN(Long.toString(System.currentTimeMillis() / 1000L + expiryTime))); + + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(itemValues) + .build(); + + dynamoDB.putItem(request); return true; } catch (Exception ex) { LOGGER.error("DynamoDB Insert Error for {}: {}/{}", primaryKey, ex.getClass(), ex.getMessage()); @@ -131,10 +158,17 @@ public boolean insertSSHCertRecord(SSHCertRecord certRecord) { public boolean deleteSSHCertRecord(String instanceId, String service) { final String primaryKey = getPrimaryKey(instanceId, service); + + HashMap keyToGet = new HashMap<>(); + keyToGet.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + + DeleteItemRequest request = DeleteItemRequest.builder() + .tableName(tableName) + .key(keyToGet) + .build(); + try { - DeleteItemSpec deleteItemSpec = new DeleteItemSpec() - .withPrimaryKey(KEY_PRIMARY, primaryKey); - table.deleteItem(deleteItemSpec); + dynamoDB.deleteItem(request); return true; } catch (Exception ex) { LOGGER.error("DynamoDB Delete Error for {}: {}/{}", primaryKey, ex.getClass(), ex.getMessage()); @@ -146,7 +180,7 @@ public boolean deleteSSHCertRecord(String instanceId, String service) { public int deleteExpiredSSHCertRecords(int expiryTimeMins) { // with dynamo db there is no need to manually expunge expired - // record since we have the TTL option enabled for our table + // record since we have the TTL option enabled for our table, // and we just need to make sure the attribute is updated with // the epoch time + timeout seconds when it should retire diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactory.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactory.java index 74bb48395c6..bf3edafef34 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactory.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactory.java @@ -15,7 +15,6 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.ssh.SSHRecordStore; import com.yahoo.athenz.common.server.ssh.SSHRecordStoreFactory; @@ -26,6 +25,7 @@ import com.yahoo.athenz.zts.notification.ZTSClientNotificationSenderImpl; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; public class DynamoDBSSHRecordStoreFactory implements SSHRecordStoreFactory { private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBSSHRecordStoreFactory.class); @@ -40,13 +40,14 @@ public SSHRecordStore create(PrivateKeyStore keyStore) { } ZTSClientNotificationSenderImpl ztsClientNotificationSender = new ZTSClientNotificationSenderImpl(); - AmazonDynamoDB client = getDynamoDBClient(ztsClientNotificationSender, keyStore); + DynamoDbClient client = getDynamoDBClient(ztsClientNotificationSender, keyStore); return new DynamoDBSSHRecordStore(client, tableName, ztsClientNotificationSender); } - AmazonDynamoDB getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { + DynamoDbClient getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { DynamoDBClientFetcher dynamoDBClientFetcher = DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); ZTSDynamoDBClientSettingsFactory ztsDynamoDBClientSettingsFactory = new ZTSDynamoDBClientSettingsFactory(keyStore); - return dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings()).getAmazonDynamoDB(); + return dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings(false)).getDynamoDbClient(); } } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusChecker.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusChecker.java index e35e7346524..fe8d79f381a 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusChecker.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusChecker.java @@ -16,7 +16,6 @@ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.status.StatusCheckException; import com.yahoo.athenz.common.server.status.StatusChecker; @@ -24,8 +23,11 @@ import com.yahoo.athenz.db.dynamodb.DynamoDBClientFetcher; import com.yahoo.athenz.db.dynamodb.DynamoDBClientFetcherFactory; import org.apache.http.HttpStatus; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.ListTablesRequest; +import software.amazon.awssdk.services.dynamodb.model.ListTablesResponse; -import java.io.IOException; +import java.util.List; public class DynamoDBStatusChecker implements StatusChecker { @@ -44,12 +46,12 @@ public void check() throws StatusCheckException { // Get DynamoDB client and temp credentials (if required) DynamoDBClientFetcher dynamoDBClientFetcher = getDynamoDBClientFetcher(); ZTSDynamoDBClientSettingsFactory ztsDynamoDBClientSettingsFactory = new ZTSDynamoDBClientSettingsFactory(keyStore); - clientAndCreds = dynamoDBClientFetcher.getDynamoDBClient(null, ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings()); - AmazonDynamoDB amazonDynamoDB = clientAndCreds.getAmazonDynamoDB(); + clientAndCreds = dynamoDBClientFetcher.getDynamoDBClient(null, + ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings(false)); + DynamoDbClient dynamoDbClient = clientAndCreds.getDynamoDbClient(); // Get list of tables and verify our table appears - boolean tableFound = amazonDynamoDB.listTables().getTableNames().stream() - .anyMatch(fetchedTableName -> fetchedTableName.equals(tableName)); + boolean tableFound = dynamoDbTablePresent(dynamoDbClient); if (!tableFound) { throw new StatusCheckException(HttpStatus.SC_OK, "Table named " + tableName + " wasn't found in DynamoDB"); @@ -59,18 +61,8 @@ public void check() throws StatusCheckException { } catch (Throwable ex) { throw new StatusCheckException(ex); } finally { - // Close resources if (clientAndCreds != null) { - try { - if (clientAndCreds.getAmazonDynamoDB() != null) { - clientAndCreds.getAmazonDynamoDB().shutdown(); - } - if (clientAndCreds.getAwsCredentialsProvider() != null) { - clientAndCreds.getAwsCredentialsProvider().close(); - } - } catch (IOException ignored) { - - } + clientAndCreds.close(); } } } @@ -78,4 +70,29 @@ public void check() throws StatusCheckException { DynamoDBClientFetcher getDynamoDBClientFetcher() { return DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); } + + boolean dynamoDbTablePresent(DynamoDbClient ddb) { + boolean moreTables = true; + String lastEvaluatedTableName = null; + + while (moreTables) { + ListTablesRequest request = ListTablesRequest.builder() + .limit(100) + .exclusiveStartTableName(lastEvaluatedTableName) + .build(); + + ListTablesResponse response = ddb.listTables(request); + List tables = response.tableNames(); + + for (String table : tables) { + if (tableName.equals(table)) { + return true; + } + } + + lastEvaluatedTableName = response.lastEvaluatedTableName(); + moreTables = lastEvaluatedTableName != null; + } + return false; + } } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerFactory.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerFactory.java index e2a49685fec..8e1151b6cbd 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerFactory.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerFactory.java @@ -44,9 +44,9 @@ private PrivateKeyStore getKeyStore() { ZTSConsts.ZTS_PKEY_STORE_FACTORY_CLASS); PrivateKeyStoreFactory pkeyFactory; try { - pkeyFactory = (PrivateKeyStoreFactory) Class.forName(pkeyFactoryClass).newInstance(); - } catch (InstantiationException | IllegalAccessException | ClassNotFoundException e) { - LOGGER.error("Invalid PrivateKeyStoreFactory class: {} error: {}", pkeyFactoryClass, e.getMessage()); + pkeyFactory = (PrivateKeyStoreFactory) Class.forName(pkeyFactoryClass).getDeclaredConstructor().newInstance(); + } catch (Exception ex) { + LOGGER.error("Invalid PrivateKeyStoreFactory class: {} error: {}", pkeyFactoryClass, ex.getMessage()); throw new IllegalArgumentException("Invalid private key store"); } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/ZTSDynamoDBClientSettingsFactory.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/ZTSDynamoDBClientSettingsFactory.java index 746855749e8..a6bac2f348b 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/ZTSDynamoDBClientSettingsFactory.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/cert/impl/ZTSDynamoDBClientSettingsFactory.java @@ -37,7 +37,7 @@ public class ZTSDynamoDBClientSettingsFactory { private final String externalId; private final Integer minExpiryTime; private final Integer maxExpiryTime; - private final String keygroupName; + private final String keyGroupName; public ZTSDynamoDBClientSettingsFactory(PrivateKeyStore keyStore) { keyPath = System.getProperty(ZTS_PROP_DYNAMODB_KEY_PATH, ""); @@ -54,12 +54,14 @@ public ZTSDynamoDBClientSettingsFactory(PrivateKeyStore keyStore) { String maxExpiryTimeStr = System.getProperty(ZTS_PROP_DYNAMODB_MAX_EXPIRY_TIME, ""); minExpiryTime = minExpiryTimeStr.isEmpty() ? null : Integer.parseInt(minExpiryTimeStr); maxExpiryTime = maxExpiryTimeStr.isEmpty() ? null : Integer.parseInt(maxExpiryTimeStr); - keygroupName = System.getProperty(ZTS_PROP_DYNAMODB_TRUSTSTORE_KEYGROUPNAME, ""); + keyGroupName = System.getProperty(ZTS_PROP_DYNAMODB_TRUSTSTORE_KEYGROUPNAME, ""); this.keyStore = keyStore; } - public DynamoDBClientSettings getDynamoDBClientSettings() { - return new DynamoDBClientSettings(certPath, domainName, roleName, trustStore, trustStorePassword, ztsURL, region, keyPath, appName, keyStore, externalId, minExpiryTime, maxExpiryTime, keygroupName); + public DynamoDBClientSettings getDynamoDBClientSettings(boolean isAsyncClient) { + return new DynamoDBClientSettings(certPath, domainName, roleName, trustStore, trustStorePassword, + ztsURL, region, keyPath, appName, keyStore, externalId, minExpiryTime, maxExpiryTime, + keyGroupName, isAsyncClient); } } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/store/CloudStore.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/store/CloudStore.java index d3c9b430a97..45a0bd15bab 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/store/CloudStore.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/store/CloudStore.java @@ -25,7 +25,6 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.ConcurrentHashMap; -import com.amazonaws.AmazonServiceException; import com.yahoo.athenz.common.server.util.ConfigProperties; import org.eclipse.jetty.client.HttpClient; import org.eclipse.jetty.client.ContentResponse; @@ -35,22 +34,20 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import com.amazonaws.regions.Regions; -import com.amazonaws.auth.BasicSessionCredentials; -import com.amazonaws.auth.AWSStaticCredentialsProvider; -import com.amazonaws.services.s3.AmazonS3; -import com.amazonaws.services.s3.AmazonS3ClientBuilder; -import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder; -import com.amazonaws.services.securitytoken.AWSSecurityTokenService; -import com.amazonaws.services.securitytoken.model.AssumeRoleRequest; -import com.amazonaws.services.securitytoken.model.AssumeRoleResult; -import com.amazonaws.services.securitytoken.model.Credentials; import com.yahoo.athenz.zts.AWSTemporaryCredentials; import com.yahoo.athenz.zts.ResourceException; import com.yahoo.athenz.zts.ZTSConsts; import com.yahoo.rdl.JSON; import com.yahoo.rdl.Struct; import com.yahoo.rdl.Timestamp; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; +import software.amazon.awssdk.services.sts.model.Credentials; +import software.amazon.awssdk.awscore.exception.AwsServiceException; import static com.yahoo.athenz.common.ServerCommonConsts.ZTS_PROP_AWS_REGION_NAME; @@ -68,7 +65,7 @@ public class CloudStore { boolean awsEnabled; int cacheTimeout; int invalidCacheTimeout; - BasicSessionCredentials credentials; + AwsBasicCredentials credentials; final private Map awsAccountCache; final private Map azureSubscriptionCache; final private Map azureTenantCache; @@ -200,24 +197,6 @@ void initializeAwsSupport() { credsUpdateTime, TimeUnit.SECONDS); } - public AmazonS3 getS3Client() { - - if (!awsEnabled) { - throw new ResourceException(ResourceException.INTERNAL_SERVER_ERROR, - "AWS Support not enabled"); - } - - if (credentials == null) { - throw new ResourceException(ResourceException.INTERNAL_SERVER_ERROR, - "AWS Role credentials are not available"); - } - - return AmazonS3ClientBuilder.standard() - .withCredentials(new AWSStaticCredentialsProvider(credentials)) - .withRegion(Regions.fromName(awsRegion)) - .build(); - } - boolean loadBootMetaData() { // first load the dynamic document @@ -228,7 +207,6 @@ boolean loadBootMetaData() { } if (!parseInstanceInfo(document)) { - LOGGER.error("CloudStore: unable to parse instance identity document: {}", document); return false; } @@ -251,7 +229,6 @@ boolean loadBootMetaData() { // report the error and return false if (!parseIamRoleInfo(iamRole)) { - LOGGER.error("CloudStore: unable to parse iam role data: {}", iamRole); return false; } @@ -285,7 +262,7 @@ boolean parseInstanceInfo(String document) { return true; } - boolean parseIamRoleInfo(String iamRole) { + boolean parseIamRoleInfo(final String iamRole) { Struct iamRoleStruct = JSON.fromString(iamRole, Struct.class); if (iamRoleStruct == null) { @@ -296,7 +273,7 @@ boolean parseIamRoleInfo(String iamRole) { // extract and parse our profile arn // "InstanceProfileArn" : "arn:aws:iam::1111111111111:instance-profile/iaas.athenz.zts,athenz", - String profileArn = iamRoleStruct.getString("InstanceProfileArn"); + final String profileArn = iamRoleStruct.getString("InstanceProfileArn"); if (StringUtil.isEmpty(profileArn)) { LOGGER.error("CloudStore: unable to extract InstanceProfileArn from iam role data: {}", iamRole); return false; @@ -305,7 +282,7 @@ boolean parseIamRoleInfo(String iamRole) { return parseInstanceProfileArn(profileArn); } - boolean parseInstanceProfileArn(String profileArn) { + boolean parseInstanceProfileArn(final String profileArn) { // "InstanceProfileArn" : "arn:aws:iam::1111111111111:instance-profile/iaas.athenz.zts,athenz", @@ -365,9 +342,12 @@ boolean fetchRoleCredentials() { String accessKeyId = credsStruct.getString("AccessKeyId"); String secretAccessKey = credsStruct.getString("SecretAccessKey"); - String token = credsStruct.getString("Token"); - credentials = new BasicSessionCredentials(accessKeyId, secretAccessKey, token); + credentials = AwsBasicCredentials.builder() + .accessKeyId(accessKeyId) + .secretAccessKey(secretAccessKey) + .build(); + return true; } @@ -440,23 +420,23 @@ AssumeRoleRequest getAssumeRoleRequest(final String account, final String roleNa final String arn = "arn:aws:iam::" + account + ":role/" + roleName; - AssumeRoleRequest req = new AssumeRoleRequest(); - req.setRoleArn(arn); - req.setRoleSessionName(getAssumeRoleSessionName(principal)); + AssumeRoleRequest.Builder builder = AssumeRoleRequest.builder() + .roleArn(arn) + .roleSessionName(getAssumeRoleSessionName(principal)); if (durationSeconds != null && durationSeconds > 0) { - req.setDurationSeconds(durationSeconds); + builder = builder.durationSeconds(durationSeconds); } if (!StringUtil.isEmpty(externalId)) { - req.setExternalId(externalId); + builder = builder.externalId(externalId); } - return req; + return builder.build(); } - AWSSecurityTokenService getTokenServiceClient() { + StsClient getTokenServiceClient() { - return AWSSecurityTokenServiceClientBuilder.standard() - .withCredentials(new AWSStaticCredentialsProvider(credentials)) - .withRegion(Regions.fromName(awsRegion)) + return StsClient.builder() + .credentialsProvider(StaticCredentialsProvider.create(credentials)) + .region(Region.of(awsRegion)) .build(); } @@ -606,35 +586,36 @@ public AWSTemporaryCredentials assumeAWSRole(String account, String roleName, St AssumeRoleRequest req = getAssumeRoleRequest(account, roleName, durationSeconds, externalId, principal); try { - AWSSecurityTokenService client = getTokenServiceClient(); - AssumeRoleResult res = client.assumeRole(req); + StsClient client = getTokenServiceClient(); + AssumeRoleResponse res = client.assumeRole(req); - Credentials awsCreds = res.getCredentials(); + Credentials awsCreds = res.credentials(); tempCreds = new AWSTemporaryCredentials() - .setAccessKeyId(awsCreds.getAccessKeyId()) - .setSecretAccessKey(awsCreds.getSecretAccessKey()) - .setSessionToken(awsCreds.getSessionToken()) - .setExpiration(Timestamp.fromMillis(awsCreds.getExpiration().getTime())); + .setAccessKeyId(awsCreds.accessKeyId()) + .setSecretAccessKey(awsCreds.secretAccessKey()) + .setSessionToken(awsCreds.sessionToken()) + .setExpiration(Timestamp.fromMillis(awsCreds.expiration().toEpochMilli())); - } catch (AmazonServiceException ex) { + } catch (AwsServiceException ex) { + int statusCode = ex.awsErrorDetails().sdkHttpResponse().statusCode(); LOGGER.error("CloudStore: assumeAWSRole - unable to assume role: {}, error: {}, status code: {}", - req.getRoleArn(), ex.getMessage(), ex.getStatusCode()); + req.roleArn(), ex.getMessage(), statusCode); // if this is access denied then we're going to cache // the failed results - if (ex.getStatusCode() == ResourceException.FORBIDDEN) { + if (statusCode == ResourceException.FORBIDDEN) { putInvalidCacheCreds(cacheKey); } - errorMessage.append(ex.getErrorMessage()); + errorMessage.append(ex.awsErrorDetails().errorMessage()); return null; } catch (Exception ex) { LOGGER.error("CloudStore: assumeAWSRole - unable to assume role: {}, error: {}", - req.getRoleArn(), ex.getMessage()); + req.roleArn(), ex.getMessage()); errorMessage.append(ex.getMessage()); return null; diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/DynamoDBUtils.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/DynamoDBUtils.java index 6dae891dfd1..23d95bcc12b 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/DynamoDBUtils.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/DynamoDBUtils.java @@ -16,12 +16,12 @@ package com.yahoo.athenz.zts.utils; -import com.amazonaws.services.dynamodbv2.document.Item; +import software.amazon.awssdk.services.dynamodb.model.AttributeAction; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.AttributeValueUpdate; import java.text.SimpleDateFormat; -import java.util.ArrayList; -import java.util.Date; -import java.util.List; +import java.util.*; import java.util.concurrent.TimeUnit; public class DynamoDBUtils { @@ -43,18 +43,81 @@ public static List getISODatesByRange(long epochBegin, long epochEnd) { return datesArray; } - public static Date getDateFromItem(Item item, String key) { - if (item.isNull(key) || item.get(key) == null) { + public static Date getDateFromItem(Map item, final String key) { + AttributeValue value = item.get(key); + if (value == null) { return null; } - return new Date(item.getLong(key)); + return new Date(Long.parseLong(value.n())); } - public static Object getLongFromDate(Date date) { + public static String getNumberFromDate(Date date) { if (date == null) { + return "0"; + } + + return String.valueOf(date.getTime()); + } + + public static String getString(Map item, final String key) { + AttributeValue value = item.get(key); + if (value == null) { return null; } + return value.s(); + } + + public static boolean getBoolean(Map item, final String key) { + AttributeValue value = item.get(key); + if (value == null || value.bool() == null) { + return false; + } + return value.bool(); + } + + public static long getLong(Map item, final String key) { + AttributeValue value = item.get(key); + if (value == null) { + return 0; + } + return Long.parseLong(value.n()); + } - return date.getTime(); + public static void updateItemStringValue(HashMap updatedValues, + final String key, final String value) { + updatedValues.put(key, AttributeValueUpdate.builder() + .value(AttributeValue.fromS(value)) + .action(AttributeAction.PUT) + .build()); + } + + public static void updateItemBoolValue(HashMap updatedValues, + final String key, final Boolean value) { + updatedValues.put(key, AttributeValueUpdate.builder() + .value(AttributeValue.fromBool(value)) + .action(AttributeAction.PUT) + .build()); + } + + public static void updateItemLongValue(HashMap updatedValues, + final String key, final Long value) { + updatedValues.put(key, AttributeValueUpdate.builder() + .value(AttributeValue.fromN(String.valueOf(value))) + .action(AttributeAction.PUT) + .build()); + } + + public static void updateItemLongValue(HashMap updatedValues, + final String key, final Date value) { + AttributeValue attributeValue; + if (value == null) { + attributeValue = AttributeValue.builder().nul(true).build(); + } else { + attributeValue = AttributeValue.fromN(String.valueOf(value.getTime())); + } + updatedValues.put(key, AttributeValueUpdate.builder() + .value(attributeValue) + .action(AttributeAction.PUT) + .build()); } } diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/RetryDynamoDBCommand.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/RetryDynamoDBCommand.java index 34600e47897..2b8f9531607 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/RetryDynamoDBCommand.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/utils/RetryDynamoDBCommand.java @@ -1,9 +1,9 @@ package com.yahoo.athenz.zts.utils; -import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughputExceededException; import com.yahoo.athenz.zts.ZTSConsts; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughputExceededException; import java.util.concurrent.TimeoutException; import java.util.function.Supplier; diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStore.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStore.java index d209e232c4a..ee6d37e66e7 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStore.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStore.java @@ -15,25 +15,20 @@ */ package com.yahoo.athenz.zts.workload.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; import com.yahoo.athenz.common.server.workload.WorkloadRecordStore; import com.yahoo.athenz.common.server.workload.WorkloadRecordStoreConnection; -import com.yahoo.athenz.zts.ResourceException; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; public class DynamoDBWorkloadRecordStore implements WorkloadRecordStore { - private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBWorkloadRecordStore.class); - private final String tableName; private final String serviceIndexName; private final String ipIndexName; - private final DynamoDB dynamoDB; + private final DynamoDbClient dynamoDB; - public DynamoDBWorkloadRecordStore(AmazonDynamoDB client, String tableName, String serviceIndexName, String ipIndexName) { - this.dynamoDB = new DynamoDB(client); + public DynamoDBWorkloadRecordStore(DynamoDbClient client, final String tableName, final String serviceIndexName, + final String ipIndexName) { + this.dynamoDB = client; this.tableName = tableName; this.serviceIndexName = serviceIndexName; this.ipIndexName = ipIndexName; @@ -41,12 +36,7 @@ public DynamoDBWorkloadRecordStore(AmazonDynamoDB client, String tableName, Stri @Override public WorkloadRecordStoreConnection getConnection() { - try { - return new DynamoDBWorkloadRecordStoreConnection(dynamoDB, tableName, serviceIndexName, ipIndexName); - } catch (Exception ex) { - LOGGER.error("getConnection: {}", ex.getMessage()); - throw new ResourceException(ResourceException.SERVICE_UNAVAILABLE, ex.getMessage()); - } + return new DynamoDBWorkloadRecordStoreConnection(dynamoDB, tableName, serviceIndexName, ipIndexName); } @Override diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnection.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnection.java index b5210a5b196..1ce8ff392a7 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnection.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnection.java @@ -15,10 +15,6 @@ */ package com.yahoo.athenz.zts.workload.impl; -import com.amazonaws.services.dynamodbv2.document.*; -import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; -import com.amazonaws.services.dynamodbv2.document.utils.ValueMap; import com.yahoo.athenz.auth.util.AthenzUtils; import com.yahoo.athenz.common.server.workload.WorkloadRecord; import com.yahoo.athenz.common.server.workload.WorkloadRecordStoreConnection; @@ -27,17 +23,18 @@ import com.yahoo.athenz.zts.utils.RetryDynamoDBCommand; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.*; -import java.util.ArrayList; -import java.util.Date; -import java.util.List; +import java.util.*; +import java.util.concurrent.TimeoutException; public class DynamoDBWorkloadRecordStoreConnection implements WorkloadRecordStoreConnection { private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBWorkloadRecordStoreConnection.class); - private final RetryDynamoDBCommand> itemCollectionRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); - private final RetryDynamoDBCommand updateItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); - private final RetryDynamoDBCommand putItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + private final RetryDynamoDBCommand itemCollectionRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + private final RetryDynamoDBCommand updateItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); + private final RetryDynamoDBCommand putItemRetryDynamoDBCommand = new RetryDynamoDBCommand<>(); private static final String KEY_PRIMARY = "primaryKey"; private static final String KEY_SERVICE = "service"; @@ -51,20 +48,23 @@ public class DynamoDBWorkloadRecordStoreConnection implements WorkloadRecordStor private static final String KEY_EXPIRY_TIME = "certExpiryTime"; private static final String DEFAULT_HOSTNAME_IF_NULL = "NA"; - // the configuration setting is in hours so we'll automatically + // the configuration setting is in hours, so we'll automatically // convert into seconds since that's what dynamoDB needs // we need to expire records in 30 days private static long expiryTime = 3660 * Long.parseLong( System.getProperty(ZTSConsts.ZTS_PROP_WORKLOAD_DYNAMODB_ITEM_TTL_HOURS, "720")); - private final Table table; - private final Index serviceIndex; - private final Index ipIndex; - - public DynamoDBWorkloadRecordStoreConnection(DynamoDB dynamoDB, final String tableName, final String serviceIndex, final String ipIndex) { - this.table = dynamoDB.getTable(tableName); - this.serviceIndex = table.getIndex(serviceIndex); - this.ipIndex = table.getIndex(ipIndex); + private final String tableName; + private final String serviceIndexName; + private final String ipIndexName; + private final DynamoDbClient dynamoDB; + + public DynamoDBWorkloadRecordStoreConnection(DynamoDbClient dynamoDB, final String tableName, + final String serviceIndexName, final String ipIndexName) { + this.dynamoDB = dynamoDB; + this.tableName = tableName; + this.serviceIndexName = serviceIndexName; + this.ipIndexName = ipIndexName; } @Override @@ -82,30 +82,45 @@ private String getPrimaryKey(final String service, final String instanceId, fina } @Override - public List getWorkloadRecordsByService(String domain, String service) { + public List getWorkloadRecordsByService(final String domain, final String service) { + try { - QuerySpec spec = new QuerySpec() - .withKeyConditionExpression("service = :v_service") - .withValueMap(new ValueMap() - .withString(":v_service", AthenzUtils.getPrincipalName(domain, service))); + HashMap attrValues = new HashMap<>(); + attrValues.put(":v_service", AttributeValue.fromS(AthenzUtils.getPrincipalName(domain, service))); + + QueryRequest request = QueryRequest.builder() + .tableName(tableName) + .indexName(serviceIndexName) + .keyConditionExpression("service = :v_service") + .expressionAttributeValues(attrValues) + .build(); + + return processWorkloadQuery(request); - return getWorkloadRecords(spec, serviceIndex); } catch (Exception ex) { - LOGGER.error("DynamoDB getWorkloadRecordsByService failed for service={}, error={}", AthenzUtils.getPrincipalName(domain, service), ex.getMessage()); + LOGGER.error("DynamoDB getWorkloadRecordsByService failed for service={}, error={}", + AthenzUtils.getPrincipalName(domain, service), ex.getMessage()); } return new ArrayList<>(); } @Override - public List getWorkloadRecordsByIp(String ip) { + public List getWorkloadRecordsByIp(final String ip) { + try { - QuerySpec spec = new QuerySpec() - .withKeyConditionExpression("ip = :v_ip") - .withValueMap(new ValueMap() - .withString(":v_ip", ip)); + HashMap attrValues = new HashMap<>(); + attrValues.put(":v_ip", AttributeValue.fromS(ip)); + + QueryRequest request = QueryRequest.builder() + .tableName(tableName) + .indexName(ipIndexName) + .keyConditionExpression("ip = :v_ip") + .expressionAttributeValues(attrValues) + .build(); + + return processWorkloadQuery(request); - return getWorkloadRecords(spec, ipIndex); } catch (Exception ex) { LOGGER.error("DynamoDB getWorkloadRecordsByIp failed for ip={}, error={}", ip, ex.getMessage()); } @@ -113,32 +128,36 @@ public List getWorkloadRecordsByIp(String ip) { return new ArrayList<>(); } - private List getWorkloadRecords(QuerySpec spec, Index tableIndex) throws java.util.concurrent.TimeoutException, InterruptedException { - ItemCollection outcome = itemCollectionRetryDynamoDBCommand.run(() -> tableIndex.query(spec)); + private List processWorkloadQuery(QueryRequest request) throws InterruptedException, TimeoutException { + + QueryResponse response = itemCollectionRetryDynamoDBCommand.run(() -> dynamoDB.query(request)); + List workloadRecords = new ArrayList<>(); - for (Item item : outcome) { - workloadRecords.add(itemToWorkloadRecord(item)); + if (response.hasItems()) { + for (Map item : response.items()) { + workloadRecords.add(itemToWorkloadRecord(item)); + } } + return workloadRecords; } - - private WorkloadRecord itemToWorkloadRecord(Item item) { + private WorkloadRecord itemToWorkloadRecord(Map item) { WorkloadRecord workloadRecord = new WorkloadRecord(); - workloadRecord.setInstanceId(item.getString(KEY_INSTANCE_ID)); - workloadRecord.setService(item.getString(KEY_SERVICE)); - workloadRecord.setIp(item.getString(KEY_IP)); + workloadRecord.setInstanceId(DynamoDBUtils.getString(item, KEY_INSTANCE_ID)); + workloadRecord.setService(DynamoDBUtils.getString(item, KEY_SERVICE)); + workloadRecord.setIp(DynamoDBUtils.getString(item, KEY_IP)); - if (item.hasAttribute(KEY_HOSTNAME)) { - workloadRecord.setHostname(item.getString(KEY_HOSTNAME)); + if (item.containsKey(KEY_HOSTNAME)) { + workloadRecord.setHostname(DynamoDBUtils.getString(item, KEY_HOSTNAME)); } else { workloadRecord.setHostname(DEFAULT_HOSTNAME_IF_NULL); } - workloadRecord.setProvider(item.getString(KEY_PROVIDER)); + workloadRecord.setProvider(DynamoDBUtils.getString(item, KEY_PROVIDER)); workloadRecord.setCreationTime(DynamoDBUtils.getDateFromItem(item, KEY_CREATION_TIME)); workloadRecord.setUpdateTime(DynamoDBUtils.getDateFromItem(item, KEY_UPDATE_TIME)); - if (item.hasAttribute(KEY_EXPIRY_TIME)) { + if (item.containsKey(KEY_EXPIRY_TIME)) { workloadRecord.setCertExpiryTime(DynamoDBUtils.getDateFromItem(item, KEY_EXPIRY_TIME)); } else { workloadRecord.setCertExpiryTime(new Date(0)); //setting default date to 01/01/1970. @@ -150,23 +169,33 @@ private WorkloadRecord itemToWorkloadRecord(Item item) { @Override public boolean updateWorkloadRecord(WorkloadRecord workloadRecord) { - //updateItem does not fail on absence of primaryKey, and behaves as insert. So we should set all attributes with update too. + // updateItem does not fail on absence of primaryKey, and behaves as insert. + // So we should set all attributes with update too. + + HashMap itemKey = new HashMap<>(); + itemKey.put(KEY_PRIMARY, AttributeValue.fromS(getPrimaryKey(workloadRecord.getService(), workloadRecord.getInstanceId(), workloadRecord.getIp()))); + try { - UpdateItemSpec updateItemSpec = new UpdateItemSpec() - .withPrimaryKey(KEY_PRIMARY, getPrimaryKey(workloadRecord.getService(), workloadRecord.getInstanceId(), workloadRecord.getIp())) - .withAttributeUpdate( - new AttributeUpdate(KEY_SERVICE).put(workloadRecord.getService()), - new AttributeUpdate(KEY_PROVIDER).put(workloadRecord.getProvider()), - new AttributeUpdate(KEY_IP).put(workloadRecord.getIp()), - new AttributeUpdate(KEY_INSTANCE_ID).put(workloadRecord.getInstanceId()), - new AttributeUpdate(KEY_CREATION_TIME).put(DynamoDBUtils.getLongFromDate(workloadRecord.getCreationTime())), - new AttributeUpdate(KEY_UPDATE_TIME).put(DynamoDBUtils.getLongFromDate(workloadRecord.getUpdateTime())), - new AttributeUpdate(KEY_EXPIRY_TIME).put(DynamoDBUtils.getLongFromDate(workloadRecord.getCertExpiryTime())), - new AttributeUpdate(KEY_HOSTNAME).put(workloadRecord.getHostname()), - new AttributeUpdate(KEY_TTL).put(workloadRecord.getUpdateTime().getTime() / 1000L + expiryTime) - ); - updateItemRetryDynamoDBCommand.run(() -> table.updateItem(updateItemSpec)); + HashMap updatedValues = new HashMap<>(); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_INSTANCE_ID, workloadRecord.getInstanceId()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_PROVIDER, workloadRecord.getProvider()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_SERVICE, workloadRecord.getService()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_CREATION_TIME, workloadRecord.getCreationTime()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_UPDATE_TIME, workloadRecord.getUpdateTime()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_EXPIRY_TIME, workloadRecord.getCertExpiryTime()); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_IP, workloadRecord.getIp()); + DynamoDBUtils.updateItemLongValue(updatedValues, KEY_TTL, workloadRecord.getUpdateTime().getTime() / 1000L + expiryTime); + DynamoDBUtils.updateItemStringValue(updatedValues, KEY_HOSTNAME, workloadRecord.getHostname()); + + UpdateItemRequest request = UpdateItemRequest.builder() + .tableName(tableName) + .key(itemKey) + .attributeUpdates(updatedValues) + .build(); + + updateItemRetryDynamoDBCommand.run(() -> dynamoDB.updateItem(request)); return true; + } catch (Exception ex) { LOGGER.error("DynamoDB Workload update Error={}: {}/{}", workloadRecord, ex.getClass(), ex.getMessage()); return false; @@ -175,20 +204,31 @@ public boolean updateWorkloadRecord(WorkloadRecord workloadRecord) { @Override public boolean insertWorkloadRecord(WorkloadRecord workloadRecord) { + + final String primaryKey = getPrimaryKey(workloadRecord.getService(), workloadRecord.getInstanceId(), + workloadRecord.getIp()); + try { - Item item = new Item() - .withPrimaryKey(KEY_PRIMARY, getPrimaryKey(workloadRecord.getService(), workloadRecord.getInstanceId(), workloadRecord.getIp())) - .withString(KEY_SERVICE, workloadRecord.getService()) - .withString(KEY_PROVIDER, workloadRecord.getProvider()) - .withString(KEY_IP, workloadRecord.getIp()) - .withString(KEY_INSTANCE_ID, workloadRecord.getInstanceId()) - .withString(KEY_HOSTNAME, workloadRecord.getHostname()) - .with(KEY_CREATION_TIME, DynamoDBUtils.getLongFromDate(workloadRecord.getCreationTime())) - .with(KEY_UPDATE_TIME, DynamoDBUtils.getLongFromDate(workloadRecord.getUpdateTime())) - .with(KEY_EXPIRY_TIME, DynamoDBUtils.getLongFromDate(workloadRecord.getCertExpiryTime())) - .withLong(KEY_TTL, workloadRecord.getUpdateTime().getTime() / 1000L + expiryTime); - putItemRetryDynamoDBCommand.run(() -> table.putItem(item)); + HashMap itemValues = new HashMap<>(); + itemValues.put(KEY_PRIMARY, AttributeValue.fromS(primaryKey)); + itemValues.put(KEY_INSTANCE_ID, AttributeValue.fromS(workloadRecord.getInstanceId())); + itemValues.put(KEY_PROVIDER, AttributeValue.fromS(workloadRecord.getProvider())); + itemValues.put(KEY_SERVICE, AttributeValue.fromS(workloadRecord.getService())); + itemValues.put(KEY_IP, AttributeValue.fromS(workloadRecord.getIp())); + itemValues.put(KEY_HOSTNAME, AttributeValue.fromS(workloadRecord.getHostname())); + itemValues.put(KEY_CREATION_TIME, AttributeValue.fromN(DynamoDBUtils.getNumberFromDate(workloadRecord.getCreationTime()))); + itemValues.put(KEY_UPDATE_TIME, AttributeValue.fromN(DynamoDBUtils.getNumberFromDate(workloadRecord.getUpdateTime()))); + itemValues.put(KEY_EXPIRY_TIME, AttributeValue.fromN(DynamoDBUtils.getNumberFromDate(workloadRecord.getCertExpiryTime()))); + itemValues.put(KEY_TTL, AttributeValue.fromN(Long.toString(workloadRecord.getUpdateTime().getTime() / 1000L + expiryTime))); + + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(itemValues) + .build(); + + putItemRetryDynamoDBCommand.run(() -> dynamoDB.putItem(request)); return true; + } catch (Exception ex) { LOGGER.error("DynamoDB Workload Insert Error={}: {}/{}", workloadRecord, ex.getClass(), ex.getMessage()); return false; diff --git a/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactory.java b/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactory.java index 16cc2ac5bf1..85e2bb4fa6d 100644 --- a/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactory.java +++ b/servers/zts/src/main/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactory.java @@ -15,7 +15,6 @@ */ package com.yahoo.athenz.zts.workload.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.workload.WorkloadRecordStore; import com.yahoo.athenz.common.server.workload.WorkloadRecordStoreFactory; @@ -28,6 +27,7 @@ import org.eclipse.jetty.util.StringUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; public class DynamoDBWorkloadRecordStoreFactory implements WorkloadRecordStoreFactory { private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBWorkloadRecordStoreFactory.class); @@ -53,14 +53,15 @@ public WorkloadRecordStore create(PrivateKeyStore keyStore) { throw new ResourceException(ResourceException.SERVICE_UNAVAILABLE, "DynamoDB index workload-ip-index not specified"); } - AmazonDynamoDB client = getDynamoDBClient(null, keyStore); + DynamoDbClient client = getDynamoDBClient(null, keyStore); return new DynamoDBWorkloadRecordStore(client, tableName, serviceIndexName, ipIndexName); } - AmazonDynamoDB getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { + DynamoDbClient getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { DynamoDBClientFetcher dynamoDBClientFetcher = DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); ZTSDynamoDBClientSettingsFactory ztsDynamoDBClientSettingsFactory = new ZTSDynamoDBClientSettingsFactory(keyStore); - return dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings()).getAmazonDynamoDB(); + return dynamoDBClientFetcher.getDynamoDBClient(ztsClientNotificationSender, + ztsDynamoDBClientSettingsFactory.getDynamoDBClientSettings(false)).getDynamoDbClient(); } } diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/ZTSTestUtils.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/ZTSTestUtils.java index 72c82edd511..f61e54533c4 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/ZTSTestUtils.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/ZTSTestUtils.java @@ -15,7 +15,6 @@ */ package com.yahoo.athenz.zts; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; import com.yahoo.athenz.auth.util.Crypto; import com.yahoo.athenz.common.server.util.ResourceUtils; import com.yahoo.athenz.common.server.workload.WorkloadRecord; @@ -30,6 +29,7 @@ import org.eclipse.jetty.util.StringUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import java.io.File; import java.security.PrivateKey; @@ -272,7 +272,7 @@ public static void setupDomainsWithGroups(DataStore store, PrivateKey privateKey "update", domainName1 + ":resource4", com.yahoo.athenz.zms.AssertionEffect.ALLOW); policies.add(policy4); - // setup our signed domains and process them + // set up our signed domains and process them List groups = new ArrayList<>(); if (group1 != null) { @@ -395,98 +395,74 @@ public static Timestamp addDays(Timestamp date, int days) { return Timestamp.fromMillis(date.millis() + TimeUnit.MILLISECONDS.convert(days, TimeUnit.DAYS)); } - public static Map generateAttributeValues(String service, - String instanceId, - String currentTime, - String lastNotifiedTime, - String lastNotifiedServer, - String expiryTime, - String hostName) { + public static Map generateAttributeValues(final String service, + final String instanceId, final String currentTime, final String lastNotifiedTime, + final String lastNotifiedServer, final String expiryTime, final String hostName) { + String provider = "provider"; String primaryKey = provider + ":" + service + ":" + instanceId; Map item = new HashMap<>(); - item.put("primaryKey", new AttributeValue(primaryKey)); - item.put("service", new AttributeValue(service)); - item.put("provider", new AttributeValue(provider)); - item.put("instanceId", new AttributeValue(instanceId)); - item.put("currentSerial", new AttributeValue("currentSerial")); + item.put("primaryKey", AttributeValue.fromS(primaryKey)); + item.put("service", AttributeValue.fromS(service)); + item.put("provider", AttributeValue.fromS(provider)); + item.put("instanceId", AttributeValue.fromS(instanceId)); + item.put("currentSerial", AttributeValue.fromS("currentSerial")); - AttributeValue currentTimeVal = new AttributeValue(); - currentTimeVal.setN(currentTime); + AttributeValue currentTimeVal = AttributeValue.fromN(currentTime); if (!StringUtil.isEmpty(currentTime)) { item.put("currentTime", currentTimeVal); item.put("prevTime", currentTimeVal); } - item.put("currentIP", new AttributeValue("currentIP")); - item.put("prevSerial", new AttributeValue("prevSerial")); - item.put("prevIP", new AttributeValue("prevIP")); + item.put("currentIP", AttributeValue.fromS("currentIP")); + item.put("prevSerial", AttributeValue.fromS("prevSerial")); + item.put("prevIP", AttributeValue.fromS("prevIP")); - AttributeValue clientCertVal = new AttributeValue(); - clientCertVal.setBOOL(false); - item.put("clientCert", clientCertVal); + item.put("clientCert", AttributeValue.fromBool(false)); if (!StringUtil.isEmpty(lastNotifiedTime)) { - AttributeValue lastNotifiedTimeVal = new AttributeValue(); - lastNotifiedTimeVal.setN(lastNotifiedTime); - item.put("lastNotifiedTime", lastNotifiedTimeVal); + item.put("lastNotifiedTime", AttributeValue.fromN(lastNotifiedTime)); } if (!StringUtil.isEmpty(lastNotifiedServer)) { - item.put("lastNotifiedServer", new AttributeValue(lastNotifiedServer)); + item.put("lastNotifiedServer", AttributeValue.fromS(lastNotifiedServer)); } if (!StringUtil.isEmpty(expiryTime)) { - AttributeValue expiryTimeVal = new AttributeValue(); - expiryTimeVal.setN(expiryTime); - item.put("expiryTime", expiryTimeVal); + item.put("expiryTime", AttributeValue.fromN(expiryTime)); } if (!StringUtil.isEmpty(hostName)) { - item.put("hostName", new AttributeValue(hostName)); + item.put("hostName", AttributeValue.fromS(hostName)); } return item; } - public static Map generateWorkloadAttributeValues(String service, - String instanceId, - String provider, - String ip, - String hostname, - String creationTime, - String updateTime, - String certExpiryTime) { + public static Map generateWorkloadAttributeValues(final String service, + final String instanceId, final String provider, final String ip, final String hostname, + final String creationTime, final String updateTime, final String certExpiryTime) { + String primaryKey = service + "#" + instanceId + "#" + ip; Map item = new HashMap<>(); - item.put("primaryKey", new AttributeValue(primaryKey)); - item.put("service", new AttributeValue(service)); - item.put("provider", new AttributeValue(provider)); - item.put("instanceId", new AttributeValue(instanceId)); - item.put("ip", new AttributeValue(ip)); - item.put("hostname", new AttributeValue(hostname)); - AttributeValue creationTimeVal = new AttributeValue(); - creationTimeVal.setN(creationTime); - AttributeValue updateTimeVal = new AttributeValue(); - updateTimeVal.setN(updateTime); - AttributeValue certExpiryTimeVal = new AttributeValue(); - certExpiryTimeVal.setN(certExpiryTime); - item.put("creationTime", creationTimeVal); - item.put("updateTime", updateTimeVal); - item.put("certExpiryTime", certExpiryTimeVal); + item.put("primaryKey", AttributeValue.fromS(primaryKey)); + item.put("service", AttributeValue.fromS(service)); + item.put("provider", AttributeValue.fromS(provider)); + item.put("instanceId", AttributeValue.fromS(instanceId)); + item.put("ip", AttributeValue.fromS(ip)); + item.put("hostname", AttributeValue.fromS(hostname)); + item.put("creationTime", AttributeValue.fromN(creationTime)); + item.put("updateTime", AttributeValue.fromN(updateTime)); + item.put("certExpiryTime", AttributeValue.fromN(certExpiryTime)); return item; } - public static WorkloadRecord createWorkloadRecord(Date creationTime, - Date updateTime, - String provider, - String instanceId, - String hostname, - String ip, - String service, - Date certExpiryTime) { + public static WorkloadRecord createWorkloadRecord(Date creationTime, Date updateTime, + final String provider, final String instanceId, final String hostname, final String ip, + final String service, Date certExpiryTime) { + WorkloadRecord workloadRecord = new WorkloadRecord(); workloadRecord.setCreationTime(creationTime); workloadRecord.setUpdateTime(updateTime); @@ -500,7 +476,7 @@ public static WorkloadRecord createWorkloadRecord(Date creationTime, } public static String getAssumeRoleResource(final String domainName, final String roleName, - boolean wildCardRole, boolean wildCardDomain) { + boolean wildCardRole, boolean wildCardDomain) { if (wildCardRole && wildCardDomain) { return "*:role.*"; } else if (wildCardDomain) { diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactoryTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactoryTest.java index d41015df967..6f3dc84aa97 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactoryTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/AWSCertRecordStoreFactoryTest.java @@ -15,11 +15,7 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.auth.InstanceProfileCredentialsProvider; -import com.amazonaws.services.rds.auth.RdsIamAuthTokenGenerator; import com.yahoo.athenz.common.server.cert.CertRecordStore; -import org.mockito.ArgumentMatchers; -import org.mockito.Mockito; import org.testng.annotations.Test; import com.yahoo.athenz.zts.ZTSConsts; @@ -28,24 +24,15 @@ public class AWSCertRecordStoreFactoryTest { - class TestAWSCertRecordStoreFactory extends AWSCertRecordStoreFactory { - - RdsIamAuthTokenGenerator generator = Mockito.mock(RdsIamAuthTokenGenerator.class); - - @Override - RdsIamAuthTokenGenerator getTokenGenerator(InstanceProfileCredentialsProvider awsCredProvider) { - - Mockito.when(generator.getAuthToken(ArgumentMatchers.any())).thenReturn("token"); - return generator; - } + static class TestAWSCertRecordStoreFactory extends AWSCertRecordStoreFactory { @Override - String getInstanceRegion() { - return "us-west-2"; + String getAuthToken(String hostname, int port, String rdsUser) { + return "token"; } } - class TestAWSCertRecordStoreFactory2 extends AWSCertRecordStoreFactory { + static class TestAWSCertRecordStoreFactory2 extends AWSCertRecordStoreFactory { boolean throwGetTokenExc = false; @@ -54,7 +41,7 @@ void setThrowGetTokenExc(boolean value) { } @Override - String getAuthToken(String hostname, int port, String rdsUser, String rdsIamRole) { + String getAuthToken(String hostname, int port, String rdsUser) { if (throwGetTokenExc) { throw new IllegalArgumentException("Unable to get token"); } else { @@ -66,7 +53,7 @@ String getAuthToken(String hostname, int port, String rdsUser, String rdsIamRole @Test public void testCreate() { - System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_INSTANCE, "instance"); + System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE, "instance"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER, "rds-user"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_IAM_ROLE, "role"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME, "1"); @@ -81,7 +68,7 @@ public void testCreate() { } assertNotNull(store); - System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_INSTANCE); + System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_IAM_ROLE); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME); @@ -90,16 +77,16 @@ public void testCreate() { @Test public void testGetTokenException() { - System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_INSTANCE, "instance"); + System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE, "instance"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER, "rds-user"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_IAM_ROLE, "role"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME, "1"); - AWSCertRecordStoreFactory factory = new TestAWSCertRecordStoreFactory2(); + TestAWSCertRecordStoreFactory2 factory = new TestAWSCertRecordStoreFactory2(); CertRecordStore store = factory.create(null); assertNotNull(store); - ((TestAWSCertRecordStoreFactory2) factory).setThrowGetTokenExc(true); + factory.setThrowGetTokenExc(true); // we should not get any exceptions even though the get token // call will throw exceptions @@ -108,7 +95,7 @@ public void testGetTokenException() { updater.run(); updater.run(); - System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_INSTANCE); + System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_IAM_ROLE); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME); @@ -117,7 +104,7 @@ public void testGetTokenException() { @Test public void testCredentialsUpdater() { - System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_INSTANCE, "instance"); + System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE, "instance"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER, "rds-user2"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_IAM_ROLE, "role"); System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME, "1"); @@ -132,9 +119,28 @@ public void testCredentialsUpdater() { } catch (InterruptedException ignored) { } - System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_MASTER_INSTANCE); + System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_IAM_ROLE); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME); } + + @Test + public void testOriginalMethods() { + + System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE, "instance"); + System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER, "rds-user"); + System.setProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME, "30000"); + + AWSCertRecordStoreFactory factory = new AWSCertRecordStoreFactory(); + + try { + factory.getAuthToken("host", 3306, "user"); + } catch (Exception ignored) { + } + + System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_PRIMARY_INSTANCE); + System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_USER); + System.clearProperty(ZTSConsts.ZTS_PROP_AWS_RDS_CREDS_REFRESH_TIME); + } } diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnectionTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnectionTest.java index 018302eed79..913d613c86f 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnectionTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreConnectionTest.java @@ -15,27 +15,20 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.document.*; -import com.amazonaws.services.dynamodbv2.document.internal.IteratorSupport; -import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec; -import com.amazonaws.services.dynamodbv2.model.AmazonDynamoDBException; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; -import com.amazonaws.services.dynamodbv2.model.TransactionConflictException; import com.yahoo.athenz.common.server.cert.X509CertRecord; -import com.yahoo.athenz.zts.ZTSConsts; import com.yahoo.athenz.zts.ZTSTestUtils; import com.yahoo.athenz.zts.utils.DynamoDBUtils; import org.mockito.*; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; -import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; +import software.amazon.awssdk.awscore.exception.AwsServiceException; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.*; import java.util.*; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.when; import static org.mockito.internal.verification.VerificationModeFactory.times; import static org.testng.Assert.*; @@ -47,23 +40,15 @@ public class DynamoDBCertRecordStoreConnectionTest { private final String currentTimeIndexName = "cert-table-currenttime-index"; private final String hostNameIndexName = "cert-table-hostname-index"; + @Mock private DynamoDbClient dynamoDB = Mockito.mock(DynamoDbClient.class); - @Mock private DynamoDB dynamoDB = Mockito.mock(DynamoDB.class); - @Mock private Table table = Mockito.mock(Table.class); - @Mock private Index currentTimeIndex = Mockito.mock(Index.class); - @Mock private Index hostNameIndex = Mockito.mock(Index.class); - @Mock private Item item = Mockito.mock(Item.class); - @Mock private PutItemOutcome putOutcome = Mockito.mock(PutItemOutcome.class); - @Mock private DeleteItemOutcome deleteOutcome = Mockito.mock(DeleteItemOutcome.class); - @Mock private UpdateItemOutcome updateOutcome = Mockito.mock(UpdateItemOutcome.class); - + @Mock private PutItemResponse putOutcome = Mockito.mock(PutItemResponse.class); + @Mock private DeleteItemResponse deleteOutcome = Mockito.mock(DeleteItemResponse.class); + @Mock private UpdateItemResponse updateOutcome = Mockito.mock(UpdateItemResponse.class); @BeforeMethod public void setUp() { - MockitoAnnotations.initMocks(this); - Mockito.doReturn(table).when(dynamoDB).getTable(tableName); - Mockito.doReturn(currentTimeIndex).when(table).getIndex(currentTimeIndexName); - Mockito.doReturn(hostNameIndex).when(table).getIndex(hostNameIndexName); + MockitoAnnotations.openMocks(this); } private DynamoDBCertRecordStoreConnection getDBConnection() { @@ -73,14 +58,26 @@ private DynamoDBCertRecordStoreConnection getDBConnection() { @Test public void testGetX509CertRecord() { + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("athenz.provider:cn:1234")); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); + + Map attrs = new HashMap<>(); + Date now = new Date(); - long tstamp = mockNonNullableColumns(now, false); - Mockito.doReturn(tstamp).when(item).getLong("lastNotifiedTime"); - Mockito.doReturn(tstamp).when(item).get("lastNotifiedTime"); - Mockito.doReturn("last-notified-server").when(item).getString("lastNotifiedServer"); - Mockito.doReturn(tstamp).when(item).getLong("expiryTime"); - Mockito.doReturn(tstamp).when(item).get("expiryTime"); - Mockito.doReturn("hostname").when(item).getString("hostName"); + long tstamp = mockNonNullableColumns(attrs, now); + + attrs.put("lastNotifiedTime", AttributeValue.fromN(String.valueOf(tstamp))); + attrs.put("hostName", AttributeValue.fromS("hostname")); + attrs.put("expiryTime", AttributeValue.fromN(String.valueOf(tstamp))); + attrs.put("lastNotifiedServer", AttributeValue.fromS("last-notified-server")); + + GetItemResponse response = GetItemResponse.builder().item(attrs).build(); + Mockito.doReturn(response).when(dynamoDB).getItem(request); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); dbConn.setOperationTimeout(10); @@ -91,7 +88,7 @@ public void testGetX509CertRecord() { assertEquals(certRecord.getLastNotifiedServer(), "last-notified-server"); assertEquals(certRecord.getExpiryTime(), now); assertEquals(certRecord.getHostName(), "hostname"); - assertEquals(certRecord.getClientCert(), false); + assertFalse(certRecord.getClientCert()); dbConn.close(); } @@ -99,12 +96,25 @@ public void testGetX509CertRecord() { @Test public void testGetX509CertRecordNullableColumns() { + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("athenz.provider:cn:1234")); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); + + Map attrs = new HashMap<>(); + Date now = new Date(); - mockNonNullableColumns(now, true); - Mockito.doReturn(true).when(item).isNull("lastNotifiedTime"); - Mockito.doReturn(true).when(item).isNull("lastNotifiedServer"); - Mockito.doReturn(true).when(item).isNull("expiryTime"); - Mockito.doReturn(true).when(item).isNull("hostName"); + mockNonNullableColumns(attrs, now); + attrs.remove("lastNotifiedTime"); + attrs.remove("lastNotifiedServer"); + attrs.remove("expiryTime"); + attrs.remove("hostName"); + + GetItemResponse response = GetItemResponse.builder().item(attrs).build(); + Mockito.doReturn(response).when(dynamoDB).getItem(request); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); dbConn.setOperationTimeout(10); @@ -122,19 +132,32 @@ public void testGetX509CertRecordNullableColumns() { @Test public void testGetX509CertRecordNotFoundNull() { - Mockito.doReturn(null).when(table).getItem("primaryKey", "athenz.provider:cn:1234"); + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("athenz.provider:cn:1234")); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); + + GetItemResponse response = GetItemResponse.builder().item(null).build(); + Mockito.when(dynamoDB.getItem(request)).thenReturn(response).thenReturn(null); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); X509CertRecord certRecord = dbConn.getX509CertRecord("athenz.provider", "1234", "cn"); assertNull(certRecord); + + certRecord = dbConn.getX509CertRecord("athenz.provider", "1234", "cn"); + assertNull(certRecord); + dbConn.close(); } @Test public void testGetX509CertRecordNotFoundException() { - Mockito.doThrow(new AmazonDynamoDBException("item not found")) - .when(table).getItem("primaryKey", "athenz.provider:cn:1234"); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).getItem(ArgumentMatchers.any(GetItemRequest.class)); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); X509CertRecord certRecord = dbConn.getX509CertRecord("athenz.provider", "1234", "cn"); @@ -145,7 +168,8 @@ public void testGetX509CertRecordNotFoundException() { @Test public void testInsertX509Record() { - DynamoDBCertRecordStoreConnection dbConn = new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, currentTimeIndexName, hostNameIndexName); + DynamoDBCertRecordStoreConnection dbConn = new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, + currentTimeIndexName, hostNameIndexName); Date now = new Date(); String dateIsoFormat = DynamoDBUtils.getIso8601FromDate(now); @@ -155,39 +179,45 @@ public void testInsertX509Record() { certRecord.setExpiryTime(now); certRecord.setHostName("hostname"); - Item item = new Item() - .withPrimaryKey("primaryKey", "athenz.provider:cn:1234") - .withString("instanceId", certRecord.getInstanceId()) - .withString("provider", certRecord.getProvider()) - .withString("service", certRecord.getService()) - .withString("currentSerial", certRecord.getCurrentSerial()) - .withString("currentIP", certRecord.getCurrentIP()) - .withLong("currentTime", certRecord.getCurrentTime().getTime()) - .withString("currentDate", dateIsoFormat) - .withString("prevSerial", certRecord.getPrevSerial()) - .withString("prevIP", certRecord.getPrevIP()) - .withLong("prevTime", certRecord.getPrevTime().getTime()) - .withBoolean("clientCert", certRecord.getClientCert()) - .withLong("ttl", certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720) - .withLong("lastNotifiedTime", certRecord.getLastNotifiedTime().getTime()) - .withString("lastNotifiedServer", certRecord.getLastNotifiedServer()) - .withLong("expiryTime", certRecord.getExpiryTime().getTime()) - .withString("hostName", certRecord.getHostName()); - - Mockito.doReturn(putOutcome).when(table).putItem(item); + HashMap itemValues = new HashMap<>(); + + itemValues.put("primaryKey", AttributeValue.fromS("athenz.provider:cn:1234")); + itemValues.put("instanceId", AttributeValue.fromS(certRecord.getInstanceId())); + itemValues.put("provider", AttributeValue.fromS(certRecord.getProvider())); + itemValues.put("service", AttributeValue.fromS(certRecord.getService())); + itemValues.put("currentSerial", AttributeValue.fromS(certRecord.getCurrentSerial())); + itemValues.put("currentIP", AttributeValue.fromS(certRecord.getCurrentIP())); + itemValues.put("currentTime", AttributeValue.fromN(String.valueOf(certRecord.getCurrentTime().getTime()))); + itemValues.put("currentDate", AttributeValue.fromS(dateIsoFormat)); + itemValues.put("prevSerial", AttributeValue.fromS(certRecord.getPrevSerial())); + itemValues.put("prevIP", AttributeValue.fromS(certRecord.getPrevIP())); + itemValues.put("prevTime", AttributeValue.fromN(String.valueOf(certRecord.getPrevTime().getTime()))); + itemValues.put("clientCert", AttributeValue.fromBool(certRecord.getClientCert())); + itemValues.put("ttl", AttributeValue.fromN(String.valueOf(certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720))); + itemValues.put("lastNotifiedTime", AttributeValue.fromN(String.valueOf(certRecord.getLastNotifiedTime().getTime()))); + itemValues.put("lastNotifiedServer", AttributeValue.fromS(certRecord.getLastNotifiedServer())); + itemValues.put("expiryTime", AttributeValue.fromN(String.valueOf(certRecord.getExpiryTime().getTime()))); + itemValues.put("hostName", AttributeValue.fromS(certRecord.getHostName())); + + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(itemValues) + .build(); + + Mockito.doReturn(putOutcome).when(dynamoDB).putItem(request); boolean requestSuccess = dbConn.insertX509CertRecord(certRecord); assertTrue(requestSuccess); - ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(Item.class); - Mockito.verify(table, times(1)).putItem(itemCaptor.capture()); - List allValues = itemCaptor.getAllValues(); + ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(PutItemRequest.class); + Mockito.verify(dynamoDB, times(1)).putItem(itemCaptor.capture()); + List allValues = itemCaptor.getAllValues(); assertEquals(1, allValues.size()); - assertEquals(allValues.get(0).get("primaryKey"), item.get("primaryKey")); - assertEquals(allValues.get(0).get("provider"), item.get("provider")); - assertEquals(allValues.get(0).get("instanceId"), item.get("instanceId")); - assertEquals(allValues.get(0).get("service"), item.get("service")); - assertEquals(allValues.get(0).get("expiryTime"), item.get("expiryTime")); - assertEquals(allValues.get(0).get("hostName"), item.get("hostName")); + Map item = allValues.get(0).item(); + assertEquals(DynamoDBUtils.getString(item, "primaryKey"), "athenz.provider:cn:1234"); + assertEquals(DynamoDBUtils.getString(item, "provider"), "athenz.provider"); + assertEquals(DynamoDBUtils.getString(item, "instanceId"),"1234"); + assertEquals(DynamoDBUtils.getString(item, "service"), "cn"); + assertEquals(DynamoDBUtils.getString(item, "hostName"), "hostname"); dbConn.close(); } @@ -195,7 +225,8 @@ public void testInsertX509Record() { @Test public void testInsertX509RecordNoHostname() { - DynamoDBCertRecordStoreConnection dbConn = new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, currentTimeIndexName, hostNameIndexName); + DynamoDBCertRecordStoreConnection dbConn = new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, + currentTimeIndexName, hostNameIndexName); Date now = new Date(); String dateIsoFormat = DynamoDBUtils.getIso8601FromDate(now); @@ -204,40 +235,46 @@ public void testInsertX509RecordNoHostname() { certRecord.setLastNotifiedServer("last-notified-server"); certRecord.setExpiryTime(now); - Item item = new Item() - .withPrimaryKey("primaryKey", "athenz.provider:cn:1234") - .withString("instanceId", certRecord.getInstanceId()) - .withString("provider", certRecord.getProvider()) - .withString("service", certRecord.getService()) - .withString("currentSerial", certRecord.getCurrentSerial()) - .withString("currentIP", certRecord.getCurrentIP()) - .withLong("currentTime", certRecord.getCurrentTime().getTime()) - .withString("currentDate", dateIsoFormat) - .withString("prevSerial", certRecord.getPrevSerial()) - .withString("prevIP", certRecord.getPrevIP()) - .withLong("prevTime", certRecord.getPrevTime().getTime()) - .withBoolean("clientCert", certRecord.getClientCert()) - .withLong("ttl", certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720) - .withLong("lastNotifiedTime", certRecord.getLastNotifiedTime().getTime()) - .withString("lastNotifiedServer", certRecord.getLastNotifiedServer()) - .withLong("expiryTime", certRecord.getExpiryTime().getTime()); - - Mockito.doReturn(putOutcome).when(table).putItem(item); + HashMap itemValues = new HashMap<>(); + + itemValues.put("primaryKey", AttributeValue.fromS("athenz.provider:cn:1234")); + itemValues.put("instanceId", AttributeValue.fromS(certRecord.getInstanceId())); + itemValues.put("provider", AttributeValue.fromS(certRecord.getProvider())); + itemValues.put("service", AttributeValue.fromS(certRecord.getService())); + itemValues.put("currentSerial", AttributeValue.fromS(certRecord.getCurrentSerial())); + itemValues.put("currentIP", AttributeValue.fromS(certRecord.getCurrentIP())); + itemValues.put("currentTime", AttributeValue.fromN(String.valueOf(certRecord.getCurrentTime().getTime()))); + itemValues.put("currentDate", AttributeValue.fromS(dateIsoFormat)); + itemValues.put("prevSerial", AttributeValue.fromS(certRecord.getPrevSerial())); + itemValues.put("prevIP", AttributeValue.fromS(certRecord.getPrevIP())); + itemValues.put("prevTime", AttributeValue.fromN(String.valueOf(certRecord.getPrevTime().getTime()))); + itemValues.put("clientCert", AttributeValue.fromBool(certRecord.getClientCert())); + itemValues.put("ttl", AttributeValue.fromN(String.valueOf(certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720))); + itemValues.put("lastNotifiedTime", AttributeValue.fromN(String.valueOf(certRecord.getLastNotifiedTime().getTime()))); + itemValues.put("lastNotifiedServer", AttributeValue.fromS(certRecord.getLastNotifiedServer())); + itemValues.put("expiryTime", AttributeValue.fromN(String.valueOf(certRecord.getExpiryTime().getTime()))); + + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(itemValues) + .build(); + + Mockito.doReturn(putOutcome).when(dynamoDB).putItem(request); boolean requestSuccess = dbConn.insertX509CertRecord(certRecord); assertTrue(requestSuccess); - ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(Item.class); - Mockito.verify(table, times(1)).putItem(itemCaptor.capture()); - List allValues = itemCaptor.getAllValues(); + ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(PutItemRequest.class); + Mockito.verify(dynamoDB, times(1)).putItem(itemCaptor.capture()); + List allValues = itemCaptor.getAllValues(); assertEquals(1, allValues.size()); - assertEquals(allValues.get(0).get("primaryKey"), item.get("primaryKey")); - assertEquals(allValues.get(0).get("provider"), item.get("provider")); - assertEquals(allValues.get(0).get("instanceId"), item.get("instanceId")); - assertEquals(allValues.get(0).get("service"), item.get("service")); - assertEquals(allValues.get(0).get("expiryTime"), item.get("expiryTime")); + Map item = allValues.get(0).item(); + assertEquals(DynamoDBUtils.getString(item, "primaryKey"), "athenz.provider:cn:1234"); + assertEquals(DynamoDBUtils.getString(item, "provider"), "athenz.provider"); + assertEquals(DynamoDBUtils.getString(item, "instanceId"), "1234"); + assertEquals(DynamoDBUtils.getString(item, "service"), "cn"); // When hostname is null, primaryKey will be used - assertEquals(allValues.get(0).get("hostName"), item.get("primaryKey")); + assertEquals(DynamoDBUtils.getString(item, "hostName"), "athenz.provider:cn:1234"); dbConn.close(); } @@ -245,7 +282,8 @@ public void testInsertX509RecordNoHostname() { @Test public void testInsertX509RecordNullableColumns() { - DynamoDBCertRecordStoreConnection dbConn = new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, currentTimeIndexName, hostNameIndexName); + DynamoDBCertRecordStoreConnection dbConn = new DynamoDBCertRecordStoreConnection(dynamoDB, tableName, + currentTimeIndexName, hostNameIndexName); Date now = new Date(); String dateIsoFormat = DynamoDBUtils.getIso8601FromDate(now); @@ -255,26 +293,31 @@ public void testInsertX509RecordNullableColumns() { certRecord.setExpiryTime(null); certRecord.setHostName(null); - Item item = new Item() - .withPrimaryKey("primaryKey", "athenz.provider:cn:1234") - .withString("instanceId", certRecord.getInstanceId()) - .withString("provider", certRecord.getProvider()) - .withString("service", certRecord.getService()) - .withString("currentSerial", certRecord.getCurrentSerial()) - .withString("currentIP", certRecord.getCurrentIP()) - .withLong("currentTime", certRecord.getCurrentTime().getTime()) - .withString("currentDate", dateIsoFormat) - .withString("prevSerial", certRecord.getPrevSerial()) - .withString("prevIP", certRecord.getPrevIP()) - .withLong("prevTime", certRecord.getPrevTime().getTime()) - .withBoolean("clientCert", certRecord.getClientCert()) - .withLong("ttl", certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720) - .with("lastNotifiedTime", null) - .with("lastNotifiedServer", null) - .with("expiryTime", null) - .with("hostName", null); - - Mockito.doReturn(putOutcome).when(table).putItem(item); + HashMap itemValues = new HashMap<>(); + + itemValues.put("primaryKey", AttributeValue.fromS("athenz.provider:cn:1234")); + itemValues.put("instanceId", AttributeValue.fromS(certRecord.getInstanceId())); + itemValues.put("provider", AttributeValue.fromS(certRecord.getProvider())); + itemValues.put("service", AttributeValue.fromS(certRecord.getService())); + itemValues.put("currentSerial", AttributeValue.fromS(certRecord.getCurrentSerial())); + itemValues.put("currentIP", AttributeValue.fromS(certRecord.getCurrentIP())); + itemValues.put("currentTime", AttributeValue.fromN(String.valueOf(certRecord.getCurrentTime().getTime()))); + itemValues.put("currentDate", AttributeValue.fromS(dateIsoFormat)); + itemValues.put("prevSerial", AttributeValue.fromS(certRecord.getPrevSerial())); + itemValues.put("prevIP", AttributeValue.fromS(certRecord.getPrevIP())); + itemValues.put("prevTime", AttributeValue.fromN(String.valueOf(certRecord.getPrevTime().getTime()))); + itemValues.put("clientCert", AttributeValue.fromBool(certRecord.getClientCert())); + itemValues.put("ttl", AttributeValue.fromN(String.valueOf(certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720))); + itemValues.remove("lastNotifiedTime"); + itemValues.remove("lastNotifiedServer"); + itemValues.remove("expiryTime"); + + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(itemValues) + .build(); + + Mockito.doReturn(putOutcome).when(dynamoDB).putItem(request); boolean requestSuccess = dbConn.insertX509CertRecord(certRecord); assertTrue(requestSuccess); @@ -287,8 +330,8 @@ public void testInsertX509RecordException() { Date now = new Date(); X509CertRecord certRecord = getRecordNonNullableColumns(now); - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).putItem(any(Item.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).putItem(ArgumentMatchers.any(PutItemRequest.class)); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); boolean requestSuccess = dbConn.insertX509CertRecord(certRecord); @@ -310,45 +353,19 @@ public void testUpdateX509Record() { certRecord.setHostName("hostname"); certRecord.setSvcDataUpdateTime(now); - UpdateItemSpec item = new UpdateItemSpec() - .withPrimaryKey("primaryKey", "athenz.provider:cn:1234") - .withAttributeUpdate( - new AttributeUpdate("instanceId").put(certRecord.getInstanceId()), - new AttributeUpdate("provider").put(certRecord.getProvider()), - new AttributeUpdate("service").put(certRecord.getService()), - new AttributeUpdate("currentSerial").put(certRecord.getCurrentSerial()), - new AttributeUpdate("currentIP").put(certRecord.getCurrentIP()), - new AttributeUpdate("currentTime").put(certRecord.getCurrentTime().getTime()), - new AttributeUpdate("currentDate").put(DynamoDBUtils.getIso8601FromDate(certRecord.getCurrentTime())), - new AttributeUpdate("prevSerial").put(certRecord.getPrevSerial()), - new AttributeUpdate("prevIP").put(certRecord.getPrevIP()), - new AttributeUpdate("prevTime").put(certRecord.getPrevTime().getTime()), - new AttributeUpdate("clientCert").put(certRecord.getClientCert()), - new AttributeUpdate("ttl").put(certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720), - new AttributeUpdate("svcDataUpdateTime").put(certRecord.getSvcDataUpdateTime().getTime()), - new AttributeUpdate("expiryTime").put(certRecord.getExpiryTime().getTime()), - new AttributeUpdate("hostName").put(certRecord.getHostName())); - - Mockito.doReturn(updateOutcome).when(table).updateItem(item); + Mockito.doReturn(updateOutcome).when(dynamoDB).updateItem(any(UpdateItemRequest.class)); boolean requestSuccess = dbConn.updateX509CertRecord(certRecord); assertTrue(requestSuccess); - ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(UpdateItemSpec.class); - Mockito.verify(table, times(1)).updateItem(itemCaptor.capture()); - List allValues = itemCaptor.getAllValues(); + ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class); + Mockito.verify(dynamoDB, times(1)).updateItem(itemCaptor.capture()); + List allValues = itemCaptor.getAllValues(); assertEquals(1, allValues.size()); - UpdateItemSpec capturedItem = allValues.get(0); - assertEquals(capturedItem.getKeyComponents().toArray()[0].toString(), item.getKeyComponents().toArray()[0].toString()); - - List capturedAttributes = capturedItem.getAttributeUpdate(); - List expectedAttributes = item.getAttributeUpdate(); - - for (int i = 0; i < expectedAttributes.size(); ++i) { - System.out.println("expected attr: " + expectedAttributes.get(i).getAttributeName() + ", value: " + expectedAttributes.get(i).getValue()); - assertEquals(capturedAttributes.get(i).getAttributeName(), expectedAttributes.get(i).getAttributeName()); - assertEquals(capturedAttributes.get(i).getValue(), expectedAttributes.get(i).getValue()); - } + assertEquals(allValues.get(0).attributeUpdates().get("provider").value().s(), "athenz.provider"); + assertEquals(allValues.get(0).attributeUpdates().get("instanceId").value().s(),"1234"); + assertEquals(allValues.get(0).attributeUpdates().get("service").value().s(), "cn"); + assertEquals(allValues.get(0).attributeUpdates().get("hostName").value().s(), "hostname"); dbConn.close(); } @@ -365,50 +382,19 @@ public void testUpdateX509RecordNoHostName() { certRecord.setExpiryTime(now); certRecord.setSvcDataUpdateTime(now); - UpdateItemSpec item = new UpdateItemSpec() - .withPrimaryKey("primaryKey", "athenz.provider:cn:1234") - .withAttributeUpdate( - new AttributeUpdate("instanceId").put(certRecord.getInstanceId()), - new AttributeUpdate("provider").put(certRecord.getProvider()), - new AttributeUpdate("service").put(certRecord.getService()), - new AttributeUpdate("currentSerial").put(certRecord.getCurrentSerial()), - new AttributeUpdate("currentIP").put(certRecord.getCurrentIP()), - new AttributeUpdate("currentTime").put(certRecord.getCurrentTime().getTime()), - new AttributeUpdate("currentDate").put(DynamoDBUtils.getIso8601FromDate(certRecord.getCurrentTime())), - new AttributeUpdate("prevSerial").put(certRecord.getPrevSerial()), - new AttributeUpdate("prevIP").put(certRecord.getPrevIP()), - new AttributeUpdate("prevTime").put(certRecord.getPrevTime().getTime()), - new AttributeUpdate("clientCert").put(certRecord.getClientCert()), - new AttributeUpdate("ttl").put(certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720), - new AttributeUpdate("svcDataUpdateTime").put(certRecord.getSvcDataUpdateTime().getTime()), - new AttributeUpdate("expiryTime").put(certRecord.getExpiryTime().getTime())); - - Mockito.doReturn(updateOutcome).when(table).updateItem(item); + Mockito.doReturn(updateOutcome).when(dynamoDB).updateItem(any(UpdateItemRequest.class)); boolean requestSuccess = dbConn.updateX509CertRecord(certRecord); assertTrue(requestSuccess); - ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(UpdateItemSpec.class); - Mockito.verify(table, times(1)).updateItem(itemCaptor.capture()); - List allValues = itemCaptor.getAllValues(); + ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class); + Mockito.verify(dynamoDB, times(1)).updateItem(itemCaptor.capture()); + List allValues = itemCaptor.getAllValues(); assertEquals(1, allValues.size()); - UpdateItemSpec capturedItem = allValues.get(0); - assertEquals(capturedItem.getKeyComponents().toArray()[0].toString(), item.getKeyComponents().toArray()[0].toString()); - - List capturedAttributes = capturedItem.getAttributeUpdate(); - List expectedAttributes = item.getAttributeUpdate(); - - // Check everyone except the hostname (it will be filled with the primaryKey value as the hostName index doesn't allow nulls) - for (int i = 0; i < capturedAttributes.size() - 1; ++i) { - System.out.println("expected attr: " + expectedAttributes.get(i).getAttributeName() + ", value: " + expectedAttributes.get(i).getValue()); - assertEquals(capturedAttributes.get(i).getAttributeName(), expectedAttributes.get(i).getAttributeName()); - assertEquals(capturedAttributes.get(i).getValue(), expectedAttributes.get(i).getValue()); - } - - // Make sure hostName received the value of the primaryKey - System.out.println("expected attr: hostName, value: athenz.provider:cn:1234"); - assertEquals(capturedAttributes.get(capturedAttributes.size() - 1).getAttributeName(), "hostName"); - assertEquals(capturedAttributes.get(capturedAttributes.size() - 1).getValue(), "athenz.provider:cn:1234"); + assertEquals(allValues.get(0).attributeUpdates().get("provider").value().s(), "athenz.provider"); + assertEquals(allValues.get(0).attributeUpdates().get("instanceId").value().s(),"1234"); + assertEquals(allValues.get(0).attributeUpdates().get("service").value().s(), "cn"); + assertEquals(allValues.get(0).attributeUpdates().get("hostName").value().s(), "athenz.provider:cn:1234"); dbConn.close(); } @@ -426,28 +412,20 @@ public void testUpdateX509RecordNullableColumns() { certRecord.setHostName(null); certRecord.setSvcDataUpdateTime(now); - UpdateItemSpec item = new UpdateItemSpec() - .withPrimaryKey("primaryKey", "athenz.provider:cn:1234") - .withAttributeUpdate( - new AttributeUpdate("instanceId").put(certRecord.getInstanceId()), - new AttributeUpdate("provider").put(certRecord.getProvider()), - new AttributeUpdate("service").put(certRecord.getService()), - new AttributeUpdate("currentSerial").put(certRecord.getCurrentSerial()), - new AttributeUpdate("currentIP").put(certRecord.getCurrentIP()), - new AttributeUpdate("currentTime").put(certRecord.getCurrentTime().getTime()), - new AttributeUpdate("currentDate").put(DynamoDBUtils.getIso8601FromDate(certRecord.getCurrentTime())), - new AttributeUpdate("prevSerial").put(certRecord.getPrevSerial()), - new AttributeUpdate("prevIP").put(certRecord.getPrevIP()), - new AttributeUpdate("prevTime").put(certRecord.getPrevTime().getTime()), - new AttributeUpdate("clientCert").put(certRecord.getClientCert()), - new AttributeUpdate("ttl").put(certRecord.getCurrentTime().getTime() / 1000L + 3660 * 720), - new AttributeUpdate("svcDataUpdateTime").put(certRecord.getSvcDataUpdateTime().getTime()), - new AttributeUpdate("expiryTime").put(null)); - - Mockito.doReturn(updateOutcome).when(table).updateItem(item); + Mockito.doReturn(updateOutcome).when(dynamoDB).updateItem(any(UpdateItemRequest.class)); boolean requestSuccess = dbConn.updateX509CertRecord(certRecord); assertTrue(requestSuccess); + ArgumentCaptor itemCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class); + Mockito.verify(dynamoDB, times(1)).updateItem(itemCaptor.capture()); + List allValues = itemCaptor.getAllValues(); + assertEquals(1, allValues.size()); + + assertEquals(allValues.get(0).attributeUpdates().get("provider").value().s(), "athenz.provider"); + assertEquals(allValues.get(0).attributeUpdates().get("instanceId").value().s(),"1234"); + assertEquals(allValues.get(0).attributeUpdates().get("service").value().s(), "cn"); + assertEquals(allValues.get(0).attributeUpdates().get("hostName").value().s(), "athenz.provider:cn:1234"); + dbConn.close(); } @@ -457,8 +435,8 @@ public void testUpdateX509RecordException() { Date now = new Date(); X509CertRecord certRecord = getRecordNonNullableColumns(now); - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).updateItem(any(UpdateItemSpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).updateItem(ArgumentMatchers.any(UpdateItemRequest.class)); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); boolean requestSuccess = dbConn.updateX509CertRecord(certRecord); @@ -469,9 +447,16 @@ public void testUpdateX509RecordException() { @Test public void testDeleteX509Record() { - DeleteItemSpec deleteItemSpec = new DeleteItemSpec() - .withPrimaryKey("primaryKey", "athenz.provider:cn:1234"); - Mockito.doReturn(deleteOutcome).when(table).deleteItem(deleteItemSpec); + + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("athenz.provider:cn:1234")); + + DeleteItemRequest request = DeleteItemRequest.builder() + .tableName(tableName) + .key(keyToGet) + .build(); + + Mockito.doReturn(deleteOutcome).when(dynamoDB).deleteItem(request); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); @@ -483,8 +468,8 @@ public void testDeleteX509Record() { @Test public void testDeleteX509RecordException() { - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).deleteItem(any(DeleteItemSpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).deleteItem(ArgumentMatchers.any(DeleteItemRequest.class)); DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); @@ -530,32 +515,20 @@ private void assertNonNullableColumns(Date now, X509CertRecord certRecord) { assertFalse(certRecord.getClientCert()); } - private long mockNonNullableColumns(Date now, boolean clientCertHasValue) { + private long mockNonNullableColumns(Map attrs, Date now) { + long tstamp = now.getTime(); - Mockito.doReturn(item).when(table).getItem("primaryKey", "athenz.provider:cn:1234"); - - Mockito.doReturn("athenz.provider").when(item).getString("provider"); - Mockito.doReturn("1234").when(item).getString("instanceId"); - Mockito.doReturn("cn").when(item).getString("service"); - - Mockito.doReturn(false).when(item).isNull("currentTime"); - Mockito.doReturn(false).when(item).isNull("prevTime"); - - Mockito.doReturn("cn").when(item).getString("service"); - Mockito.doReturn("current-serial").when(item).getString("currentSerial"); - Mockito.doReturn("current-ip").when(item).getString("currentIP"); - Mockito.doReturn(tstamp).when(item).getLong("currentTime"); - Mockito.doReturn(tstamp).when(item).get("currentTime"); - Mockito.doReturn("prev-serial").when(item).getString("prevSerial"); - Mockito.doReturn("prev-ip").when(item).getString("prevIP"); - Mockito.doReturn(tstamp).when(item).getLong("prevTime"); - Mockito.doReturn(tstamp).when(item).get("prevTime"); - if (clientCertHasValue) { - Mockito.doReturn(false).when(item).getBoolean("clientCert"); - } else { - Mockito.when(item.getBoolean(anyString())).thenThrow(new IncompatibleTypeException("Value not found")); - } + attrs.put("provider", AttributeValue.fromS("athenz.provider")); + attrs.put("instanceId", AttributeValue.fromS("1234")); + attrs.put("service", AttributeValue.fromS("cn")); + attrs.put("currentSerial", AttributeValue.fromS("current-serial")); + attrs.put("currentTime", AttributeValue.fromN(String.valueOf(tstamp))); + attrs.put("currentIP", AttributeValue.fromS("current-ip")); + attrs.put("prevSerial", AttributeValue.fromS("prev-serial")); + attrs.put("prevIP", AttributeValue.fromS("prev-ip")); + attrs.put("prevTime", AttributeValue.fromN(String.valueOf(tstamp))); + return tstamp; } @@ -568,77 +541,40 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestamp() { long sevenDaysAgo = nowL - 7 * 24 * 60 * 60 * 1000; Map unNotified = ZTSTestUtils.generateAttributeValues( - "home.test.service2", - "unNotified", - null, - null, - null, - null, - "testHost1"); + "home.test.service2", "unNotified", null, null, null, null, "testHost1"); Map reNotified = ZTSTestUtils.generateAttributeValues( - "home.test.service3", - "reNotified", - Long.toString(fiveDaysAgo), - Long.toString(fiveDaysAgo), - "testServer", - null, - "testHost2"); + "home.test.service3", "reNotified", Long.toString(fiveDaysAgo), Long.toString(fiveDaysAgo), + "testServer", null, "testHost2"); Map rebootstrapped = ZTSTestUtils.generateAttributeValues( - "home.test.service3", - "rebootstrapped", - Long.toString(sevenDaysAgo), - Long.toString(sevenDaysAgo), - "testServer", - null, - "testHost2"); + "home.test.service3", "rebootstrapped", Long.toString(sevenDaysAgo), Long.toString(sevenDaysAgo), + "testServer", null, "testHost2"); Map willBeUpdatedByOtherZts = ZTSTestUtils.generateAttributeValues( - "home.test.service4", - "willBeUpdatedByOtherZts", - Long.toString(fiveDaysAgo), - Long.toString(fiveDaysAgo), - "testServer", - null, - "testHost3"); - - Item item1 = ItemUtils.toItem(unNotified); - Item item2 = ItemUtils.toItem(reNotified); - Item item3 = ItemUtils.toItem(willBeUpdatedByOtherZts); - Item item4 = ItemUtils.toItem(rebootstrapped); - - ItemCollection itemCollection = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport = Mockito.mock(IteratorSupport.class); - when(itemCollection.iterator()).thenReturn(iteratorSupport); - when(iteratorSupport.hasNext()).thenReturn(true, true, true, true, false); - when(iteratorSupport.next()).thenReturn(item1).thenReturn(item2).thenReturn(item3).thenReturn(item4); - - Mockito.doReturn(itemCollection).when(currentTimeIndex).query(any(QuerySpec.class)); - - ItemCollection itemCollection2 = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport2 = Mockito.mock(IteratorSupport.class); - when(itemCollection2.iterator()).thenReturn(iteratorSupport2); - when(iteratorSupport2.hasNext()).thenReturn( - true, false, // One record with host testHost1 - true, true, false, // Two records with host testHost2 - true, false, // One record with host testHost3 - true, true, false); // Two records with host testHost2 - - when(iteratorSupport2.next()) - .thenReturn(item1) - .thenReturn(item2).thenReturn(item4) - .thenReturn(item3) - .thenReturn(item2).thenReturn(item4); - - Mockito.doReturn(itemCollection2).when(hostNameIndex).query(any(QuerySpec.class)); - - AttributeValue lastNotifiedTimeAttrValue = new AttributeValue(); - lastNotifiedTimeAttrValue.setN(Long.toString(nowL)); - AttributeValue lastNotifiedServerAttrValue = new AttributeValue(); - lastNotifiedServerAttrValue.setS("localhost"); - AttributeValue lastNotifiedOtherServerAttrValue = new AttributeValue(); - lastNotifiedOtherServerAttrValue.setS("SomeOtherZTS"); + "home.test.service4", "willBeUpdatedByOtherZts", Long.toString(fiveDaysAgo), Long.toString(fiveDaysAgo), + "testServer", null, "testHost3"); + + QueryResponse response1 = QueryResponse.builder().items(Collections.singleton(unNotified)).build(); + QueryResponse response2 = QueryResponse.builder().items(List.of(reNotified, rebootstrapped)).build(); + QueryResponse response3 = QueryResponse.builder().items(Collections.singleton(willBeUpdatedByOtherZts)).build(); + QueryResponse response4 = QueryResponse.builder().items(List.of(reNotified, rebootstrapped)).build(); + + QueryResponse responseEmpty = QueryResponse.builder().build(); + QueryResponse responseFull = QueryResponse.builder().items(List.of(unNotified, reNotified, + willBeUpdatedByOtherZts, rebootstrapped)).build(); + Mockito.when(dynamoDB.query(any(QueryRequest.class))).thenReturn(responseFull) + .thenReturn(responseEmpty).thenReturn(responseEmpty).thenReturn(responseEmpty) + .thenReturn(responseEmpty).thenReturn(responseEmpty).thenReturn(responseEmpty) + .thenReturn(responseEmpty).thenReturn(responseEmpty).thenReturn(responseEmpty) + .thenReturn(responseEmpty).thenReturn(responseEmpty).thenReturn(responseEmpty) + .thenReturn(responseEmpty).thenReturn(responseEmpty).thenReturn(responseEmpty) + .thenReturn(responseEmpty).thenReturn(response1).thenReturn(response2) + .thenReturn(response3).thenReturn(response4); + + AttributeValue lastNotifiedTimeAttrValue = AttributeValue.fromN(Long.toString(nowL)); + AttributeValue lastNotifiedServerAttrValue = AttributeValue.fromS("localhost"); + AttributeValue lastNotifiedOtherServerAttrValue = AttributeValue.fromS("SomeOtherZTS"); unNotified.put("lastNotifiedTime", lastNotifiedTimeAttrValue); unNotified.put("lastNotifiedServer", lastNotifiedServerAttrValue); @@ -649,34 +585,29 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestamp() { willBeUpdatedByOtherZts.put("lastNotifiedTime", lastNotifiedTimeAttrValue); willBeUpdatedByOtherZts.put("lastNotifiedServer", lastNotifiedOtherServerAttrValue); - Item updatedItem1 = ItemUtils.toItem(unNotified); - Item updatedItem2 = ItemUtils.toItem(reNotified); - Item updatedItem3 = ItemUtils.toItem(willBeUpdatedByOtherZts); - - UpdateItemOutcome updateItemOutcome1 = Mockito.mock(UpdateItemOutcome.class); - when(updateItemOutcome1.getItem()).thenReturn(updatedItem1); + UpdateItemResponse updateItemOutcome1 = UpdateItemResponse.builder().attributes(unNotified).build(); + UpdateItemResponse updateItemOutcome2 = UpdateItemResponse.builder().attributes(reNotified).build(); + UpdateItemResponse updateItemOutcome3 = UpdateItemResponse.builder().attributes(willBeUpdatedByOtherZts).build(); - UpdateItemOutcome updateItemOutcome2 = Mockito.mock(UpdateItemOutcome.class); - when(updateItemOutcome2.getItem()).thenReturn(updatedItem2); + when(dynamoDB.updateItem(any(UpdateItemRequest.class))) + .thenReturn(updateItemOutcome1) + .thenReturn(updateItemOutcome2) + .thenReturn(updateItemOutcome3); - UpdateItemOutcome updateItemOutcome3 = Mockito.mock(UpdateItemOutcome.class); - when(updateItemOutcome3.getItem()).thenReturn(updatedItem3); - - when(table.updateItem(any(UpdateItemSpec.class))).thenReturn(updateItemOutcome1).thenReturn(updateItemOutcome2).thenReturn(updateItemOutcome3); List records = dbConn.updateUnrefreshedCertificatesNotificationTimestamp( "localhost", nowL, "provider"); - ArgumentCaptor updateArguments = ArgumentCaptor.forClass(UpdateItemSpec.class); - Mockito.verify(table, Mockito.times(3)).updateItem(updateArguments.capture()); + ArgumentCaptor updateArguments = ArgumentCaptor.forClass(UpdateItemRequest.class); + Mockito.verify(dynamoDB, Mockito.times(3)).updateItem(updateArguments.capture()); // Assert get filtered records - List allUpdateArguments = updateArguments.getAllValues(); - assertEquals(3, allUpdateArguments.size()); - assertEquals("{primaryKey: provider:home.test.service2:unNotified}", allUpdateArguments.get(0).getKeyComponents().toArray()[0].toString()); - assertEquals("{primaryKey: provider:home.test.service3:reNotified}", allUpdateArguments.get(1).getKeyComponents().toArray()[0].toString()); - assertEquals("{primaryKey: provider:home.test.service4:willBeUpdatedByOtherZts}", allUpdateArguments.get(2).getKeyComponents().toArray()[0].toString()); + List allUpdateArguments = updateArguments.getAllValues(); + assertEquals(allUpdateArguments.size(), 3); + assertEquals(allUpdateArguments.get(0).key().get("primaryKey").toString(), "AttributeValue(S=provider:home.test.service2:unNotified)"); + assertEquals(allUpdateArguments.get(1).key().get("primaryKey").toString(), "AttributeValue(S=provider:home.test.service3:reNotified)"); + assertEquals(allUpdateArguments.get(2).key().get("primaryKey").toString(), "AttributeValue(S=provider:home.test.service4:willBeUpdatedByOtherZts)"); // Assert Update assertEquals(records.size(), 2); @@ -695,6 +626,7 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestampUpdateDynamoDB long nowL = now.getTime(); long fiveDaysAgo = nowL - 5 * 24 * 60 * 60 * 1000; + List> items = new ArrayList<>(); Map reNotified = ZTSTestUtils.generateAttributeValues( "home.test.service3", "reNotified", @@ -703,28 +635,13 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestampUpdateDynamoDB "testServer", null, "testHost2"); + items.add(reNotified); - Item item1 = ItemUtils.toItem(reNotified); - - ItemCollection itemCollection = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport = Mockito.mock(IteratorSupport.class); - when(itemCollection.iterator()).thenReturn(iteratorSupport); - when(iteratorSupport.hasNext()).thenReturn(true, false); - when(iteratorSupport.next()).thenReturn(item1); - - Mockito.doReturn(itemCollection).when(currentTimeIndex).query(any(QuerySpec.class)); + QueryResponse response = QueryResponse.builder().items(items).build(); + Mockito.doReturn(response).when(dynamoDB).query(any(QueryRequest.class)); - ItemCollection itemCollection2 = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport2 = Mockito.mock(IteratorSupport.class); - when(itemCollection2.iterator()).thenReturn(iteratorSupport2); - when(iteratorSupport2.hasNext()).thenReturn(true, false); - - when(iteratorSupport2.next()).thenReturn(item1); - - Mockito.doReturn(itemCollection2).when(hostNameIndex).query(any(QuerySpec.class)); - - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).updateItem(any(UpdateItemSpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).updateItem(ArgumentMatchers.any(UpdateItemRequest.class)); List result = dbConn.updateUnrefreshedCertificatesNotificationTimestamp( "serverTest", @@ -732,7 +649,6 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestampUpdateDynamoDB "providerTest"); assertEquals(result.size(), 0); - dbConn.close(); } @@ -743,6 +659,7 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestampUpdateExceptio long nowL = now.getTime(); long fiveDaysAgo = nowL - 5 * 24 * 60 * 60 * 1000; + List> items = new ArrayList<>(); Map reNotified = ZTSTestUtils.generateAttributeValues( "home.test.service3", "reNotified", @@ -751,28 +668,13 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestampUpdateExceptio "testServer", null, "testHost2"); + items.add(reNotified); - Item item1 = ItemUtils.toItem(reNotified); - - ItemCollection itemCollection = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport = Mockito.mock(IteratorSupport.class); - when(itemCollection.iterator()).thenReturn(iteratorSupport); - when(iteratorSupport.hasNext()).thenReturn(true, false); - when(iteratorSupport.next()).thenReturn(item1); - - Mockito.doReturn(itemCollection).when(currentTimeIndex).query(any(QuerySpec.class)); + QueryResponse response = QueryResponse.builder().items(items).build(); + Mockito.doReturn(response).when(dynamoDB).query(any(QueryRequest.class)); - ItemCollection itemCollection2 = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport2 = Mockito.mock(IteratorSupport.class); - when(itemCollection2.iterator()).thenReturn(iteratorSupport2); - when(iteratorSupport2.hasNext()).thenReturn(true, false); - - when(iteratorSupport2.next()).thenReturn(item1); - - Mockito.doReturn(itemCollection2).when(hostNameIndex).query(any(QuerySpec.class)); - - Mockito.doThrow(new TransactionConflictException("error")) - .when(table).updateItem(any(UpdateItemSpec.class)); + Mockito.doThrow(TransactionConflictException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).updateItem(ArgumentMatchers.any(UpdateItemRequest.class)); List result = dbConn.updateUnrefreshedCertificatesNotificationTimestamp( "serverTest", @@ -780,7 +682,6 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestampUpdateExceptio "providerTest"); assertEquals(result.size(), 0); - dbConn.close(); } @@ -788,47 +689,8 @@ public void testUpdateUnrefreshedCertificatesNotificationTimestampUpdateExceptio public void testUpdateUnrefreshedCertificatesNotificationTimestampTimeException() { DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); - Mockito.doThrow(new TransactionConflictException("error")) - .when(currentTimeIndex).query(any(QuerySpec.class)); - - List result = dbConn.updateUnrefreshedCertificatesNotificationTimestamp( - "serverTest", - 1591706189000L, - "providerTest"); - - assertEquals(result.size(), 0); - - dbConn.close(); - } - - @Test - public void testUpdateUnrefreshedCertificatesNotificationTimestampHostException() { - DynamoDBCertRecordStoreConnection dbConn = getDBConnection(); - Date now = new Date(1591706189000L); - long nowL = now.getTime(); - long fiveDaysAgo = nowL - 5 * 24 * 60 * 60 * 1000; - - Map reNotified = ZTSTestUtils.generateAttributeValues( - "home.test.service3", - "reNotified", - Long.toString(fiveDaysAgo), - Long.toString(fiveDaysAgo), - "testServer", - null, - "testHost2"); - - Item item1 = ItemUtils.toItem(reNotified); - - ItemCollection itemCollection = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport = Mockito.mock(IteratorSupport.class); - when(itemCollection.iterator()).thenReturn(iteratorSupport); - when(iteratorSupport.hasNext()).thenReturn(true, false); - when(iteratorSupport.next()).thenReturn(item1); - - Mockito.doReturn(itemCollection).when(currentTimeIndex).query(any(QuerySpec.class)); - - Mockito.doThrow(new TransactionConflictException("error")) - .when(hostNameIndex).query(any(QuerySpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).query(ArgumentMatchers.any(QueryRequest.class)); List result = dbConn.updateUnrefreshedCertificatesNotificationTimestamp( "serverTest", diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactoryTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactoryTest.java index 32fadd4e7e2..d8808f409de 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactoryTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreFactoryTest.java @@ -15,8 +15,6 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; -import com.amazonaws.services.dynamodbv2.document.Table; import com.yahoo.athenz.common.server.cert.CertRecordStore; import com.yahoo.athenz.zts.ResourceException; import com.yahoo.athenz.zts.notification.ZTSClientNotificationSenderImpl; @@ -29,7 +27,7 @@ import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.zts.ZTSConsts; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import static com.yahoo.athenz.zts.ZTSConsts.*; import static org.mockito.Mockito.when; @@ -39,22 +37,19 @@ public class DynamoDBCertRecordStoreFactoryTest { - @Mock private AmazonDynamoDB dbClient; - @Mock private Table table; - @Mock private DynamoDB dynamoDB; + @Mock private DynamoDbClient dbClient; class TestDynamoDBCertRecordStoreFactory extends DynamoDBCertRecordStoreFactory { @Override - AmazonDynamoDB getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { + DynamoDbClient getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { return dbClient; } } @BeforeMethod public void setUp() { - MockitoAnnotations.initMocks(this); - Mockito.doReturn(table).when(dynamoDB).getTable("Athenz-ZTS-Table"); + MockitoAnnotations.openMocks(this); } @Test @@ -188,7 +183,7 @@ public void testGetDynamoDBClient() { DynamoDBCertRecordStoreFactory factory = new DynamoDBCertRecordStoreFactory(); ZTSClientNotificationSenderImpl ztsClientNotificationSender = Mockito.mock(ZTSClientNotificationSenderImpl.class); PrivateKeyStore privateKeyStore = Mockito.mock(PrivateKeyStore.class); - AmazonDynamoDB dynamoDBClient = factory.getDynamoDBClient(ztsClientNotificationSender, privateKeyStore); + DynamoDbClient dynamoDBClient = factory.getDynamoDBClient(ztsClientNotificationSender, privateKeyStore); assertNotNull(dynamoDBClient); System.clearProperty(ZTS_PROP_DYNAMODB_KEY_PATH); diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreTest.java index 06822ae6aa8..f5741de6059 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBCertRecordStoreTest.java @@ -15,8 +15,6 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; - import com.yahoo.athenz.auth.Principal; import com.yahoo.athenz.auth.impl.SimplePrincipal; import com.yahoo.athenz.auth.util.Crypto; @@ -29,6 +27,7 @@ import org.mockito.MockitoAnnotations; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import java.io.File; import java.io.IOException; @@ -41,11 +40,11 @@ public class DynamoDBCertRecordStoreTest { - @Mock private AmazonDynamoDB dbClient; + @Mock private DynamoDbClient dbClient; @BeforeMethod public void setUp() { - MockitoAnnotations.initMocks(this); + MockitoAnnotations.openMocks(this); } @Test @@ -66,33 +65,6 @@ public void testGetConnection() { store.clearConnections(); } - @Test - public void testGetConnectionException() { - - // passing null for table name to get exception - DynamoDBCertRecordStore store = new DynamoDBCertRecordStore( - dbClient, - null, - "Athenz-ZTS-Current-Time-Index", - "Athenz-ZTS-Host-Name-Index", - null); - - try { - store.getConnection(); - fail(); - } catch (Exception ignored) { - } - - // passing null for index name to get exception - store = new DynamoDBCertRecordStore(dbClient, "Athenz-ZTS-Table", null, null, null); - - try { - store.getConnection(); - fail(); - } catch (Exception ignored) { - } - } - @Test public void testLog() { diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelperTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelperTest.java index a575f116468..3ded7a77081 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelperTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBNotificationsHelperTest.java @@ -1,16 +1,30 @@ +/* + * Copyright The Athenz Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.document.Item; -import com.amazonaws.services.dynamodbv2.document.ItemUtils; -import com.amazonaws.services.dynamodbv2.document.Table; -import com.amazonaws.services.dynamodbv2.document.UpdateItemOutcome; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; -import com.amazonaws.services.dynamodbv2.model.AttributeValue; -import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughputExceededException; import com.yahoo.athenz.zts.ZTSConsts; import com.yahoo.athenz.zts.ZTSTestUtils; import org.mockito.Mockito; import org.testng.annotations.Test; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughputExceededException; +import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest; +import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse; import java.util.*; import java.util.concurrent.TimeUnit; @@ -18,7 +32,6 @@ import static org.testng.Assert.*; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.when; public class DynamoDBNotificationsHelperTest { @@ -65,18 +78,19 @@ public void testIsMostUpdatedHostRecord() { null, "testHost1"); - Item itemNoCurrentTime = ItemUtils.toItem(noCurrentTime); - Item itemEmptyCurrentTime = ItemUtils.toItem(emptyCurrentTime); - Item itemFiveDaysAgo = ItemUtils.toItem(fiveDaysAgoMap); - Item itemThreeDaysAgo = ItemUtils.toItem(threeDaysAgoMap); - List allItems = Arrays.asList(itemEmptyCurrentTime, itemFiveDaysAgo, itemThreeDaysAgo, itemNoCurrentTime); + List> allItems = Arrays.asList(emptyCurrentTime, fiveDaysAgoMap, + threeDaysAgoMap, noCurrentTime); DynamoDBNotificationsHelper dynamoDBNotificationsHelper = new DynamoDBNotificationsHelper(); - assertFalse(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(itemFiveDaysAgo, allItems, "currentTime", "primaryKey")); - assertFalse(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(itemNoCurrentTime, allItems, "currentTime", "primaryKey")); - assertFalse(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(itemEmptyCurrentTime, allItems, "currentTime", "primaryKey")); - assertTrue(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(itemThreeDaysAgo, allItems, "currentTime", "primaryKey")); + assertFalse(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(fiveDaysAgoMap, + allItems, "currentTime", "primaryKey")); + assertFalse(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(noCurrentTime, + allItems, "currentTime", "primaryKey")); + assertFalse(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(emptyCurrentTime, + allItems, "currentTime", "primaryKey")); + assertTrue(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(threeDaysAgoMap, + allItems, "currentTime", "primaryKey")); } @Test @@ -94,11 +108,11 @@ public void testIsMostUpdatedHostRecordSingleRecord() { null, "testHost1"); - Item itemThreeDaysAgo = ItemUtils.toItem(threeDaysAgoMap); - List allItems = Collections.singletonList(itemThreeDaysAgo); + List> allItems = Collections.singletonList(threeDaysAgoMap); DynamoDBNotificationsHelper dynamoDBNotificationsHelper = new DynamoDBNotificationsHelper(); - assertTrue(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(itemThreeDaysAgo, allItems, "currentTime", "primaryKey")); + assertTrue(dynamoDBNotificationsHelper.isMostUpdatedRecordBasedOnAttribute(threeDaysAgoMap, + allItems, "currentTime", "primaryKey")); } @Test @@ -119,15 +133,14 @@ public void testUpdateLastNotifiedItem() throws TimeoutException, InterruptedExc null, "testHost2"); - Item item = ItemUtils.toItem(reNotified); - - UpdateItemOutcome updateItemOutcome1 = Mockito.mock(UpdateItemOutcome.class); - when(updateItemOutcome1.getItem()).thenReturn(item); - Table table = Mockito.mock(Table.class); - Mockito.when(table.updateItem(any(UpdateItemSpec.class))).thenReturn(updateItemOutcome1); - Item updatedItem = dynamoDBNotificationsHelper.updateLastNotifiedItem("lastNotifiedServer", lastNotifiedTime, yesterday, item, "primaryKey", table); + DynamoDbClient dynamoDbClient = Mockito.mock(DynamoDbClient.class); + UpdateItemResponse response = UpdateItemResponse.builder().attributes(reNotified).build(); + Mockito.when(dynamoDbClient.updateItem(any(UpdateItemRequest.class))).thenReturn(response); + Map updatedItem = dynamoDBNotificationsHelper.updateLastNotifiedItem( + "lastNotifiedServer", lastNotifiedTime, yesterday, reNotified, "primaryKey", + "tableName", dynamoDbClient); - assertEquals(updatedItem, item); + assertEquals(updatedItem, reNotified); } @Test @@ -148,15 +161,17 @@ public void testUpdateLastNotifiedItemRetry() throws TimeoutException, Interrupt null, "testHost2"); - Item item = ItemUtils.toItem(reNotified); + DynamoDbClient dynamoDbClient = Mockito.mock(DynamoDbClient.class); + UpdateItemResponse response = UpdateItemResponse.builder().attributes(reNotified).build(); + Mockito.when(dynamoDbClient.updateItem(any(UpdateItemRequest.class))) + .thenThrow(ProvisionedThroughputExceededException.builder().build()) + .thenReturn(response); - UpdateItemOutcome updateItemOutcome1 = Mockito.mock(UpdateItemOutcome.class); - when(updateItemOutcome1.getItem()).thenReturn(item); - Table table = Mockito.mock(Table.class); - Mockito.when(table.updateItem(any(UpdateItemSpec.class))).thenThrow(new ProvisionedThroughputExceededException("Provisioned Throughput Exceeded")).thenReturn(updateItemOutcome1); - Item updatedItem = dynamoDBNotificationsHelper.updateLastNotifiedItem("lastNotifiedServer", lastNotifiedTime, yesterday, item, "primaryKey", table); + Map updatedItem = dynamoDBNotificationsHelper.updateLastNotifiedItem( + "lastNotifiedServer", lastNotifiedTime, yesterday, reNotified, "primaryKey", + "tableName", dynamoDbClient); - assertEquals(updatedItem, item); + assertEquals(updatedItem, reNotified); } @Test @@ -179,18 +194,13 @@ public void testUpdateLastNotifiedItemRetryFailed() throws InterruptedException null, "testHost2"); - Item item = ItemUtils.toItem(reNotified); - - UpdateItemOutcome updateItemOutcome1 = Mockito.mock(UpdateItemOutcome.class); - when(updateItemOutcome1.getItem()).thenReturn(item); - Table table = Mockito.mock(Table.class); - // After getting this error twice, we stop retrying - Mockito.when(table.updateItem(any(UpdateItemSpec.class))) - .thenThrow(new ProvisionedThroughputExceededException("Provisioned Throughput Exceeded")) - .thenThrow(new ProvisionedThroughputExceededException("Provisioned Throughput Exceeded")); + DynamoDbClient dynamoDbClient = Mockito.mock(DynamoDbClient.class); + Mockito.when(dynamoDbClient.updateItem(any(UpdateItemRequest.class))) + .thenThrow(ProvisionedThroughputExceededException.builder().build()); try { - dynamoDBNotificationsHelper.updateLastNotifiedItem("lastNotifiedServer", lastNotifiedTime, yesterday, item, "primaryKey", table); + dynamoDBNotificationsHelper.updateLastNotifiedItem("lastNotifiedServer", lastNotifiedTime, yesterday, + reNotified, "primaryKey", "tableName", dynamoDbClient); fail(); } catch (TimeoutException ex) { assertEquals("Failed too many retries. Check table provisioned throughput settings.", ex.getMessage()); @@ -198,6 +208,5 @@ public void testUpdateLastNotifiedItemRetryFailed() throws InterruptedException System.clearProperty(ZTSConsts.ZTS_PROP_CERT_DYNAMODB_RETRIES); System.clearProperty(ZTSConsts.ZTS_PROP_CERT_DYNAMODB_RETRIES_SLEEP_MILLIS); - } } diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnectionTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnectionTest.java index cdb5e8faf60..107a57a4286 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnectionTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreConnectionTest.java @@ -15,17 +15,21 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.document.*; -import com.amazonaws.services.dynamodbv2.document.spec.DeleteItemSpec; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; -import com.amazonaws.services.dynamodbv2.model.AmazonDynamoDBException; import com.yahoo.athenz.common.server.ssh.SSHCertRecord; +import com.yahoo.athenz.zts.utils.DynamoDBUtils; import org.mockito.ArgumentMatchers; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.MockitoAnnotations; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import software.amazon.awssdk.awscore.exception.AwsServiceException; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.*; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; import static org.testng.Assert.*; @@ -33,29 +37,36 @@ public class DynamoDBSSHRecordStoreConnectionTest { private final String tableName = "cert-table"; - @Mock private DynamoDB dynamoDB; - @Mock private Table table; - @Mock private Item item; - @Mock private PutItemOutcome putOutcome; - @Mock private DeleteItemOutcome deleteOutcome; - @Mock private UpdateItemOutcome updateOutcome; + @Mock private DynamoDbClient dynamoDB; + @Mock private PutItemResponse putOutcome = Mockito.mock(PutItemResponse.class); + @Mock private DeleteItemResponse deleteOutcome = Mockito.mock(DeleteItemResponse.class); + @Mock private UpdateItemResponse updateOutcome = Mockito.mock(UpdateItemResponse.class); @BeforeMethod public void setUp() { - MockitoAnnotations.initMocks(this); - Mockito.doReturn(table).when(dynamoDB).getTable(tableName); + MockitoAnnotations.openMocks(this); } @Test public void testGetSSHCertRecord() { - Mockito.doReturn(item).when(table).getItem("primaryKey", "cn:1234"); + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("cn:1234")); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); - Mockito.doReturn("1234").when(item).getString("instanceId"); - Mockito.doReturn("cn").when(item).getString("service"); - Mockito.doReturn("host1,host2").when(item).getString("principals"); - Mockito.doReturn("10.10.10.11").when(item).getString("clientIP"); - Mockito.doReturn("10.10.10.12").when(item).getString("privateIP"); + Map attrs = new HashMap<>(); + attrs.put("instanceId", AttributeValue.fromS("1234")); + attrs.put("service", AttributeValue.fromS("cn")); + attrs.put("principals", AttributeValue.fromS("host1,host2")); + attrs.put("clientIP", AttributeValue.fromS("10.10.10.11")); + attrs.put("privateIP", AttributeValue.fromS("10.10.10.12")); + + GetItemResponse response = GetItemResponse.builder().item(attrs).build(); + Mockito.doReturn(response).when(dynamoDB).getItem(request); DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); dbConn.setOperationTimeout(10); @@ -73,19 +84,40 @@ public void testGetSSHCertRecord() { @Test public void testGetSSHCertRecordNotFoundNull() { - Mockito.doReturn(null).when(table).getItem("primaryKey", "cn:1234"); + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("cn:1234")); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); + + GetItemResponse response1 = GetItemResponse.builder().item(null).build(); + GetItemResponse response2 = GetItemResponse.builder().item(Collections.emptyMap()).build(); + Mockito.when(dynamoDB.getItem(request)).thenReturn(response1).thenReturn(response2); DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); + // first time we should get null item SSHCertRecord certRecord = dbConn.getSSHCertRecord("1234", "cn"); assertNull(certRecord); + // second time we should get empty map + certRecord = dbConn.getSSHCertRecord("1234", "cn"); + assertNull(certRecord); dbConn.close(); } @Test public void testGetSSHCertRecordNotFoundException() { - Mockito.doThrow(new AmazonDynamoDBException("item not found")) - .when(table).getItem("primaryKey", "cn:1234"); + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("cn:1234")); + + GetItemRequest request = GetItemRequest.builder() + .key(keyToGet) + .tableName(tableName) + .build(); + + Mockito.when(dynamoDB.getItem(request)).thenThrow(ResourceNotFoundException.builder().build()); DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); SSHCertRecord certRecord = dbConn.getSSHCertRecord("1234", "cn"); @@ -105,15 +137,20 @@ public void testInsertSSHRecord() { certRecord.setClientIP("10.10.10.11"); certRecord.setPrivateIP("10.10.10.12"); - Item item = new Item() - .withPrimaryKey("primaryKey", "cn:1234") - .withString("instanceId", certRecord.getInstanceId()) - .withString("service", certRecord.getService()) - .withString("principals", certRecord.getPrincipals()) - .withString("clientIP", certRecord.getClientIP()) - .withString("privateIP", certRecord.getPrivateIP()); + HashMap itemValues = new HashMap<>(); + itemValues.put("primaryKey", AttributeValue.fromS("cn:1234")); + itemValues.put("instanceId", AttributeValue.fromS(certRecord.getInstanceId())); + itemValues.put("service", AttributeValue.fromS(certRecord.getService())); + itemValues.put("clientIP", AttributeValue.fromS(certRecord.getClientIP())); + itemValues.put("principals", AttributeValue.fromS(certRecord.getPrincipals())); + itemValues.put("privateIP", AttributeValue.fromS(certRecord.getPrivateIP())); + + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(itemValues) + .build(); - Mockito.doReturn(putOutcome).when(table).putItem(item); + Mockito.doReturn(putOutcome).when(dynamoDB).putItem(request); boolean requestSuccess = dbConn.insertSSHCertRecord(certRecord); assertTrue(requestSuccess); @@ -125,8 +162,8 @@ public void testInsertSSHRecordException() { SSHCertRecord certRecord = new SSHCertRecord(); - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).putItem(ArgumentMatchers.any(Item.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).putItem(ArgumentMatchers.any(PutItemRequest.class)); DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); boolean requestSuccess = dbConn.insertSSHCertRecord(certRecord); @@ -147,16 +184,23 @@ public void testUpdateSSHRecord() { certRecord.setClientIP("10.10.10.11"); certRecord.setPrivateIP("10.10.10.12"); - UpdateItemSpec item = new UpdateItemSpec() - .withPrimaryKey("primaryKey", "cn:1234") - .withAttributeUpdate( - new AttributeUpdate("instanceId").put(certRecord.getInstanceId()), - new AttributeUpdate("service").put(certRecord.getService()), - new AttributeUpdate("principals").put(certRecord.getPrincipals()), - new AttributeUpdate("clientIP").put(certRecord.getClientIP()), - new AttributeUpdate("privateIP").put(certRecord.getPrivateIP())); + HashMap updatedValues = new HashMap<>(); + DynamoDBUtils.updateItemStringValue(updatedValues, "instanceId", certRecord.getInstanceId()); + DynamoDBUtils.updateItemStringValue(updatedValues, "service", certRecord.getService()); + DynamoDBUtils.updateItemStringValue(updatedValues, "clientIP", certRecord.getClientIP()); + DynamoDBUtils.updateItemStringValue(updatedValues, "principals", certRecord.getPrincipals()); + DynamoDBUtils.updateItemStringValue(updatedValues, "privateIP", certRecord.getPrivateIP()); + + HashMap itemKey = new HashMap<>(); + itemKey.put("primaryKey", AttributeValue.fromS("cn:1234")); + + UpdateItemRequest request = UpdateItemRequest.builder() + .tableName(tableName) + .key(itemKey) + .attributeUpdates(updatedValues) + .build(); - Mockito.doReturn(updateOutcome).when(table).updateItem(item); + Mockito.doReturn(updateOutcome).when(dynamoDB).updateItem(request); boolean requestSuccess = dbConn.updateSSHCertRecord(certRecord); assertTrue(requestSuccess); @@ -168,8 +212,8 @@ public void testUpdateSSHRecordException() { SSHCertRecord certRecord = new SSHCertRecord(); - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).updateItem(ArgumentMatchers.any(UpdateItemSpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).updateItem(ArgumentMatchers.any(UpdateItemRequest.class)); DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); boolean requestSuccess = dbConn.updateSSHCertRecord(certRecord); @@ -180,9 +224,16 @@ public void testUpdateSSHRecordException() { @Test public void testDeleteSSHRecord() { - DeleteItemSpec deleteItemSpec = new DeleteItemSpec() - .withPrimaryKey("primaryKey", "cn:1234"); - Mockito.doReturn(deleteOutcome).when(table).deleteItem(deleteItemSpec); + + HashMap keyToGet = new HashMap<>(); + keyToGet.put("primaryKey", AttributeValue.fromS("cn:1234")); + + DeleteItemRequest request = DeleteItemRequest.builder() + .tableName(tableName) + .key(keyToGet) + .build(); + + Mockito.doReturn(deleteOutcome).when(dynamoDB).deleteItem(request); DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); @@ -194,8 +245,8 @@ public void testDeleteSSHRecord() { @Test public void testDeleteSSHRecordException() { - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).deleteItem(ArgumentMatchers.any(DeleteItemSpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).deleteItem(ArgumentMatchers.any(DeleteItemRequest.class)); DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); @@ -205,7 +256,7 @@ public void testDeleteSSHRecordException() { } @Test - public void testdeleteExpiredSSHCertRecords() { + public void testDeleteExpiredSSHCertRecords() { DynamoDBSSHRecordStoreConnection dbConn = new DynamoDBSSHRecordStoreConnection(dynamoDB, tableName); assertEquals(0, dbConn.deleteExpiredSSHCertRecords(100)); assertEquals(0, dbConn.deleteExpiredSSHCertRecords(100000)); diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactoryTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactoryTest.java index 39059039bbb..76c4a61b833 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactoryTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreFactoryTest.java @@ -15,9 +15,6 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; -import com.amazonaws.services.dynamodbv2.document.Table; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.ssh.SSHRecordStore; import com.yahoo.athenz.zts.ResourceException; @@ -28,6 +25,7 @@ import org.mockito.MockitoAnnotations; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import static com.yahoo.athenz.zts.ZTSConsts.*; import static org.mockito.Mockito.when; @@ -35,22 +33,18 @@ public class DynamoDBSSHRecordStoreFactoryTest { - @Mock private AmazonDynamoDB dbClient; - @Mock private Table table; - - @Mock private DynamoDB dynamoDB; + @Mock private DynamoDbClient dbClient; class TestDynamoDBSSHRecordStoreFactory extends DynamoDBSSHRecordStoreFactory { @Override - AmazonDynamoDB getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { + DynamoDbClient getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { return dbClient; } } @BeforeMethod public void setUp() { - MockitoAnnotations.initMocks(this); - Mockito.doReturn(table).when(dynamoDB).getTable("Athenz-ZTS-Table"); + MockitoAnnotations.openMocks(this); } @Test @@ -121,7 +115,7 @@ public void testGetDynamoDBClient() { DynamoDBSSHRecordStoreFactory factory = new DynamoDBSSHRecordStoreFactory(); ZTSClientNotificationSenderImpl ztsClientNotificationSender = Mockito.mock(ZTSClientNotificationSenderImpl.class); PrivateKeyStore privateKeyStore = Mockito.mock(PrivateKeyStore.class); - AmazonDynamoDB dynamoDBClient = factory.getDynamoDBClient(ztsClientNotificationSender, privateKeyStore); + DynamoDbClient dynamoDBClient = factory.getDynamoDBClient(ztsClientNotificationSender, privateKeyStore); assertNotNull(dynamoDBClient); System.clearProperty(ZTS_PROP_DYNAMODB_KEY_PATH); diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreTest.java index 3791cdf2ecf..4585809b0db 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBSSHRecordStoreTest.java @@ -15,7 +15,6 @@ */ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.yahoo.athenz.auth.Principal; import com.yahoo.athenz.auth.impl.SimplePrincipal; import com.yahoo.athenz.common.server.db.RolesProvider; @@ -27,19 +26,19 @@ import org.mockito.MockitoAnnotations; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.fail; import static org.testng.AssertJUnit.assertFalse; import static org.testng.AssertJUnit.assertTrue; public class DynamoDBSSHRecordStoreTest { - @Mock private AmazonDynamoDB dbClient; + @Mock private DynamoDbClient dbClient; @BeforeMethod public void setUp() { - MockitoAnnotations.initMocks(this); + MockitoAnnotations.openMocks(this); } @Test @@ -55,19 +54,6 @@ public void testGetConnection() { store.clearConnections(); } - @Test - public void testGetConnectionException() { - - // passing null for table name to get exception - DynamoDBSSHRecordStore store = new DynamoDBSSHRecordStore(dbClient, null, null); - - try { - store.getConnection(); - fail(); - } catch (Exception ignored) { - } - } - @Test public void testLog() { diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerTest.java index 2b9d604f904..77c0006fed1 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/cert/impl/DynamoDBStatusCheckerTest.java @@ -16,19 +16,21 @@ package com.yahoo.athenz.zts.cert.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.model.ListTablesResult; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.status.StatusCheckException; import com.yahoo.athenz.db.dynamodb.DynamoDBClientAndCredentials; import com.yahoo.athenz.db.dynamodb.DynamoDBClientFetcher; import com.yahoo.athenz.db.dynamodb.DynamoDBClientSettings; -import com.yahoo.athenz.zts.AWSCredentialsProviderImpl; +import com.yahoo.athenz.zts.AWSCredentialsProviderImplV2; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.MockitoAnnotations; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.ListTablesRequest; +import software.amazon.awssdk.services.dynamodb.model.ListTablesResponse; import java.io.IOException; import java.util.Collections; @@ -40,9 +42,9 @@ public class DynamoDBStatusCheckerTest { @Mock private DynamoDBClientFetcher dynamoDBClientFetcher; - @Mock private AmazonDynamoDB amazonDynamoDB; - @Mock private AWSCredentialsProviderImpl awsCredentialsProvider; - @Mock private ListTablesResult listTablesResult; + @Mock private DynamoDbClient amazonDynamoDB; + @Mock private DynamoDbAsyncClient amazonDynamoAsyncDB; + @Mock private AWSCredentialsProviderImplV2 awsCredentialsProvider; @Mock private PrivateKeyStore keyStore; public class DynamoDBStatusCheckerTestClass extends DynamoDBStatusChecker { @@ -59,21 +61,25 @@ DynamoDBClientFetcher getDynamoDBClientFetcher() { @BeforeMethod public void setUp() { - MockitoAnnotations.initMocks(this); + MockitoAnnotations.openMocks(this); } @Test public void testCheck() throws StatusCheckException, IOException { // Mock getting client and credentials successfully and table exists String tableName = "testTable"; - when(listTablesResult.getTableNames()).thenReturn(Collections.singletonList(tableName)); - when(amazonDynamoDB.listTables()).thenReturn(listTablesResult); - DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, null, awsCredentialsProvider); - when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))).thenReturn(dynamoDBClientAndCredentials); + + ListTablesResponse response = Mockito.mock(ListTablesResponse.class); + Mockito.when(amazonDynamoDB.listTables((ListTablesRequest) any())).thenReturn(response); + when(response.tableNames()).thenReturn(Collections.singletonList(tableName)); + DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, + amazonDynamoAsyncDB, awsCredentialsProvider); + when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))) + .thenReturn(dynamoDBClientAndCredentials); DynamoDBStatusCheckerTestClass dynamoDBStatusChecker = new DynamoDBStatusCheckerTestClass(tableName); dynamoDBStatusChecker.check(); - Mockito.verify(amazonDynamoDB, times(1)).shutdown(); + Mockito.verify(amazonDynamoDB, times(1)).close(); Mockito.verify(awsCredentialsProvider, times(1)).close(); } @@ -81,34 +87,20 @@ public void testCheck() throws StatusCheckException, IOException { public void testCheckNoCredentialsProvider() throws StatusCheckException, IOException { // Mock getting client and credentials successfully and table exists String tableName = "testTable"; - when(listTablesResult.getTableNames()).thenReturn(Collections.singletonList(tableName)); - when(amazonDynamoDB.listTables()).thenReturn(listTablesResult); - DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, null, null); - when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))).thenReturn(dynamoDBClientAndCredentials); - DynamoDBStatusCheckerTestClass dynamoDBStatusChecker = new DynamoDBStatusCheckerTestClass(tableName); - dynamoDBStatusChecker.check(); - Mockito.verify(amazonDynamoDB, times(1)).shutdown(); - Mockito.verify(awsCredentialsProvider, times(0)).close(); - } + ListTablesResponse response = Mockito.mock(ListTablesResponse.class); + Mockito.when(amazonDynamoDB.listTables((ListTablesRequest) any())).thenReturn(response); + when(response.tableNames()).thenReturn(Collections.singletonList(tableName)); - @Test - public void testCheckNoDynamoDBClient() throws StatusCheckException, IOException { - // Mock getting client and credentials successfully and table exists - String tableName = "testTable"; - when(listTablesResult.getTableNames()).thenReturn(Collections.singletonList(tableName)); - DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(null, null, awsCredentialsProvider); - when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))).thenReturn(dynamoDBClientAndCredentials); + DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, + amazonDynamoAsyncDB, null); + when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))) + .thenReturn(dynamoDBClientAndCredentials); DynamoDBStatusCheckerTestClass dynamoDBStatusChecker = new DynamoDBStatusCheckerTestClass(tableName); - try { - dynamoDBStatusChecker.check(); - fail(); - } catch (StatusCheckException ex) { - assertEquals(500, ex.getCode()); - Mockito.verify(amazonDynamoDB, times(0)).shutdown(); - Mockito.verify(awsCredentialsProvider, times(1)).close(); - } + dynamoDBStatusChecker.check(); + Mockito.verify(amazonDynamoDB, times(1)).close(); + Mockito.verify(awsCredentialsProvider, times(0)).close(); } @Test @@ -116,10 +108,15 @@ public void testTableNotFound() throws IOException { String requestedTable = "requestedTable"; String tableNameInAws = "someExistingTable"; // Mock getting client and credentials successfully but table doesn't exist - when(listTablesResult.getTableNames()).thenReturn(Collections.singletonList(tableNameInAws)); - when(amazonDynamoDB.listTables()).thenReturn(listTablesResult); - DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, null, awsCredentialsProvider); - when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))).thenReturn(dynamoDBClientAndCredentials); + + ListTablesResponse response = Mockito.mock(ListTablesResponse.class); + Mockito.when(amazonDynamoDB.listTables((ListTablesRequest) any())).thenReturn(response); + when(response.tableNames()).thenReturn(Collections.singletonList(tableNameInAws)); + + DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, + amazonDynamoAsyncDB, awsCredentialsProvider); + when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))) + .thenReturn(dynamoDBClientAndCredentials); DynamoDBStatusCheckerTestClass dynamoDBStatusChecker = new DynamoDBStatusCheckerTestClass(requestedTable); try { @@ -129,7 +126,7 @@ public void testTableNotFound() throws IOException { assertEquals("Table named " + requestedTable + " wasn't found in DynamoDB", ex.getMsg()); assertEquals(200, ex.getCode()); } - Mockito.verify(amazonDynamoDB, times(1)).shutdown(); + Mockito.verify(amazonDynamoDB, times(1)).close(); Mockito.verify(awsCredentialsProvider, times(1)).close(); } @@ -137,10 +134,12 @@ public void testTableNotFound() throws IOException { public void testClientNullTables() throws IOException { // Mock getting client and credentials successfully but client returns null instead of tables String tableName = "testTable"; - when(amazonDynamoDB.listTables()).thenReturn(null); + Mockito.when(amazonDynamoDB.listTables((ListTablesRequest) any())).thenReturn(null); - DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, null, awsCredentialsProvider); - when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))).thenReturn(dynamoDBClientAndCredentials); + DynamoDBClientAndCredentials dynamoDBClientAndCredentials = new DynamoDBClientAndCredentials(amazonDynamoDB, + amazonDynamoAsyncDB, awsCredentialsProvider); + when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))) + .thenReturn(dynamoDBClientAndCredentials); DynamoDBStatusCheckerTestClass dynamoDBStatusChecker = new DynamoDBStatusCheckerTestClass(tableName); try { @@ -151,29 +150,10 @@ public void testClientNullTables() throws IOException { assertEquals(500, ex.getCode()); } - Mockito.verify(amazonDynamoDB, times(1)).shutdown(); + Mockito.verify(amazonDynamoDB, times(1)).close(); Mockito.verify(awsCredentialsProvider, times(1)).close(); } - @Test - public void testClientNull() throws IOException { - String tableName = "testTable"; - // Mock getting client failed (returns null) - when(dynamoDBClientFetcher.getDynamoDBClient(any(), any(DynamoDBClientSettings.class))).thenReturn(null); - DynamoDBStatusCheckerTestClass dynamoDBStatusChecker = new DynamoDBStatusCheckerTestClass(tableName); - - try { - dynamoDBStatusChecker.check(); - fail(); - } catch (StatusCheckException ex) { - assertNull(ex.getMessage()); - assertEquals(500, ex.getCode()); - } - - Mockito.verify(amazonDynamoDB, times(0)).shutdown(); - Mockito.verify(awsCredentialsProvider, times(0)).close(); - } - @Test public void testGetDynamoDBClientFetcher() { DynamoDBStatusChecker dynamoDBStatusChecker = new DynamoDBStatusChecker("testTable", keyStore); diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/store/CloudStoreTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/store/CloudStoreTest.java index 33510a530c0..8c0673a0d6b 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/store/CloudStoreTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/store/CloudStoreTest.java @@ -16,7 +16,6 @@ package com.yahoo.athenz.zts.store; -import static com.yahoo.athenz.common.ServerCommonConsts.ZTS_PROP_AWS_PUBLIC_CERT; import static com.yahoo.athenz.common.ServerCommonConsts.ZTS_PROP_AWS_REGION_NAME; import static org.testng.Assert.*; @@ -33,13 +32,13 @@ import org.mockito.Mockito; import org.testng.annotations.Test; -import com.amazonaws.auth.BasicSessionCredentials; -import com.amazonaws.services.securitytoken.model.AssumeRoleRequest; -import com.amazonaws.services.securitytoken.model.AssumeRoleResult; -import com.amazonaws.services.securitytoken.model.Credentials; import com.yahoo.athenz.zts.AWSTemporaryCredentials; import com.yahoo.athenz.zts.ResourceException; import com.yahoo.athenz.zts.ZTSConsts; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; +import software.amazon.awssdk.services.sts.model.Credentials; public class CloudStoreTest { @@ -67,51 +66,10 @@ public class CloudStoreTest { + "\"InstanceProfileId\" : \"AIPAJAVNLUGEWFWTIDPRA\"\n" + "}"; - @Test - public void testGetS3ClientNullCreds() { - CloudStore store = new CloudStore(); - store.awsEnabled = true; - store.credentials = null; - try { - store.getS3Client(); - fail(); - } catch (ResourceException ex) { - assertEquals(500, ex.getCode()); - } - store.close(); - } - - @Test - public void testGetS3ClientAWSNotEnabled() { - CloudStore store = new CloudStore(); - store.credentials = null; - try { - store.getS3Client(); - fail(); - } catch (ResourceException ex) { - assertEquals(500, ex.getCode()); - } - store.close(); - } - - @Test - public void testGetS3Client() { - - System.setProperty(ZTS_PROP_AWS_PUBLIC_CERT, "src/test/resources/aws_public.crt"); - CloudStore store = new CloudStore(); - store.credentials = new BasicSessionCredentials("accessKey", "secretKey", "token"); - store.awsEnabled = true; - store.awsRegion = "us-west-2"; - assertNotNull(store.getS3Client()); - - assertNotNull(store.getS3Client()); - store.close(); - } - @Test public void testGetTokenServiceClient() { CloudStore store = new CloudStore(); - store.credentials = new BasicSessionCredentials("accessKey", "secretKey", "token"); + store.credentials = AwsBasicCredentials.builder().accessKeyId("accessKey").secretAccessKey("secretKey").build(); store.awsEnabled = true; store.awsRegion = "us-west-2"; assertNotNull(store.getTokenServiceClient()); @@ -168,44 +126,44 @@ public void testGetAssumeRoleRequest() { CloudStore store = new CloudStore(); AssumeRoleRequest req = store.getAssumeRoleRequest("1234", "admin", null, null, "athenz.api"); - assertEquals("arn:aws:iam::1234:role/admin", req.getRoleArn()); - assertEquals("athenz.api", req.getRoleSessionName()); - assertNull(req.getDurationSeconds()); - assertNull(req.getExternalId()); + assertEquals("arn:aws:iam::1234:role/admin", req.roleArn()); + assertEquals("athenz.api", req.roleSessionName()); + assertNull(req.durationSeconds()); + assertNull(req.externalId()); req = store.getAssumeRoleRequest("12345", "adminuser", 101, "external", "athenz.api"); - assertEquals("arn:aws:iam::12345:role/adminuser", req.getRoleArn()); - assertEquals("athenz.api", req.getRoleSessionName()); - assertEquals(Integer.valueOf(101), req.getDurationSeconds()); - assertEquals("external", req.getExternalId()); + assertEquals("arn:aws:iam::12345:role/adminuser", req.roleArn()); + assertEquals("athenz.api", req.roleSessionName()); + assertEquals(Integer.valueOf(101), req.durationSeconds()); + assertEquals("external", req.externalId()); req = store.getAssumeRoleRequest("12345", "adminuser", 101, "external", "athenz.api-service"); - assertEquals("arn:aws:iam::12345:role/adminuser", req.getRoleArn()); - assertEquals("athenz.api-service", req.getRoleSessionName()); - assertEquals(Integer.valueOf(101), req.getDurationSeconds()); - assertEquals("external", req.getExternalId()); + assertEquals("arn:aws:iam::12345:role/adminuser", req.roleArn()); + assertEquals("athenz.api-service", req.roleSessionName()); + assertEquals(Integer.valueOf(101), req.durationSeconds()); + assertEquals("external", req.externalId()); req = store.getAssumeRoleRequest("12345", "adminuser", 101, "external", "athenz.api_service-test"); - assertEquals("arn:aws:iam::12345:role/adminuser", req.getRoleArn()); - assertEquals("athenz.api_service-test", req.getRoleSessionName()); - assertEquals(Integer.valueOf(101), req.getDurationSeconds()); - assertEquals("external", req.getExternalId()); + assertEquals("arn:aws:iam::12345:role/adminuser", req.roleArn()); + assertEquals("athenz.api_service-test", req.roleSessionName()); + assertEquals(Integer.valueOf(101), req.durationSeconds()); + assertEquals("external", req.externalId()); final String principalLongerThan64Chars = "athenz.environment.production.regions.us-west-2.services.zts-service"; req = store.getAssumeRoleRequest("12345", "adminuser", 101, "external", principalLongerThan64Chars); - assertEquals("arn:aws:iam::12345:role/adminuser", req.getRoleArn()); - assertEquals("athenz.environment.production....us-west-2.services.zts-service", req.getRoleSessionName()); - assertEquals(Integer.valueOf(101), req.getDurationSeconds()); - assertEquals("external", req.getExternalId()); + assertEquals("arn:aws:iam::12345:role/adminuser", req.roleArn()); + assertEquals("athenz.environment.production....us-west-2.services.zts-service", req.roleSessionName()); + assertEquals(Integer.valueOf(101), req.durationSeconds()); + assertEquals("external", req.externalId()); store.close(); System.setProperty(ZTSConsts.ZTS_PROP_AWS_ROLE_SESSION_NAME, "athenz-zts-service"); store = new CloudStore(); req = store.getAssumeRoleRequest("12345", "adminuser", 101, "external", "athenz.api-service"); - assertEquals("arn:aws:iam::12345:role/adminuser", req.getRoleArn()); - assertEquals("athenz-zts-service", req.getRoleSessionName()); - assertEquals(Integer.valueOf(101), req.getDurationSeconds()); - assertEquals("external", req.getExternalId()); + assertEquals("arn:aws:iam::12345:role/adminuser", req.roleArn()); + assertEquals("athenz-zts-service", req.roleSessionName()); + assertEquals(Integer.valueOf(101), req.durationSeconds()); + assertEquals("external", req.externalId()); store.close(); System.clearProperty(ZTSConsts.ZTS_PROP_AWS_ROLE_SESSION_NAME); } @@ -969,7 +927,7 @@ public void testInitializeAwsSupport() throws ExecutionException, TimeoutExcepti store.initializeAwsSupport(); // sleep a couple of seconds for the background thread to run - // before we try to shutting it down + // before we try to shut it down try { Thread.sleep(2000); @@ -997,13 +955,13 @@ public void testAssumeAWSRoleAWSNotEnabled() { public void testAssumeAWSRole() { MockCloudStore cloudStore = new MockCloudStore(); cloudStore.awsEnabled = true; - AssumeRoleResult mockResult = Mockito.mock(AssumeRoleResult.class); + AssumeRoleResponse mockResult = Mockito.mock(AssumeRoleResponse.class); Credentials creds = Mockito.mock(Credentials.class); - Mockito.when(creds.getAccessKeyId()).thenReturn("accesskeyid"); - Mockito.when(creds.getSecretAccessKey()).thenReturn("secretaccesskey"); - Mockito.when(creds.getSessionToken()).thenReturn("sessiontoken"); - Mockito.when(creds.getExpiration()).thenReturn(new Date()); - Mockito.when(mockResult.getCredentials()).thenReturn(creds); + Mockito.when(creds.accessKeyId()).thenReturn("accesskeyid"); + Mockito.when(creds.secretAccessKey()).thenReturn("secretaccesskey"); + Mockito.when(creds.sessionToken()).thenReturn("sessiontoken"); + Mockito.when(creds.expiration()).thenReturn(new Date().toInstant()); + Mockito.when(mockResult.credentials()).thenReturn(creds); cloudStore.setAssumeRoleResult(mockResult); cloudStore.setReturnSuperAWSRole(true); @@ -1020,13 +978,13 @@ public void testAssumeAWSRole() { public void testAssumeAWSRoleFailedCreds() { MockCloudStore cloudStore = new MockCloudStore(); cloudStore.awsEnabled = true; - AssumeRoleResult mockResult = Mockito.mock(AssumeRoleResult.class); + AssumeRoleResponse mockResult = Mockito.mock(AssumeRoleResponse.class); Credentials creds = Mockito.mock(Credentials.class); - Mockito.when(creds.getAccessKeyId()).thenReturn("accesskeyid"); - Mockito.when(creds.getSecretAccessKey()).thenReturn("secretaccesskey"); - Mockito.when(creds.getSessionToken()).thenReturn("sessiontoken"); - Mockito.when(creds.getExpiration()).thenReturn(new Date()); - Mockito.when(mockResult.getCredentials()).thenReturn(creds); + Mockito.when(creds.accessKeyId()).thenReturn("accesskeyid"); + Mockito.when(creds.secretAccessKey()).thenReturn("secretaccesskey"); + Mockito.when(creds.sessionToken()).thenReturn("sessiontoken"); + Mockito.when(creds.expiration()).thenReturn(new Date().toInstant()); + Mockito.when(mockResult.credentials()).thenReturn(creds); cloudStore.setAssumeRoleResult(mockResult); cloudStore.setReturnSuperAWSRole(true); @@ -1038,7 +996,7 @@ public void testAssumeAWSRoleFailedCreds() { errorMessage.setLength(0); assertNull(cloudStore.assumeAWSRole("account", "syncer", "athenz.syncer", null, null, errorMessage)); - // now set the timeout to 1 second and sleep that long and after + // now set the timeout to 1-second and sleep that long and after // that our test case should work as before cloudStore.invalidCacheTimeout = 1; @@ -1076,7 +1034,7 @@ public void testAssumeAWSRoleFailedCredsCache() { assertNull(cloudStore.awsInvalidCredsCache.get(cloudStore.getCacheKey("account", "syncer", "athenz.syncer", null, null))); // finally we're going to return access denied - 403 - // amazon exception and we should cache the failed creds + // amazon exception, and we should cache the failed creds cloudStore.setGetServiceException(403, true); errorMessage.setLength(0); @@ -1239,7 +1197,7 @@ public void testInvalidCacheCreds() { assertTrue(cloudStore.isFailedTempCredsRequest("cacheKey")); assertFalse(cloudStore.isFailedTempCredsRequest("unknown-key")); - // now set the timeout to only 1 second + // now set the timeout to only 1-second // and sleep so our records are expired cloudStore.invalidCacheTimeout = 1; diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/store/MockCloudStore.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/store/MockCloudStore.java index b5e0a998511..7adf1b78a7f 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/store/MockCloudStore.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/store/MockCloudStore.java @@ -16,22 +16,24 @@ package com.yahoo.athenz.zts.store; -import com.amazonaws.AmazonServiceException; -import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient; -import com.amazonaws.services.securitytoken.model.AssumeRoleRequest; -import com.amazonaws.services.securitytoken.model.AssumeRoleResult; -import com.amazonaws.services.securitytoken.model.GetCallerIdentityRequest; -import com.amazonaws.services.securitytoken.model.GetCallerIdentityResult; import com.yahoo.athenz.zts.AWSTemporaryCredentials; import org.mockito.Mockito; +import software.amazon.awssdk.awscore.exception.AwsErrorDetails; +import software.amazon.awssdk.http.SdkHttpResponse; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; +import software.amazon.awssdk.services.sts.model.GetCallerIdentityRequest; +import software.amazon.awssdk.services.sts.model.GetCallerIdentityResponse; +import software.amazon.awssdk.awscore.exception.AwsServiceException; public class MockCloudStore extends CloudStore { private String account = null; private String roleName = null; private String principal = null; private boolean returnSuperAWSRole = false; - private AssumeRoleResult assumeRoleResult = null; - private GetCallerIdentityResult callerIdentityResult = null; + private AssumeRoleResponse assumeRoleResult = null; + private GetCallerIdentityResponse callerIdentityResult = null; private int exceptionStatusCode = 0; private boolean amazonException = true; @@ -50,28 +52,31 @@ public void setMockFields(String account, String roleName, String principal) { this.principal = principal; } - void setAssumeRoleResult(AssumeRoleResult assumeRoleResult) { + void setAssumeRoleResult(AssumeRoleResponse assumeRoleResult) { this.assumeRoleResult = assumeRoleResult; } - void setGetCallerIdentityResult(GetCallerIdentityResult callerIdentityResult) { + void setGetCallerIdentityResult(GetCallerIdentityResponse callerIdentityResult) { this.callerIdentityResult = callerIdentityResult; } @Override - public AWSSecurityTokenServiceClient getTokenServiceClient() { + public StsClient getTokenServiceClient() { if (exceptionStatusCode != 0) { if (amazonException) { - AmazonServiceException ex = new AmazonServiceException("Error"); - ex.setStatusCode(exceptionStatusCode); - throw ex; + throw AwsServiceException.builder().awsErrorDetails( + AwsErrorDetails.builder().sdkHttpResponse( + SdkHttpResponse.builder().statusCode(exceptionStatusCode).build() + ).build() + ).build(); } else { throw new IllegalArgumentException("Error"); } } else { - AWSSecurityTokenServiceClient client = Mockito.mock(AWSSecurityTokenServiceClient.class); + StsClient client = Mockito.mock(StsClient.class); Mockito.when(client.assumeRole(Mockito.any(AssumeRoleRequest.class))).thenReturn(assumeRoleResult); - Mockito.when(client.getCallerIdentity(Mockito.any(GetCallerIdentityRequest.class))).thenReturn(callerIdentityResult); + Mockito.when(client.getCallerIdentity(Mockito.any(GetCallerIdentityRequest.class))) + .thenReturn(callerIdentityResult); return client; } } @@ -82,7 +87,7 @@ void setReturnSuperAWSRole(boolean returnSuperAWSRole) { @Override public AWSTemporaryCredentials assumeAWSRole(String account, String roleName, String principal, - Integer durationSeconds, String externalId, StringBuilder errorMessage) { + Integer durationSeconds, String externalId, StringBuilder errorMessage) { if (!returnSuperAWSRole) { AWSTemporaryCredentials tempCreds = null; diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnectionTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnectionTest.java index 4d36aed42ff..33232117225 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnectionTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreConnectionTest.java @@ -15,13 +15,9 @@ */ package com.yahoo.athenz.zts.workload.impl; -import com.amazonaws.services.dynamodbv2.document.*; -import com.amazonaws.services.dynamodbv2.document.internal.IteratorSupport; -import com.amazonaws.services.dynamodbv2.document.spec.QuerySpec; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; -import com.amazonaws.services.dynamodbv2.model.AmazonDynamoDBException; import com.yahoo.athenz.common.server.workload.WorkloadRecord; import com.yahoo.athenz.zts.ZTSTestUtils; +import com.yahoo.athenz.zts.utils.DynamoDBUtils; import org.mockito.ArgumentMatchers; import org.mockito.Mock; import org.mockito.Mockito; @@ -29,11 +25,14 @@ import org.testng.Assert; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import software.amazon.awssdk.awscore.exception.AwsServiceException; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.*; import java.util.Date; +import java.util.HashMap; import java.util.List; - -import static org.mockito.Mockito.when; +import java.util.Map; public class DynamoDBWorkloadRecordStoreConnectionTest { @@ -41,21 +40,13 @@ public class DynamoDBWorkloadRecordStoreConnectionTest { private final String serviceIndexName = "service-index"; private final String ipIndexName = "ip-index"; - @Mock private DynamoDB dynamoDB = Mockito.mock(DynamoDB.class); - @Mock private Table table = Mockito.mock(Table.class); - @Mock private Index serviceIndex = Mockito.mock(Index.class); - @Mock private Index ipIndex = Mockito.mock(Index.class); - @Mock private Item item = Mockito.mock(Item.class); - @Mock private PutItemOutcome putOutcome = Mockito.mock(PutItemOutcome.class); - @Mock private DeleteItemOutcome deleteOutcome = Mockito.mock(DeleteItemOutcome.class); - @Mock private UpdateItemOutcome updateOutcome = Mockito.mock(UpdateItemOutcome.class); + @Mock private DynamoDbClient dynamoDB; + @Mock private PutItemResponse putOutcome = Mockito.mock(PutItemResponse.class); + @Mock private UpdateItemResponse updateOutcome = Mockito.mock(UpdateItemResponse.class); @BeforeMethod public void setUp() { MockitoAnnotations.openMocks(this); - Mockito.doReturn(table).when(dynamoDB).getTable(tableName); - Mockito.doReturn(serviceIndex).when(table).getIndex(serviceIndexName); - Mockito.doReturn(ipIndex).when(table).getIndex(ipIndexName); } private DynamoDBWorkloadRecordStoreConnection getDBConnection() { @@ -66,25 +57,29 @@ private DynamoDBWorkloadRecordStoreConnection getDBConnection() { public void testGetWorkloadRecordsByService() { long currTime = System.currentTimeMillis(); - Mockito.doReturn("1234").when(item).getString("instanceId"); - Mockito.doReturn("openstack").when(item).getString("provider"); - Mockito.doReturn("10.10.10.11").when(item).getString("ip"); - Mockito.doReturn("test-host").when(item).getString("hostname"); - Mockito.doReturn(true).when(item).hasAttribute("hostname"); - Mockito.doReturn(currTime).when(item).get("creationTime"); - Mockito.doReturn(currTime).when(item).get("updateTime"); - Mockito.doReturn(currTime).when(item).get("certExpiryTime"); - Mockito.doReturn(currTime).when(item).getLong("creationTime"); - Mockito.doReturn(currTime).when(item).getLong("updateTime"); - Mockito.doReturn(currTime).when(item).getLong("certExpiryTime"); - Mockito.doReturn(true).when(item).hasAttribute("certExpiryTime"); - - ItemCollection itemCollection = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport = Mockito.mock(IteratorSupport.class); - when(itemCollection.iterator()).thenReturn(iteratorSupport); - when(iteratorSupport.hasNext()).thenReturn(true, false); - when(iteratorSupport.next()).thenReturn(item); - Mockito.doReturn(itemCollection).when(serviceIndex).query(Mockito.any(QuerySpec.class)); + Map attrs = new HashMap<>(); + attrs.put("instanceId", AttributeValue.fromS("1234")); + attrs.put("provider", AttributeValue.fromS("openstack")); + attrs.put("ip", AttributeValue.fromS("10.10.10.11")); + attrs.put("hostname", AttributeValue.fromS("test-host")); + attrs.put("creationTime", AttributeValue.fromN(String.valueOf(currTime))); + attrs.put("updateTime", AttributeValue.fromN(String.valueOf(currTime))); + attrs.put("certExpiryTime", AttributeValue.fromN(String.valueOf(currTime))); + + HashMap attrValues = new HashMap<>(); + attrValues.put(":v_service", AttributeValue.fromS("athenz.api")); + + QueryRequest request = QueryRequest.builder() + .tableName(tableName) + .indexName(serviceIndexName) + .keyConditionExpression("service = :v_service") + .expressionAttributeValues(attrValues) + .build(); + + List> items = new java.util.ArrayList<>(); + items.add(attrs); + QueryResponse response = QueryResponse.builder().items(items).build(); + Mockito.doReturn(response).when(dynamoDB).query(request); DynamoDBWorkloadRecordStoreConnection dbConn = getDBConnection(); dbConn.setOperationTimeout(10); @@ -104,9 +99,20 @@ public void testGetWorkloadRecordsByService() { @Test public void testGetWorkloadRecordsByServiceNotFoundNull() { - Mockito.doReturn(null).when(serviceIndex).query(Mockito.any(QuerySpec.class)); + HashMap attrValues = new HashMap<>(); + attrValues.put(":v_service", AttributeValue.fromS("athenz.api")); + + QueryRequest request = QueryRequest.builder() + .tableName(tableName) + .indexName(serviceIndexName) + .keyConditionExpression("service = :v_service") + .expressionAttributeValues(attrValues) + .build(); + + Mockito.doReturn(null).when(dynamoDB).query(request); - DynamoDBWorkloadRecordStoreConnection dbConn = new DynamoDBWorkloadRecordStoreConnection(dynamoDB, tableName, "service-index", "ip-index"); + DynamoDBWorkloadRecordStoreConnection dbConn = new DynamoDBWorkloadRecordStoreConnection(dynamoDB, + tableName, "service-index", "ip-index"); List wlRecordList = dbConn.getWorkloadRecordsByService("athenz", "api"); Assert.assertTrue(wlRecordList.isEmpty()); dbConn.close(); @@ -116,23 +122,27 @@ public void testGetWorkloadRecordsByServiceNotFoundNull() { public void testGetWorkloadRecordsByIp() { long currTime = System.currentTimeMillis(); - Mockito.doReturn("1234").when(item).getString("instanceId"); - Mockito.doReturn("openstack").when(item).getString("provider"); - Mockito.doReturn("athenz.api").when(item).getString("service"); - Mockito.doReturn("test-host").when(item).getString("hostname"); - Mockito.doReturn(currTime).when(item).get("creationTime"); - Mockito.doReturn(currTime).when(item).get("updateTime"); - Mockito.doReturn(currTime).when(item).get("certExpiryTime"); - Mockito.doReturn(currTime).when(item).getLong("creationTime"); - Mockito.doReturn(currTime).when(item).getLong("updateTime"); - Mockito.doReturn(currTime).when(item).getLong("certExpiryTime"); - - ItemCollection itemCollection = Mockito.mock(ItemCollection.class); - IteratorSupport iteratorSupport = Mockito.mock(IteratorSupport.class); - when(itemCollection.iterator()).thenReturn(iteratorSupport); - when(iteratorSupport.hasNext()).thenReturn(true, false); - when(iteratorSupport.next()).thenReturn(item); - Mockito.doReturn(itemCollection).when(ipIndex).query(Mockito.any(QuerySpec.class)); + Map attrs = new HashMap<>(); + attrs.put("instanceId", AttributeValue.fromS("1234")); + attrs.put("provider", AttributeValue.fromS("openstack")); + attrs.put("service", AttributeValue.fromS("athenz.api")); + attrs.put("creationTime", AttributeValue.fromN(String.valueOf(currTime))); + attrs.put("updateTime", AttributeValue.fromN(String.valueOf(currTime))); + + HashMap attrValues = new HashMap<>(); + attrValues.put(":v_ip", AttributeValue.fromS("10.0.0.1")); + + QueryRequest request = QueryRequest.builder() + .tableName(tableName) + .indexName(ipIndexName) + .keyConditionExpression("ip = :v_ip") + .expressionAttributeValues(attrValues) + .build(); + + List> items = new java.util.ArrayList<>(); + items.add(attrs); + QueryResponse response = QueryResponse.builder().items(items).build(); + Mockito.doReturn(response).when(dynamoDB).query(request); DynamoDBWorkloadRecordStoreConnection dbConn = getDBConnection(); dbConn.setOperationTimeout(10); @@ -152,7 +162,7 @@ public void testGetWorkloadRecordsByIp() { @Test public void testGetWorkloadRecordsByIpNotFoundNull() { - Mockito.doReturn(null).when(ipIndex).query(Mockito.any(QuerySpec.class)); + Mockito.doReturn(null).when(dynamoDB).query(ArgumentMatchers.any(QueryRequest.class)); DynamoDBWorkloadRecordStoreConnection dbConn = getDBConnection(); List wlRecordList = dbConn.getWorkloadRecordsByIp("10.0.0.1"); @@ -163,8 +173,8 @@ public void testGetWorkloadRecordsByIpNotFoundNull() { @Test public void testGetWorkloadRecordsByServiceNotFoundException() { - Mockito.doThrow(new AmazonDynamoDBException("item not found")) - .when(serviceIndex).query(Mockito.any(QuerySpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).query(ArgumentMatchers.any(QueryRequest.class)); DynamoDBWorkloadRecordStoreConnection dbConn = getDBConnection(); List wlRecordList = dbConn.getWorkloadRecordsByService("athenz", "api"); @@ -175,8 +185,8 @@ public void testGetWorkloadRecordsByServiceNotFoundException() { @Test public void testGetWorkloadRecordsByIpNotFoundException() { - Mockito.doThrow(new AmazonDynamoDBException("item not found")) - .when(ipIndex).query(Mockito.any(QuerySpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).query(ArgumentMatchers.any(QueryRequest.class)); DynamoDBWorkloadRecordStoreConnection dbConn = getDBConnection(); List wlRecordList = dbConn.getWorkloadRecordsByIp("10.0.0.1"); @@ -200,10 +210,15 @@ public void testInsertWorkloadRecord() { workloadRecord.setCreationTime(currDate); workloadRecord.setUpdateTime(currDate); - Item item = ItemUtils.toItem(ZTSTestUtils.generateWorkloadAttributeValues("athenz.api", "1234", "opensack", "10.0.0.1", "test-host.corp.yahoo.com", - Long.toString(currTime), Long.toString(currTime),Long.toString(currTime))); + PutItemRequest request = PutItemRequest.builder() + .tableName(tableName) + .item(ZTSTestUtils.generateWorkloadAttributeValues("athenz.api", "1234", "opensack", + "10.0.0.1", "test-host.corp.yahoo.com", Long.toString(currTime), Long.toString(currTime), + Long.toString(currTime))) + .build(); + + Mockito.doReturn(putOutcome).when(dynamoDB).putItem(request); - Mockito.doReturn(putOutcome).when(table).putItem(item); boolean requestSuccess = dbConn.insertWorkloadRecord(workloadRecord); Assert.assertTrue(requestSuccess); @@ -215,8 +230,8 @@ public void testInsertWorkloadRecordException() { WorkloadRecord workloadRecord = new WorkloadRecord(); - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).putItem(ArgumentMatchers.any(Item.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).putItem(ArgumentMatchers.any(PutItemRequest.class)); DynamoDBWorkloadRecordStoreConnection dbConn = getDBConnection(); boolean requestSuccess = dbConn.insertWorkloadRecord(workloadRecord); @@ -236,13 +251,20 @@ public void testUpdateWorkloadRecord() { Date currDate = new Date(currTime); workloadRecord.setUpdateTime(currDate); - UpdateItemSpec item = new UpdateItemSpec() - .withPrimaryKey("primaryKey", "athenz.api#1234#10.0.0.1") - .withAttributeUpdate( - new AttributeUpdate("provider").put(workloadRecord.getProvider()), - new AttributeUpdate("updateTime").put(workloadRecord.getUpdateTime())); + HashMap updatedValues = new HashMap<>(); + DynamoDBUtils.updateItemStringValue(updatedValues, "provider", workloadRecord.getProvider()); + DynamoDBUtils.updateItemLongValue(updatedValues, "updateTime", workloadRecord.getUpdateTime()); + + HashMap itemKey = new HashMap<>(); + itemKey.put("primaryKey", AttributeValue.fromS("athenz.api#1234#10.0.0.1")); + + UpdateItemRequest request = UpdateItemRequest.builder() + .tableName(tableName) + .key(itemKey) + .attributeUpdates(updatedValues) + .build(); - Mockito.doReturn(updateOutcome).when(table).updateItem(item); + Mockito.doReturn(updateOutcome).when(dynamoDB).updateItem(request); boolean requestSuccess = dbConn.updateWorkloadRecord(workloadRecord); Assert.assertTrue(requestSuccess); @@ -254,8 +276,8 @@ public void testUpdateWorkloadRecordException() { WorkloadRecord workloadRecord = new WorkloadRecord(); - Mockito.doThrow(new AmazonDynamoDBException("invalid operation")) - .when(table).updateItem(ArgumentMatchers.any(UpdateItemSpec.class)); + Mockito.doThrow(AwsServiceException.create("invalid operation", new Throwable("invalid operation"))) + .when(dynamoDB).updateItem(ArgumentMatchers.any(UpdateItemRequest.class)); DynamoDBWorkloadRecordStoreConnection dbConn = getDBConnection(); boolean requestSuccess = dbConn.updateWorkloadRecord(workloadRecord); diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactoryTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactoryTest.java index 88e3967beb9..73c0b578a1a 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactoryTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreFactoryTest.java @@ -15,9 +15,6 @@ */ package com.yahoo.athenz.zts.workload.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.DynamoDB; -import com.amazonaws.services.dynamodbv2.document.Table; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.common.server.workload.WorkloadRecordStore; import com.yahoo.athenz.zts.ResourceException; @@ -29,20 +26,19 @@ import org.testng.Assert; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import static com.yahoo.athenz.zts.ZTSConsts.*; import static org.mockito.Mockito.when; public class DynamoDBWorkloadRecordStoreFactoryTest { - @Mock - private AmazonDynamoDB dbClient; - @Mock private Table table; - @Mock private DynamoDB dynamoDB; + + @Mock private DynamoDbClient dbClient; class TestDynamoDBWorkloadRecordStoreFactory extends DynamoDBWorkloadRecordStoreFactory { @Override - AmazonDynamoDB getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { + DynamoDbClient getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotificationSender, PrivateKeyStore keyStore) { return dbClient; } } @@ -50,7 +46,6 @@ AmazonDynamoDB getDynamoDBClient(ZTSClientNotificationSenderImpl ztsClientNotifi @BeforeMethod public void setUp() { MockitoAnnotations.openMocks(this); - Mockito.doReturn(table).when(dynamoDB).getTable("Workloads-Table"); } @Test @@ -184,7 +179,7 @@ public void testGetDynamoDBClient() { DynamoDBWorkloadRecordStoreFactory factory = new DynamoDBWorkloadRecordStoreFactory(); ZTSClientNotificationSenderImpl ztsClientNotificationSender = Mockito.mock(ZTSClientNotificationSenderImpl.class); PrivateKeyStore privateKeyStore = Mockito.mock(PrivateKeyStore.class); - AmazonDynamoDB dynamoDBClient = factory.getDynamoDBClient(ztsClientNotificationSender, privateKeyStore); + DynamoDbClient dynamoDBClient = factory.getDynamoDBClient(ztsClientNotificationSender, privateKeyStore); Assert.assertNotNull(dynamoDBClient); System.clearProperty(ZTS_PROP_DYNAMODB_KEY_PATH); diff --git a/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreTest.java b/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreTest.java index 2c8009af144..c0364571828 100644 --- a/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreTest.java +++ b/servers/zts/src/test/java/com/yahoo/athenz/zts/workload/impl/DynamoDBWorkloadRecordStoreTest.java @@ -15,19 +15,17 @@ */ package com.yahoo.athenz.zts.workload.impl; -import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; import com.yahoo.athenz.common.server.workload.WorkloadRecordStoreConnection; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import org.testng.Assert; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; - +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; public class DynamoDBWorkloadRecordStoreTest { - @Mock - private AmazonDynamoDB dbClient; + @Mock private DynamoDbClient dbClient; @BeforeMethod public void setUp() { @@ -37,7 +35,8 @@ public void setUp() { @Test public void testGetConnection() { - DynamoDBWorkloadRecordStore store = new DynamoDBWorkloadRecordStore(dbClient, "Workload-Table", "service-index", "ip-index"); + DynamoDBWorkloadRecordStore store = new DynamoDBWorkloadRecordStore(dbClient, "Workload-Table", + "service-index", "ip-index"); WorkloadRecordStoreConnection dbConn = store.getConnection(); Assert.assertNotNull(dbConn); @@ -46,18 +45,4 @@ public void testGetConnection() { store.setOperationTimeout(10); store.clearConnections(); } - - @Test - public void testGetConnectionException() { - - // passing null for table name to get exception - DynamoDBWorkloadRecordStore store = new DynamoDBWorkloadRecordStore(dbClient, null, null, null); - - try { - store.getConnection(); - Assert.fail(); - } catch (Exception ignored) { - } - } - } diff --git a/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySender.java b/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySender.java index d05db4d04d9..1d70eff21e4 100644 --- a/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySender.java +++ b/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySender.java @@ -53,7 +53,7 @@ public class DynamoDBAuthHistorySender implements AuthHistorySender { private final DynamoDbEnhancedAsyncClient enhancedClient; private final DynamoDbAsyncTable mappedTable; - public DynamoDBAuthHistorySender(DynamoDbAsyncClient dynamoDB) throws InterruptedException { + public DynamoDBAuthHistorySender(DynamoDbAsyncClient dynamoDB) { enhancedClient = DynamoDbEnhancedAsyncClient.builder() .dynamoDbClient(dynamoDB) .build(); @@ -75,9 +75,7 @@ public void pushRecords(Set logs) throws RuntimeExcep LOGGER.info("Started Pushing to DB {} records", logs.size()); Iterable> batchPartitions = splitToBatchSizePartitions(logs); List> futures = new ArrayList<>(); - batchPartitions.forEach(partition -> { - futures.add(putBatchPartition(partition, 1)); - }); + batchPartitions.forEach(partition -> futures.add(putBatchPartition(partition, 1))); futures.forEach(CompletableFuture::join); for (CompletableFuture future : futures) { @@ -93,7 +91,7 @@ public void pushRecords(Set logs) throws RuntimeExcep LOGGER.info("Finished Pushing to DB {} records", logs.size()); } - private static void createTableIfNotExists(DynamoDbAsyncClient dynamoDbAsyncClient, String tableName) throws InterruptedException { + private static void createTableIfNotExists(DynamoDbAsyncClient dynamoDbAsyncClient, String tableName) { GlobalSecondaryIndex principalDomainIndex = GlobalSecondaryIndex.builder() .indexName(PRINCIPAL_DOMAIN_INDEX_NAME) .keySchema( @@ -141,7 +139,7 @@ private static void createTableIfNotExists(DynamoDbAsyncClient dynamoDbAsyncClie try { LOGGER.info("Trying to create table: {}", tableName); dynamoDbAsyncClient.createTable(createTableRequest).get(); - } catch (ExecutionException | DynamoDbException ex) { + } catch (Exception ex) { LOGGER.error("Table {} creation failed. Error: {}", tableName, ex.getMessage(), ex); // It is possible that the table already exists so if creation fails we will only log the error. // If the table doesn't exist, we will fail later during push. @@ -159,7 +157,7 @@ private static void createTableIfNotExists(DynamoDbAsyncClient dynamoDbAsyncClie try { dynamoDbAsyncClient.updateTimeToLive(updateTimeToLiveRequest).get(); LOGGER.info("Table {} TTL enabled successfully", tableName); - } catch (ExecutionException | DynamoDbException ex) { + } catch (Exception ex) { LOGGER.error("Table {} ttl update failed. Error: {}", tableName, ex.getMessage(), ex); } } @@ -189,7 +187,7 @@ private void retryIfItemsRemaining(int retryCount, CompletableFuture unprocessedPutItems = result.unprocessedPutItemsForTable(mappedTable); if (unprocessedPutItems != null && !unprocessedPutItems.isEmpty()) { try { - TimeUnit.SECONDS.sleep(retryCount * 2); + TimeUnit.SECONDS.sleep(retryCount * 2L); } catch (InterruptedException e) { CompletableFuture failResult = new CompletableFuture<>(); String error = "Failed to write batch for " + retryCount + " times due to InterruptedException: " + e.getMessage(); @@ -213,8 +211,7 @@ private CompletableFuture getBatchWriteResultCompletableFuture }); WriteBatch writeBatch = writeBatchBuilder.build(); BatchWriteItemEnhancedRequest writeItemEnhancedRequest = BatchWriteItemEnhancedRequest.builder().addWriteBatch(writeBatch).build(); - CompletableFuture batchWriteResultCompletableFuture = enhancedClient.batchWriteItem(writeItemEnhancedRequest); - return batchWriteResultCompletableFuture; + return enhancedClient.batchWriteItem(writeItemEnhancedRequest); } private Iterable> splitToBatchSizePartitions(Set logs) { diff --git a/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactory.java b/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactory.java index 3fdc72804d2..fa7ee9045b3 100644 --- a/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactory.java +++ b/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactory.java @@ -30,25 +30,21 @@ public class DynamoDBAuthHistorySenderFactory implements AuthHistorySenderFactory { @Override - public AuthHistorySender create(PrivateKeyStore pkeyStore, String region) { - try { - DynamoDBClientFetcher dynamoDBClientFetcher = DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); - DynamoDBClientSettings clientSettings = getClientSettings(pkeyStore); - DynamoDBClientAndCredentials dynamoDBClient = dynamoDBClientFetcher.getDynamoDBClient(null, clientSettings); - return new DynamoDBAuthHistorySender(dynamoDBClient.getAmazonDynamoAsyncDB()); - } catch (InterruptedException e) { - throw new RuntimeException("Failed to instantiate AuthHistorySender: ", e); - } + public AuthHistorySender create(PrivateKeyStore pkeyStore, final String region) { + DynamoDBClientFetcher dynamoDBClientFetcher = DynamoDBClientFetcherFactory.getDynamoDBClientFetcher(); + DynamoDBClientSettings clientSettings = getClientSettings(pkeyStore, region); + DynamoDBClientAndCredentials dynamoDBClient = dynamoDBClientFetcher.getDynamoDBClient(null, clientSettings); + return new DynamoDBAuthHistorySender(dynamoDBClient.getDynamoDbAsyncClient()); } - private DynamoDBClientSettings getClientSettings(PrivateKeyStore pkeyStore) { + private DynamoDBClientSettings getClientSettings(PrivateKeyStore pkeyStore, final String defaultRegion) { String keyPath = System.getProperty(PROP_DYNAMODB_KEY_PATH, ""); String certPath = System.getProperty(PROP_DYNAMODB_CERT_PATH, ""); String domainName = System.getProperty(PROP_DYNAMODB_DOMAIN, ""); String roleName = System.getProperty(PROP_DYNAMODB_ROLE, ""); String trustStore = System.getProperty(PROP_DYNAMODB_TRUSTSTORE, ""); String trustStorePassword = System.getProperty(PROP_DYNAMODB_TRUSTSTORE_PASSWORD, ""); - String region = System.getProperty(PROP_DYNAMODB_REGION, ""); + String region = System.getProperty(PROP_DYNAMODB_REGION, defaultRegion); String appName = System.getProperty(PROP_DYNAMODB_TRUSTSTORE_APPNAME, ""); String ztsURL = System.getProperty(PROP_DYNAMODB_ZTS_URL, ""); String externalId = System.getProperty(PROP_DYNAMODB_EXTERNAL_ID, null); @@ -56,8 +52,10 @@ private DynamoDBClientSettings getClientSettings(PrivateKeyStore pkeyStore) { String maxExpiryTimeStr = System.getProperty(PROP_DYNAMODB_MAX_EXPIRY_TIME, ""); Integer minExpiryTime = minExpiryTimeStr.isEmpty() ? null : Integer.parseInt(minExpiryTimeStr); Integer maxExpiryTime = maxExpiryTimeStr.isEmpty() ? null : Integer.parseInt(maxExpiryTimeStr); - String keygroupName = System.getProperty(PROP_DYNAMODB_TRUSTSTORE_KEYGROUPNAME, ""); + String keyGroupName = System.getProperty(PROP_DYNAMODB_TRUSTSTORE_KEYGROUPNAME, ""); - return new DynamoDBClientSettings(certPath, domainName, roleName, trustStore, trustStorePassword, ztsURL, region, keyPath, appName, pkeyStore, externalId, minExpiryTime, maxExpiryTime, keygroupName); + return new DynamoDBClientSettings(certPath, domainName, roleName, trustStore, trustStorePassword, + ztsURL, region, keyPath, appName, pkeyStore, externalId, minExpiryTime, maxExpiryTime, + keyGroupName, true); } } diff --git a/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/LocalAuthHistoryFetcher.java b/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/LocalAuthHistoryFetcher.java index bf987899149..14abf2a8fbc 100644 --- a/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/LocalAuthHistoryFetcher.java +++ b/syncers/auth_history_syncer/src/main/java/com/yahoo/athenz/syncer/auth/history/impl/LocalAuthHistoryFetcher.java @@ -71,7 +71,6 @@ private Set getRecordsFromFile(String logFileName, Lo } } catch (Exception e) { LOGGER.error("Failed to parse log event", e); - return null; } return records; diff --git a/syncers/auth_history_syncer/src/test/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactoryTest.java b/syncers/auth_history_syncer/src/test/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactoryTest.java index 0960c50fd29..0c57634905d 100644 --- a/syncers/auth_history_syncer/src/test/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactoryTest.java +++ b/syncers/auth_history_syncer/src/test/java/com/yahoo/athenz/syncer/auth/history/impl/DynamoDBAuthHistorySenderFactoryTest.java @@ -18,7 +18,6 @@ package com.yahoo.athenz.syncer.auth.history.impl; -import com.amazonaws.SdkClientException; import com.google.common.io.Resources; import com.yahoo.athenz.auth.PrivateKeyStore; import com.yahoo.athenz.syncer.auth.history.AuthHistorySender; @@ -60,15 +59,11 @@ public void testCreate() { } @Test - public void testCreateFailRegion() { - try { - DynamoDBAuthHistorySenderFactory dynamoDBAuthHistorySenderFactory = new DynamoDBAuthHistorySenderFactory(); - PrivateKeyStore pkeyStore = Mockito.mock(PrivateKeyStore.class); - dynamoDBAuthHistorySenderFactory.create(pkeyStore, "us-west-2"); - fail(); - } catch (SdkClientException ex) { - assertEquals("Could not find region information for 'null' in SDK metadata.", ex.getMessage()); - } + public void testCreateRegionArgument() { + DynamoDBAuthHistorySenderFactory dynamoDBAuthHistorySenderFactory = new DynamoDBAuthHistorySenderFactory(); + PrivateKeyStore pkeyStore = Mockito.mock(PrivateKeyStore.class); + AuthHistorySender authHistorySender = dynamoDBAuthHistorySenderFactory.create(pkeyStore, "us-west-2"); + assertNotNull(authHistorySender); } @Test