diff --git a/.build/cassandra-build-deps-template.xml b/.build/cassandra-build-deps-template.xml index cbf2f8c59d..b5db6846da 100644 --- a/.build/cassandra-build-deps-template.xml +++ b/.build/cassandra-build-deps-template.xml @@ -123,5 +123,9 @@ com.fasterxml.jackson.dataformat jackson-dataformat-yaml + + com.github.tomakehurst + wiremock-jre8 + diff --git a/.build/parent-pom-template.xml b/.build/parent-pom-template.xml index 9203d4ad89..a88764ee75 100644 --- a/.build/parent-pom-template.xml +++ b/.build/parent-pom-template.xml @@ -492,6 +492,12 @@ 0.0.15 test + + com.github.tomakehurst + wiremock-jre8 + 2.35.0 + test + net.java.dev.jna jna diff --git a/CHANGES.txt b/CHANGES.txt index 2c3666fcf4..cf87114f26 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -164,6 +164,7 @@ Merged from 3.11: * Fix error message handling when trying to use CLUSTERING ORDER with non-clustering column (CASSANDRA-17818 * Add keyspace and table name to exception message during ColumnSubselection deserialization (CASSANDRA-18346) Merged from 3.0: + * Add support for AWS Ec2 IMDSv2 (CASSANDRA-16555) * Suppress CVE-2023-35116 (CASSANDRA-18630) * Pass taskId from CompactionTask to system.compaction_history (CASSANDRA-12183) * Backport CASSANDRA-10508: Remove hard-coded SSL cipher suites (CASSANDRA-18575) diff --git a/NEWS.txt b/NEWS.txt index 11eafa231b..1ecd60ab29 100644 --- a/NEWS.txt +++ b/NEWS.txt @@ -193,6 +193,11 @@ Upgrading removed. Any deployments that depend upon the code default to this password value without explicitly specifying it in cassandra.yaml will fail on upgrade. Please specify your keystore_password and truststore_password elements in cassandra.yaml with appropriate values to prevent this failure. + - Please beware that if you use Ec2Snitch or Ec2MultiRegionSnitch, by default it will + communicate with AWS IMDS of version 2. This change is transparent, there does not need + to be done anything upon upgrade. Furthermore, IMDS of version 2 can be configured to be required in AWS EC2 console. + Consult cassandra-rackdc.properties for more details. (CASSANDRA-16555) + Deprecation ----------- diff --git a/conf/cassandra-rackdc.properties b/conf/cassandra-rackdc.properties index cc472b4e61..2d17808d30 100644 --- a/conf/cassandra-rackdc.properties +++ b/conf/cassandra-rackdc.properties @@ -37,3 +37,9 @@ rack=rack1 # rack name is the region plus the availability zone letter. # Examples: us-west-1a => dc: us-west-1, rack: us-west-1a; us-west-2b => dc: us-west-2, rack: us-west-2b; # ec2_naming_scheme=standard +# +# Type of AWS IMDS, value might be either "v1" or "v2". Defaults to "v2". +# ec2_metadata_type=v2 +# If AWS IMDS of v2 is configured, ec2_metadata_token_ttl_seconds says how many seconds a token will be valid until +# it is refreshed. Defaults to 21600. Can not be smaller than 30 and bigger than 21600. Has to be an integer. +# ec2_metadata_token_ttl_seconds=21600 diff --git a/src/java/org/apache/cassandra/locator/AbstractCloudMetadataServiceConnector.java b/src/java/org/apache/cassandra/locator/AbstractCloudMetadataServiceConnector.java new file mode 100644 index 0000000000..6210ce40b5 --- /dev/null +++ b/src/java/org/apache/cassandra/locator/AbstractCloudMetadataServiceConnector.java @@ -0,0 +1,99 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.cassandra.locator; + +import java.io.DataInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.net.HttpURLConnection; +import java.net.URL; +import java.util.Map; + +import com.google.common.collect.ImmutableMap; + +import static java.nio.charset.StandardCharsets.UTF_8; + +abstract class AbstractCloudMetadataServiceConnector +{ + protected final String metadataServiceUrl; + + protected AbstractCloudMetadataServiceConnector(String metadataServiceUrl) + { + this.metadataServiceUrl = metadataServiceUrl; + } + + public String apiCall(String query) throws IOException + { + return apiCall(metadataServiceUrl, query, 200); + } + + public String apiCall(String url, String query, int expectedResponseCode) throws IOException + { + return apiCall(url, query, "GET", ImmutableMap.of(), expectedResponseCode); + } + + public String apiCall(String url, + String query, + String method, + Map extraHeaders, + int expectedResponseCode) throws IOException + { + HttpURLConnection conn = null; + try + { + // Populate the region and zone by introspection, fail if 404 on metadata + conn = (HttpURLConnection) new URL(url + query).openConnection(); + extraHeaders.forEach(conn::setRequestProperty); + conn.setRequestMethod(method); + if (conn.getResponseCode() != expectedResponseCode) + throw new HttpException(conn.getResponseCode(), conn.getResponseMessage()); + + // Read the information. I wish I could say (String) conn.getContent() here... + int cl = conn.getContentLength(); + + if (cl == -1) + return null; + + byte[] b = new byte[cl]; + try (DataInputStream d = new DataInputStream((InputStream) conn.getContent())) + { + d.readFully(b); + } + return new String(b, UTF_8); + } + finally + { + if (conn != null) + conn.disconnect(); + } + } + + public static final class HttpException extends IOException + { + public final int responseCode; + public final String responseMessage; + + public HttpException(int responseCode, String responseMessage) + { + super("HTTP response code: " + responseCode + " (" + responseMessage + ')'); + this.responseCode = responseCode; + this.responseMessage = responseMessage; + } + } +} diff --git a/src/java/org/apache/cassandra/locator/Ec2MetadataServiceConnector.java b/src/java/org/apache/cassandra/locator/Ec2MetadataServiceConnector.java new file mode 100644 index 0000000000..b5ab2b75d7 --- /dev/null +++ b/src/java/org/apache/cassandra/locator/Ec2MetadataServiceConnector.java @@ -0,0 +1,215 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.cassandra.locator; + +import java.io.IOException; +import java.time.Duration; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableMap; + +import org.apache.cassandra.exceptions.ConfigurationException; +import org.apache.cassandra.utils.Clock; +import org.apache.cassandra.utils.Pair; + +import static java.lang.String.format; +import static java.util.Arrays.stream; +import static java.util.stream.Collectors.joining; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.EC2MetadataType.v2; + +abstract class Ec2MetadataServiceConnector extends AbstractCloudMetadataServiceConnector +{ + static final String EC2_METADATA_TYPE_PROPERTY = "ec2_metadata_type"; + static final String EC2_METADATA_URL_PROPERTY = "ec2_metadata_url"; + static final String DEFAULT_EC2_METADATA_URL = "http://169.254.169.254"; + + Ec2MetadataServiceConnector(SnitchProperties properties) + { + super(properties.get(EC2_METADATA_URL_PROPERTY, DEFAULT_EC2_METADATA_URL)); + } + + enum EC2MetadataType + { + v1(props -> Ec2MetadataServiceConnector.V1Connector.create(props)), + v2(props -> Ec2MetadataServiceConnector.V2Connector.create(props)); + + private final Function provider; + + EC2MetadataType(Function provider) + { + this.provider = provider; + } + + Ec2MetadataServiceConnector create(SnitchProperties properties) + { + return provider.apply(properties); + } + } + + static Ec2MetadataServiceConnector create(SnitchProperties props) + { + try + { + return EC2MetadataType.valueOf(props.get(EC2_METADATA_TYPE_PROPERTY, v2.name())).create(props); + } + catch (IllegalArgumentException ex) + { + throw new ConfigurationException(format("%s must be one of %s", EC2_METADATA_TYPE_PROPERTY, + stream(EC2MetadataType.values()).map(Enum::name).collect(joining(", ")))); + } + } + + static class V1Connector extends Ec2MetadataServiceConnector + { + static V1Connector create(SnitchProperties props) + { + return new V1Connector(props); + } + + V1Connector(SnitchProperties props) + { + super(props); + } + + @Override + public String toString() + { + return String.format("%s{%s=%s}", V1Connector.class.getName(), EC2_METADATA_URL_PROPERTY, metadataServiceUrl); + } + } + + static class V2Connector extends Ec2MetadataServiceConnector + { + static final int MAX_TOKEN_TIME_IN_SECONDS = 21600; + static final int MIN_TOKEN_TIME_IN_SECONDS = 30; + + static final String AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY = "ec2_metadata_token_ttl_seconds"; + static final String AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER = "X-aws-ec2-metadata-token-ttl-seconds"; + static final String AWS_EC2_METADATA_TOKEN_HEADER = "X-aws-ec2-metadata-token"; + static final String TOKEN_QUERY = "/latest/api/token"; + + @VisibleForTesting + static int HTTP_REQUEST_RETRIES = 1; + + private Pair token; + @VisibleForTesting + final Duration tokenTTL; + + static V2Connector create(SnitchProperties props) + { + String tokenTTLString = props.get(AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, + Integer.toString(MAX_TOKEN_TIME_IN_SECONDS)); + + Duration tokenTTL; + try + { + tokenTTL = Duration.ofSeconds(Integer.parseInt(tokenTTLString)); + + if (tokenTTL.getSeconds() < MIN_TOKEN_TIME_IN_SECONDS || tokenTTL.getSeconds() > MAX_TOKEN_TIME_IN_SECONDS) + { + throw new ConfigurationException(format("property %s was set to %s seconds which is not in allowed range of [%s..%s]", + AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, + tokenTTL.getSeconds(), + MIN_TOKEN_TIME_IN_SECONDS, + MAX_TOKEN_TIME_IN_SECONDS)); + } + } + catch (NumberFormatException ex) + { + throw new ConfigurationException(format("Unable to parse integer from property %s, value to parse: %s", + AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, tokenTTLString)); + } + + return new V2Connector(props, tokenTTL); + } + + V2Connector(SnitchProperties properties, Duration tokenTTL) + { + super(properties); + this.tokenTTL = tokenTTL; + } + + @Override + public String apiCall(String url, String query, String method, Map extraHeaders, int expectedResponseCode) throws IOException + { + Map headers = new HashMap<>(extraHeaders); + for (int retry = 0; retry <= HTTP_REQUEST_RETRIES; retry++) + { + String resolvedToken; + if (token != null && token.right > Clock.Global.currentTimeMillis()) + resolvedToken = token.left; + else + resolvedToken = getToken(); + + try + { + headers.put(AWS_EC2_METADATA_TOKEN_HEADER, resolvedToken); + return super.apiCall(url, query, method, headers, expectedResponseCode); + } + catch (HttpException ex) + { + if (retry == HTTP_REQUEST_RETRIES) + throw ex; + + if (ex.responseCode == 401) // invalidate token if it is 401 + this.token = null; + } + } + + throw new AssertionError(); + } + + @Override + public String toString() + { + return String.format("%s{%s=%s,%s=%s}", + V2Connector.class.getName(), + EC2_METADATA_URL_PROPERTY, metadataServiceUrl, + AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, + tokenTTL.getSeconds()); + } + + /** + * Get a session token to enable requests to the meta data service. + * IMDSv2 + */ + @VisibleForTesting + public synchronized String getToken() + { + try + { + token = Pair.create(super.apiCall(metadataServiceUrl, + TOKEN_QUERY, + "PUT", + ImmutableMap.of(AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER, String.valueOf(tokenTTL.getSeconds())), + 200), + Clock.Global.currentTimeMillis() + tokenTTL.toMillis() - TimeUnit.SECONDS.toMillis(5)); + return token.left; + } + catch (IOException e) + { + throw new RuntimeException(e); + } + } + } +} diff --git a/src/java/org/apache/cassandra/locator/Ec2MultiRegionSnitch.java b/src/java/org/apache/cassandra/locator/Ec2MultiRegionSnitch.java index 45c387da19..1ddb2b8962 100644 --- a/src/java/org/apache/cassandra/locator/Ec2MultiRegionSnitch.java +++ b/src/java/org/apache/cassandra/locator/Ec2MultiRegionSnitch.java @@ -21,8 +21,10 @@ import java.io.IOException; import java.net.InetAddress; import java.net.UnknownHostException; -import org.apache.cassandra.exceptions.ConfigurationException; +import com.google.common.annotations.VisibleForTesting; + import org.apache.cassandra.config.DatabaseDescriptor; +import org.apache.cassandra.exceptions.ConfigurationException; import org.apache.cassandra.gms.ApplicationState; import org.apache.cassandra.gms.Gossiper; import org.apache.cassandra.service.StorageService; @@ -40,16 +42,28 @@ import org.apache.cassandra.service.StorageService; */ public class Ec2MultiRegionSnitch extends Ec2Snitch { - private static final String PUBLIC_IP_QUERY_URL = "http://169.254.169.254/latest/meta-data/public-ipv4"; - private static final String PRIVATE_IP_QUERY_URL = "http://169.254.169.254/latest/meta-data/local-ipv4"; + @VisibleForTesting + static final String PUBLIC_IP_QUERY = "/latest/meta-data/public-ipv4"; + @VisibleForTesting + static final String PRIVATE_IP_QUERY = "/latest/meta-data/local-ipv4"; private final String localPrivateAddress; public Ec2MultiRegionSnitch() throws IOException, ConfigurationException { - super(); - InetAddress localPublicAddress = InetAddress.getByName(awsApiCall(PUBLIC_IP_QUERY_URL)); + this(new SnitchProperties()); + } + + public Ec2MultiRegionSnitch(SnitchProperties props) throws IOException, ConfigurationException + { + this(props, Ec2MetadataServiceConnector.create(props)); + } + + Ec2MultiRegionSnitch(SnitchProperties props, Ec2MetadataServiceConnector connector) throws IOException + { + super(props, connector); + InetAddress localPublicAddress = InetAddress.getByName(connector.apiCall(PUBLIC_IP_QUERY)); logger.info("EC2Snitch using publicIP as identifier: {}", localPublicAddress); - localPrivateAddress = awsApiCall(PRIVATE_IP_QUERY_URL); + localPrivateAddress = connector.apiCall(PRIVATE_IP_QUERY); // use the Public IP to broadcast Address to other nodes. DatabaseDescriptor.setBroadcastAddress(localPublicAddress); if (DatabaseDescriptor.getBroadcastRpcAddress() == null) diff --git a/src/java/org/apache/cassandra/locator/Ec2Snitch.java b/src/java/org/apache/cassandra/locator/Ec2Snitch.java index 5e51408669..bf72a5f75d 100644 --- a/src/java/org/apache/cassandra/locator/Ec2Snitch.java +++ b/src/java/org/apache/cassandra/locator/Ec2Snitch.java @@ -17,29 +17,40 @@ */ package org.apache.cassandra.locator; -import java.io.DataInputStream; -import java.io.FilterInputStream; import java.io.IOException; -import java.net.HttpURLConnection; -import java.net.URL; -import java.nio.charset.StandardCharsets; import java.util.Map; import java.util.Set; import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; import org.slf4j.LoggerFactory; + import org.apache.cassandra.db.SystemKeyspace; import org.apache.cassandra.exceptions.ConfigurationException; import org.apache.cassandra.gms.ApplicationState; import org.apache.cassandra.gms.EndpointState; import org.apache.cassandra.gms.Gossiper; -import org.apache.cassandra.io.util.FileUtils; import org.apache.cassandra.utils.FBUtilities; /** * A snitch that assumes an EC2 region is a DC and an EC2 availability_zone - * is a rack. This information is available in the config for the node. + * is a rack. This information is available in the config for the node. + + * Since CASSANDRA-16555, it is possible to choose version of AWS IMDS. + * + * By default, since CASSANDRA-16555, IMDSv2 is used. + * + * The version of IMDS is driven by property {@link Ec2MetadataServiceConnector#EC2_METADATA_TYPE_PROPERTY} and + * can be of value either 'v1' or 'v2'. + * + * It is possible to specify custom URL of IMDS by {@link Ec2MetadataServiceConnector#EC2_METADATA_URL_PROPERTY}. + * A user is not meant to change this under normal circumstances, it is suitable for testing only. + * + * IMDSv2 is secured by a token which needs to be fetched from IDMSv2 first, and it has to be passed in a header + * for the actual queries to IDMSv2. Ec2Snitch is doing this automatically. The only configuration parameter exposed + * to a user is {@link Ec2MetadataServiceConnector.V2Connector#AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY} + * which is by default set to {@link Ec2MetadataServiceConnector.V2Connector#MAX_TOKEN_TIME_IN_SECONDS}. TTL has + * to be an integer from the range [30, 21600]. */ public class Ec2Snitch extends AbstractNetworkTopologySnitch { @@ -49,7 +60,9 @@ public class Ec2Snitch extends AbstractNetworkTopologySnitch static final String EC2_NAMING_LEGACY = "legacy"; private static final String EC2_NAMING_STANDARD = "standard"; - private static final String ZONE_NAME_QUERY_URL = "http://169.254.169.254/latest/meta-data/placement/availability-zone"; + @VisibleForTesting + public static final String ZONE_NAME_QUERY = "/latest/meta-data/placement/availability-zone"; + private static final String DEFAULT_DC = "UNKNOWN-DC"; private static final String DEFAULT_RACK = "UNKNOWN-RACK"; @@ -59,6 +72,8 @@ public class Ec2Snitch extends AbstractNetworkTopologySnitch private Map> savedEndpoints; + protected final Ec2MetadataServiceConnector connector; + public Ec2Snitch() throws IOException, ConfigurationException { this(new SnitchProperties()); @@ -66,7 +81,13 @@ public class Ec2Snitch extends AbstractNetworkTopologySnitch public Ec2Snitch(SnitchProperties props) throws IOException, ConfigurationException { - String az = awsApiCall(ZONE_NAME_QUERY_URL); + this(props, Ec2MetadataServiceConnector.create(props)); + } + + Ec2Snitch(SnitchProperties props, Ec2MetadataServiceConnector connector) throws IOException + { + this.connector = connector; + String az = connector.apiCall(ZONE_NAME_QUERY); // if using the full naming scheme, region name is created by removing letters from the // end of the availability zone and zone is the full zone name @@ -93,7 +114,7 @@ public class Ec2Snitch extends AbstractNetworkTopologySnitch String datacenterSuffix = props.get("dc_suffix", ""); ec2region = region.concat(datacenterSuffix); - logger.info("EC2Snitch using region: {}, zone: {}.", ec2region, ec2zone); + logger.info("EC2Snitch using region: {}, zone: {}, properties: {}", ec2region, ec2zone, connector); } private static boolean isUsingLegacyNaming(SnitchProperties props) @@ -101,31 +122,6 @@ public class Ec2Snitch extends AbstractNetworkTopologySnitch return props.get(SNITCH_PROP_NAMING_SCHEME, EC2_NAMING_STANDARD).equalsIgnoreCase(EC2_NAMING_LEGACY); } - String awsApiCall(String url) throws IOException, ConfigurationException - { - // Populate the region and zone by introspection, fail if 404 on metadata - HttpURLConnection conn = (HttpURLConnection) new URL(url).openConnection(); - DataInputStream d = null; - try - { - conn.setRequestMethod("GET"); - if (conn.getResponseCode() != 200) - throw new ConfigurationException("Ec2Snitch was unable to execute the API call. Not an ec2 node?"); - - // Read the information. I wish I could say (String) conn.getContent() here... - int cl = conn.getContentLength(); - byte[] b = new byte[cl]; - d = new DataInputStream((FilterInputStream) conn.getContent()); - d.readFully(b); - return new String(b, StandardCharsets.UTF_8); - } - finally - { - FileUtils.close(d); - conn.disconnect(); - } - } - public String getRack(InetAddressAndPort endpoint) { if (endpoint.equals(FBUtilities.getBroadcastAddressAndPort())) diff --git a/src/java/org/apache/cassandra/locator/SnitchProperties.java b/src/java/org/apache/cassandra/locator/SnitchProperties.java index b0bfe4c2ef..efd29300ac 100644 --- a/src/java/org/apache/cassandra/locator/SnitchProperties.java +++ b/src/java/org/apache/cassandra/locator/SnitchProperties.java @@ -21,6 +21,8 @@ import java.io.InputStream; import java.net.URL; import java.util.Properties; +import com.google.common.annotations.VisibleForTesting; + import org.apache.cassandra.io.util.FileUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -61,6 +63,12 @@ public class SnitchProperties } } + @VisibleForTesting + public SnitchProperties(Properties properties) + { + this.properties = properties; + } + /** * Get a snitch property value or return defaultValue if not defined. */ diff --git a/test/unit/org/apache/cassandra/locator/Ec2ConnectorTest.java b/test/unit/org/apache/cassandra/locator/Ec2ConnectorTest.java new file mode 100644 index 0000000000..756e13cef3 --- /dev/null +++ b/test/unit/org/apache/cassandra/locator/Ec2ConnectorTest.java @@ -0,0 +1,106 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.cassandra.locator; + +import java.time.Duration; +import java.util.Properties; + +import org.junit.Test; + +import org.apache.cassandra.exceptions.ConfigurationException; +import org.apache.cassandra.locator.Ec2MetadataServiceConnector.EC2MetadataType; +import org.apache.cassandra.locator.Ec2MetadataServiceConnector.V1Connector; +import org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector; + +import static java.lang.String.format; +import static java.util.Arrays.stream; +import static java.util.stream.Collectors.joining; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.DEFAULT_EC2_METADATA_URL; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.EC2_METADATA_TYPE_PROPERTY; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector.AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector.MAX_TOKEN_TIME_IN_SECONDS; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector.MIN_TOKEN_TIME_IN_SECONDS; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class Ec2ConnectorTest +{ + @Test + public void testV1Configuration() + { + Properties p = new Properties(); + p.setProperty(EC2_METADATA_TYPE_PROPERTY, EC2MetadataType.v1.name()); + Ec2MetadataServiceConnector ec2Connector = Ec2MetadataServiceConnector.create(new SnitchProperties(p)); + + assertTrue(ec2Connector instanceof V1Connector); + + assertEquals(DEFAULT_EC2_METADATA_URL, ec2Connector.metadataServiceUrl); + } + + @Test + public void testV2Configuration() + { + Ec2MetadataServiceConnector ec2Connector = Ec2MetadataServiceConnector.create(new SnitchProperties(new Properties())); + + // v2 connector by default + assertTrue(ec2Connector instanceof V2Connector); + assertEquals(DEFAULT_EC2_METADATA_URL, ec2Connector.metadataServiceUrl); + assertEquals(Duration.ofSeconds(MAX_TOKEN_TIME_IN_SECONDS), ((V2Connector) ec2Connector).tokenTTL); + } + + @Test + public void testInvalidConfiguration() + { + assertThatExceptionOfType(ConfigurationException.class) + .describedAs("it should not be possible create a connector of type non-existent") + .isThrownBy(() -> { + Properties p = new Properties(); + p.setProperty(EC2_METADATA_TYPE_PROPERTY, "non-existent"); + Ec2MetadataServiceConnector.create(new SnitchProperties(p)); + }).withMessage(format("%s must be one of %s", EC2_METADATA_TYPE_PROPERTY, + stream(EC2MetadataType.values()).map(Enum::name).collect(joining(", ")))); + + assertThatExceptionOfType(ConfigurationException.class) + .describedAs(format("it should not be possible to set %s to more than %s", + AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, + MAX_TOKEN_TIME_IN_SECONDS)) + .isThrownBy(() -> { + Properties p = new Properties(); + p.setProperty(EC2_METADATA_TYPE_PROPERTY, EC2MetadataType.v2.name()); + p.setProperty(AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, "50000"); + + Ec2MetadataServiceConnector.create(new SnitchProperties(p)); + }).withMessage(format("property %s was set to 50000 seconds which is not in allowed range of [%s..%s]", + AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, + MIN_TOKEN_TIME_IN_SECONDS, + MAX_TOKEN_TIME_IN_SECONDS)); + + assertThatExceptionOfType(ConfigurationException.class) + .describedAs("ttl value should be integer") + .isThrownBy(() -> { + Properties p = new Properties(); + p.setProperty(EC2_METADATA_TYPE_PROPERTY, EC2MetadataType.v2.name()); + p.setProperty(AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, "abc"); + + Ec2MetadataServiceConnector.create(new SnitchProperties(p)); + }).withMessage(format("Unable to parse integer from property %s, value to parse: abc", + AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY)); + } +} diff --git a/test/unit/org/apache/cassandra/locator/EC2SnitchTest.java b/test/unit/org/apache/cassandra/locator/Ec2SnitchTest.java similarity index 61% rename from test/unit/org/apache/cassandra/locator/EC2SnitchTest.java rename to test/unit/org/apache/cassandra/locator/Ec2SnitchTest.java index 78c5d40dfc..abd9bb0989 100644 --- a/test/unit/org/apache/cassandra/locator/EC2SnitchTest.java +++ b/test/unit/org/apache/cassandra/locator/Ec2SnitchTest.java @@ -18,38 +18,41 @@ package org.apache.cassandra.locator; - -import java.io.IOException; import java.util.Collections; import java.util.EnumMap; import java.util.HashSet; import java.util.Map; +import java.util.Properties; import java.util.Set; import org.junit.AfterClass; -import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; import org.apache.cassandra.config.DatabaseDescriptor; import org.apache.cassandra.db.Keyspace; import org.apache.cassandra.db.commitlog.CommitLog; -import org.apache.cassandra.exceptions.ConfigurationException; import org.apache.cassandra.gms.ApplicationState; import org.apache.cassandra.gms.Gossiper; import org.apache.cassandra.gms.VersionedValue; import org.apache.cassandra.service.StorageService; +import org.mockito.stubbing.Answer; import static org.apache.cassandra.ServerTestUtils.cleanup; import static org.apache.cassandra.ServerTestUtils.mkdirs; import static org.apache.cassandra.config.CassandraRelevantProperties.GOSSIP_DISABLE_THREAD_VALIDATION; +import static org.apache.cassandra.locator.Ec2MultiRegionSnitch.PRIVATE_IP_QUERY; +import static org.apache.cassandra.locator.Ec2MultiRegionSnitch.PUBLIC_IP_QUERY; import static org.apache.cassandra.locator.Ec2Snitch.EC2_NAMING_LEGACY; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; -public class EC2SnitchTest +public class Ec2SnitchTest { - private static String az; - private final SnitchProperties legacySnitchProps = new SnitchProperties() { public String get(String propertyName, String defaultValue) @@ -71,30 +74,179 @@ public class EC2SnitchTest StorageService.instance.initServer(0); } - private class TestEC2Snitch extends Ec2Snitch + + @AfterClass + public static void tearDown() { - public TestEC2Snitch() throws IOException, ConfigurationException - { - super(); - } - - public TestEC2Snitch(SnitchProperties props) throws IOException, ConfigurationException - { - super(props); - } - - @Override - String awsApiCall(String url) throws IOException, ConfigurationException - { - return az; - } + StorageService.instance.stopClient(); } @Test - public void testLegacyRac() throws IOException, ConfigurationException + public void testLegacyRac() throws Exception + { + Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class); + when(connectorMock.apiCall(anyString())).thenReturn("us-east-1d"); + Ec2Snitch snitch = new Ec2Snitch(legacySnitchProps, connectorMock); + testLegacyRacInternal(snitch); + } + + @Test + public void testMultiregionLegacyRac() throws Exception + { + Ec2MetadataServiceConnector multiRegionConnectorMock = mock(Ec2MetadataServiceConnector.class); + when(multiRegionConnectorMock.apiCall(anyString())).then((Answer) invocation -> { + String query = invocation.getArgument(0); + return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-1d"; + }); + + Ec2Snitch snitch = new Ec2MultiRegionSnitch(legacySnitchProps, multiRegionConnectorMock); + testLegacyRacInternal(snitch); + } + + @Test + public void testLegacyNewRegions() throws Exception + { + Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class); + when(connectorMock.apiCall(anyString())).thenReturn("us-east-2d"); + testLegacyNewRegionsInternal(new Ec2Snitch(legacySnitchProps, connectorMock)); + } + + @Test + public void testLegacyMultiRegionNewRegions() throws Exception + { + Ec2MetadataServiceConnector multiRegionConnectorMock = mock(Ec2MetadataServiceConnector.class); + when(multiRegionConnectorMock.apiCall(anyString())).then((Answer) invocation -> { + String query = invocation.getArgument(0); + return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-2d"; + }); + + testLegacyNewRegionsInternal(new Ec2MultiRegionSnitch(legacySnitchProps, multiRegionConnectorMock)); + } + + + @Test + public void testFullNamingScheme() throws Exception + { + Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class); + when(connectorMock.apiCall(anyString())).thenReturn("us-east-2d"); + Ec2Snitch snitch = new Ec2Snitch(new SnitchProperties(new Properties()), connectorMock); + + InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1"); + + assertEquals("us-east-2", snitch.getDatacenter(local)); + assertEquals("us-east-2d", snitch.getRack(local)); + + Ec2MetadataServiceConnector multiRegionConnectorMock = mock(Ec2MetadataServiceConnector.class); + when(multiRegionConnectorMock.apiCall(anyString())).then((Answer) invocation -> { + String query = invocation.getArgument(0); + return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-2d"; + }); + + assertEquals("us-east-2", snitch.getDatacenter(local)); + assertEquals("us-east-2d", snitch.getRack(local)); + } + + @Test + public void validateDatacenter_RequiresLegacy_CorrectAmazonName() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-east-1"); + assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true)); + } + + @Test + public void validateDatacenter_RequiresLegacy_LegacyName() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-east"); + assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true)); + } + + @Test + public void validate_RequiresLegacy_HappyPath() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-east"); + Set racks = new HashSet<>(); + racks.add("1a"); + assertTrue(Ec2Snitch.validate(datacenters, racks, true)); + } + + @Test + public void validate_RequiresLegacy_HappyPathWithDCSuffix() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-east_CUSTOM_SUFFIX"); + Set racks = new HashSet<>(); + racks.add("1a"); + assertTrue(Ec2Snitch.validate(datacenters, racks, true)); + } + + @Test + public void validateRack_RequiresAmazonName_CorrectAmazonName() + { + Set racks = new HashSet<>(); + racks.add("us-east-1a"); + assertTrue(Ec2Snitch.validate(Collections.emptySet(), racks, false)); + } + + @Test + public void validateRack_RequiresAmazonName_LegacyName() + { + Set racks = new HashSet<>(); + racks.add("1a"); + assertFalse(Ec2Snitch.validate(Collections.emptySet(), racks, false)); + } + + @Test + public void validate_RequiresAmazonName_HappyPath() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-east-1"); + Set racks = new HashSet<>(); + racks.add("us-east-1a"); + assertTrue(Ec2Snitch.validate(datacenters, racks, false)); + } + + @Test + public void validate_RequiresAmazonName_HappyPathWithDCSuffix() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-east-1_CUSTOM_SUFFIX"); + Set racks = new HashSet<>(); + racks.add("us-east-1a"); + assertTrue(Ec2Snitch.validate(datacenters, racks, false)); + } + + /** + * Validate upgrades in legacy mode for regions that didn't change name between the standard and legacy modes. + */ + @Test + public void validate_RequiresLegacy_DCValidStandardAndLegacy() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-west-2"); + Set racks = new HashSet<>(); + racks.add("2a"); + racks.add("2b"); + assertTrue(Ec2Snitch.validate(datacenters, racks, true)); + } + + /** + * Check that racks names are enough to detect a mismatch in naming conventions. + */ + @Test + public void validate_RequiresLegacy_RackInvalidForLegacy() + { + Set datacenters = new HashSet<>(); + datacenters.add("us-west-2"); + Set racks = new HashSet<>(); + racks.add("us-west-2a"); + assertFalse(Ec2Snitch.validate(datacenters, racks, true)); + } + + private void testLegacyRacInternal(Ec2Snitch snitch) throws Exception { - az = "us-east-1d"; - Ec2Snitch snitch = new TestEC2Snitch(legacySnitchProps); InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1"); InetAddressAndPort nonlocal = InetAddressAndPort.getByName("127.0.0.7"); @@ -111,135 +263,10 @@ public class EC2SnitchTest assertEquals("1d", snitch.getRack(local)); } - @Test - public void testLegacyNewRegions() throws IOException, ConfigurationException + private void testLegacyNewRegionsInternal(Ec2Snitch snitch) throws Exception { - az = "us-east-2d"; - Ec2Snitch snitch = new TestEC2Snitch(legacySnitchProps); InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1"); assertEquals("us-east-2", snitch.getDatacenter(local)); assertEquals("2d", snitch.getRack(local)); } - - @Test - public void testFullNamingScheme() throws IOException, ConfigurationException - { - InetAddressAndPort local = InetAddressAndPort.getByName("127.0.0.1"); - az = "us-east-2d"; - Ec2Snitch snitch = new TestEC2Snitch(); - - assertEquals("us-east-2", snitch.getDatacenter(local)); - assertEquals("us-east-2d", snitch.getRack(local)); - - az = "us-west-1a"; - snitch = new TestEC2Snitch(); - - assertEquals("us-west-1", snitch.getDatacenter(local)); - assertEquals("us-west-1a", snitch.getRack(local)); - } - - @Test - public void validateDatacenter_RequiresLegacy_CorrectAmazonName() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-east-1"); - Assert.assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true)); - } - - @Test - public void validateDatacenter_RequiresLegacy_LegacyName() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-east"); - Assert.assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true)); - } - - @Test - public void validate_RequiresLegacy_HappyPath() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-east"); - Set racks = new HashSet<>(); - racks.add("1a"); - Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, true)); - } - - @Test - public void validate_RequiresLegacy_HappyPathWithDCSuffix() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-east_CUSTOM_SUFFIX"); - Set racks = new HashSet<>(); - racks.add("1a"); - Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, true)); - } - - @Test - public void validateRack_RequiresAmazonName_CorrectAmazonName() - { - Set racks = new HashSet<>(); - racks.add("us-east-1a"); - Assert.assertTrue(Ec2Snitch.validate(Collections.emptySet(), racks, false)); - } - - @Test - public void validateRack_RequiresAmazonName_LegacyName() - { - Set racks = new HashSet<>(); - racks.add("1a"); - Assert.assertFalse(Ec2Snitch.validate(Collections.emptySet(), racks, false)); - } - - @Test - public void validate_RequiresAmazonName_HappyPath() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-east-1"); - Set racks = new HashSet<>(); - racks.add("us-east-1a"); - Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, false)); - } - - @Test - public void validate_RequiresAmazonName_HappyPathWithDCSuffix() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-east-1_CUSTOM_SUFFIX"); - Set racks = new HashSet<>(); - racks.add("us-east-1a"); - Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, false)); - } - - /** - * Validate upgrades in legacy mode for regions that didn't change name between the standard and legacy modes. - */ - @Test - public void validate_RequiresLegacy_DCValidStandardAndLegacy() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-west-2"); - Set racks = new HashSet<>(); - racks.add("2a"); - racks.add("2b"); - Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, true)); - } - - /** - * Check that racks names are enough to detect a mismatch in naming conventions. - */ - @Test - public void validate_RequiresLegacy_RackInvalidForLegacy() - { - Set datacenters = new HashSet<>(); - datacenters.add("us-west-2"); - Set racks = new HashSet<>(); - racks.add("us-west-2a"); - Assert.assertFalse(Ec2Snitch.validate(datacenters, racks, true)); - } - - @AfterClass - public static void tearDown() - { - StorageService.instance.stopClient(); - } } diff --git a/test/unit/org/apache/cassandra/locator/Ec2V2ConnectorMockingTest.java b/test/unit/org/apache/cassandra/locator/Ec2V2ConnectorMockingTest.java new file mode 100644 index 0000000000..8fdbdfa093 --- /dev/null +++ b/test/unit/org/apache/cassandra/locator/Ec2V2ConnectorMockingTest.java @@ -0,0 +1,169 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.cassandra.locator; + +import java.util.Properties; + +import org.junit.Rule; +import org.junit.Test; + +import com.github.tomakehurst.wiremock.client.MappingBuilder; +import com.github.tomakehurst.wiremock.junit.WireMockRule; +import org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.equalTo; +import static com.github.tomakehurst.wiremock.client.WireMock.get; +import static com.github.tomakehurst.wiremock.client.WireMock.put; +import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.EC2MetadataType.v2; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.EC2_METADATA_TYPE_PROPERTY; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V1Connector.EC2_METADATA_URL_PROPERTY; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector.AWS_EC2_METADATA_TOKEN_HEADER; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector.AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector.AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY; +import static org.apache.cassandra.locator.Ec2MetadataServiceConnector.V2Connector.TOKEN_QUERY; +import static org.apache.cassandra.locator.Ec2Snitch.ZONE_NAME_QUERY; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +public class Ec2V2ConnectorMockingTest +{ + private static final String token = "thisismytoken"; + private static final String az = "us-east-1a"; + + @Rule + public WireMockRule v2Service = new WireMockRule(wireMockConfig().bindAddress("127.0.0.1").port(8080)); + + @Test + public void testV2Connector() throws Throwable + { + v2Service.stubFor(tokenRequest(100, 200, token)); + v2Service.stubFor(azRequest(az, 200, token)); + + assertEquals(az, getConnector(100).apiCall(ZONE_NAME_QUERY)); + } + + @Test + public void testV2ConnectorWhenUnauthorized() + { + String unauthorizedBody = "\n" + + "\n" + + "\n" + + " \n" + + " 401 - Unauthorized\n" + + " \n" + + " \n" + + "

401 - Unauthorized

\n" + + " \n" + + "\n"; + + v2Service.stubFor(tokenRequest(100, 200, token)); + + v2Service.stubFor(get(urlEqualTo(ZONE_NAME_QUERY)).withHeader(AWS_EC2_METADATA_TOKEN_HEADER, equalTo(token)) + .willReturn(aResponse().withStatus(401) + .withStatusMessage("Unauthorized") + .withHeader("Content-Type", "text/html") + .withHeader("Content-Length", String.valueOf(unauthorizedBody.getBytes(UTF_8).length)))); + + V2Connector.HTTP_REQUEST_RETRIES = 0; + Ec2MetadataServiceConnector v2Connector = getConnector(100); + + assertThatExceptionOfType(AbstractCloudMetadataServiceConnector.HttpException.class) + .isThrownBy(() -> v2Connector.apiCall(ZONE_NAME_QUERY)) + .matches(ex -> ex.responseCode == 401 && "Unauthorized".equals(ex.responseMessage), + "exception should have response code 401 with response message 'Unauthorized'"); + } + + @Test + public void testCachedToken() throws Throwable + { + v2Service.stubFor(tokenRequest(30, 200, token)); + v2Service.stubFor(azRequest(az, 200, token)); + + V2Connector spiedConnector = (V2Connector) spy(getConnector(30)); + + spiedConnector.apiCall(ZONE_NAME_QUERY); + verify(spiedConnector, times(1)).getToken(); + + // lets wait 10 seconds and make a call again + // which will use cached token and will not call token endpoint + // for the second time + Thread.sleep(10000); + + // as token is not expired yet, another call to getToken will not be done + spiedConnector.apiCall(ZONE_NAME_QUERY); + + // here we still have just 1 call to getToken method because we used a cached token + verify(spiedConnector, times(1)).getToken(); + } + + @Test + public void testExpiredTokenInteraction() throws Throwable + { + v2Service.stubFor(tokenRequest(30, 200, token)); + v2Service.stubFor(azRequest(az, 200, token)); + + V2Connector spiedConnector = (V2Connector) spy(getConnector(30)); + + spiedConnector.apiCall(ZONE_NAME_QUERY); + verify(spiedConnector, times(1)).getToken(); + + // lets expire the token + Thread.sleep(35000); + + // as token is expired, another call to getToken will be done + spiedConnector.apiCall(ZONE_NAME_QUERY); + verify(spiedConnector, times(2)).getToken(); + } + + private Ec2MetadataServiceConnector getConnector(int tokenTTL) + { + Properties p = new Properties(); + p.setProperty(EC2_METADATA_TYPE_PROPERTY, v2.toString()); + p.setProperty(EC2_METADATA_URL_PROPERTY, "http://127.0.0.1:8080"); + p.setProperty(AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER_PROPERTY, String.valueOf(tokenTTL)); + + return Ec2MetadataServiceConnector.create(new SnitchProperties(p)); + } + + private MappingBuilder tokenRequest(int ttl, int status, String tokenToReturn) + { + return put(urlEqualTo(TOKEN_QUERY)).withHeader(AWS_EC2_METADATA_TOKEN_TTL_SECONDS_HEADER, equalTo(String.valueOf(ttl))) + .willReturn(aResponse().withBody(tokenToReturn) + .withStatus(status) + .withHeader("Content-Type", "text/plain") + .withHeader("Content-Length", String.valueOf(tokenToReturn.getBytes(UTF_8).length))); + } + + private MappingBuilder azRequest(String az, int status, String token) + { + return get(urlEqualTo(ZONE_NAME_QUERY)).withHeader(AWS_EC2_METADATA_TOKEN_HEADER, equalTo(token)) + .willReturn(aResponse().withBody(az) + .withStatus(status) + .withHeader("Content-Type", "text/plain") + .withHeader("Content-Length", String.valueOf(az.getBytes(UTF_8).length))); + } +}