Merge branch 'cassandra-4.0' into cassandra-4.1

CHANGES.txt

@@ -10,6 +10,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)

NEWS.txt

@@ -51,6 +51,15 @@ restore snapshots created with the previous major version using the
 'sstableloader' tool. You can upgrade the file format of your snapshots
 using the provided 'sstableupgrade' tool.
 
+4.1.3
+=====
+
+Upgrading
+---------
+    - 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)
 
 4.1.1
 =====

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

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<String, String> 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;
+        }
+    }
+}

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<SnitchProperties, Ec2MetadataServiceConnector> provider;
+
+        EC2MetadataType(Function<SnitchProperties, Ec2MetadataServiceConnector> 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<String, Long> 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<String, String> extraHeaders, int expectedResponseCode) throws IOException
+        {
+            Map<String, String> 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.
+         * <a href="https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html">IMDSv2</a>
+         */
+        @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);
+            }
+        }
+    }
+}

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";
+    @VisibleForTesting
-    private static final String PRIVATE_IP_QUERY_URL = "http://169.254.169.254/latest/meta-data/local-ipv4";
+    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();
+        this(new SnitchProperties());
-        InetAddress localPublicAddress = InetAddress.getByName(awsApiCall(PUBLIC_IP_QUERY_URL));
+    }
+
+    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)

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<InetAddressAndPort, Map<String, String>> 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()))

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;
@@ -59,6 +61,12 @@ public class SnitchProperties
         }
     }
 
+    @VisibleForTesting
+    public SnitchProperties(Properties properties)
+    {
+        this.properties = properties;
+    }
+
     /**
      * Get a snitch property value or return defaultValue if not defined.
      */

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));
+    }
+}

test/unit/org/apache/cassandra/locator/Ec2SnitchTest.java

@@ -18,37 +18,40 @@
 
 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.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)
@@ -70,30 +73,179 @@ public class EC2SnitchTest
         StorageService.instance.initServer(0);
     }
 
-    private class TestEC2Snitch extends Ec2Snitch
+
+    @AfterClass
+    public static void tearDown()
     {
-        public TestEC2Snitch() throws IOException, ConfigurationException
+        StorageService.instance.stopClient();
-        {
-            super();
-        }
-
-        public TestEC2Snitch(SnitchProperties props) throws IOException, ConfigurationException
-        {
-            super(props);
-        }
-
-        @Override
-        String awsApiCall(String url) throws IOException, ConfigurationException
-        {
-            return az;
-        }
     }
 
     @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<String>) 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<String>) 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<String>) 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<String> datacenters = new HashSet<>();
+        datacenters.add("us-east-1");
+        assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true));
+    }
+
+    @Test
+    public void validateDatacenter_RequiresLegacy_LegacyName()
+    {
+        Set<String> datacenters = new HashSet<>();
+        datacenters.add("us-east");
+        assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true));
+    }
+
+    @Test
+    public void validate_RequiresLegacy_HappyPath()
+    {
+        Set<String> datacenters = new HashSet<>();
+        datacenters.add("us-east");
+        Set<String> racks = new HashSet<>();
+        racks.add("1a");
+        assertTrue(Ec2Snitch.validate(datacenters, racks, true));
+    }
+
+    @Test
+    public void validate_RequiresLegacy_HappyPathWithDCSuffix()
+    {
+        Set<String> datacenters = new HashSet<>();
+        datacenters.add("us-east_CUSTOM_SUFFIX");
+        Set<String> racks = new HashSet<>();
+        racks.add("1a");
+        assertTrue(Ec2Snitch.validate(datacenters, racks, true));
+    }
+
+    @Test
+    public void validateRack_RequiresAmazonName_CorrectAmazonName()
+    {
+        Set<String> racks = new HashSet<>();
+        racks.add("us-east-1a");
+        assertTrue(Ec2Snitch.validate(Collections.emptySet(), racks, false));
+    }
+
+    @Test
+    public void validateRack_RequiresAmazonName_LegacyName()
+    {
+        Set<String> racks = new HashSet<>();
+        racks.add("1a");
+        assertFalse(Ec2Snitch.validate(Collections.emptySet(), racks, false));
+    }
+
+    @Test
+    public void validate_RequiresAmazonName_HappyPath()
+    {
+        Set<String> datacenters = new HashSet<>();
+        datacenters.add("us-east-1");
+        Set<String> racks = new HashSet<>();
+        racks.add("us-east-1a");
+        assertTrue(Ec2Snitch.validate(datacenters, racks, false));
+    }
+
+    @Test
+    public void validate_RequiresAmazonName_HappyPathWithDCSuffix()
+    {
+        Set<String> datacenters = new HashSet<>();
+        datacenters.add("us-east-1_CUSTOM_SUFFIX");
+        Set<String> 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<String> datacenters = new HashSet<>();
+        datacenters.add("us-west-2");
+        Set<String> 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<String> datacenters = new HashSet<>();
+        datacenters.add("us-west-2");
+        Set<String> 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");
 
@@ -110,135 +262,10 @@ public class EC2SnitchTest
         assertEquals("1d", snitch.getRack(local));
     }
 
-    @Test
+    private void testLegacyNewRegionsInternal(Ec2Snitch snitch) throws Exception
-    public void testLegacyNewRegions() throws IOException, ConfigurationException
     {
-        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<String> datacenters = new HashSet<>();
-        datacenters.add("us-east-1");
-        Assert.assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true));
-    }
-
-    @Test
-    public void validateDatacenter_RequiresLegacy_LegacyName()
-    {
-        Set<String> datacenters = new HashSet<>();
-        datacenters.add("us-east");
-        Assert.assertTrue(Ec2Snitch.validate(datacenters, Collections.emptySet(), true));
-    }
-
-    @Test
-    public void validate_RequiresLegacy_HappyPath()
-    {
-        Set<String> datacenters = new HashSet<>();
-        datacenters.add("us-east");
-        Set<String> racks = new HashSet<>();
-        racks.add("1a");
-        Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, true));
-    }
-
-    @Test
-    public void validate_RequiresLegacy_HappyPathWithDCSuffix()
-    {
-        Set<String> datacenters = new HashSet<>();
-        datacenters.add("us-east_CUSTOM_SUFFIX");
-        Set<String> racks = new HashSet<>();
-        racks.add("1a");
-        Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, true));
-    }
-
-    @Test
-    public void validateRack_RequiresAmazonName_CorrectAmazonName()
-    {
-        Set<String> racks = new HashSet<>();
-        racks.add("us-east-1a");
-        Assert.assertTrue(Ec2Snitch.validate(Collections.emptySet(), racks, false));
-    }
-
-    @Test
-    public void validateRack_RequiresAmazonName_LegacyName()
-    {
-        Set<String> racks = new HashSet<>();
-        racks.add("1a");
-        Assert.assertFalse(Ec2Snitch.validate(Collections.emptySet(), racks, false));
-    }
-
-    @Test
-    public void validate_RequiresAmazonName_HappyPath()
-    {
-        Set<String> datacenters = new HashSet<>();
-        datacenters.add("us-east-1");
-        Set<String> racks = new HashSet<>();
-        racks.add("us-east-1a");
-        Assert.assertTrue(Ec2Snitch.validate(datacenters, racks, false));
-    }
-
-    @Test
-    public void validate_RequiresAmazonName_HappyPathWithDCSuffix()
-    {
-        Set<String> datacenters = new HashSet<>();
-        datacenters.add("us-east-1_CUSTOM_SUFFIX");
-        Set<String> 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<String> datacenters = new HashSet<>();
-        datacenters.add("us-west-2");
-        Set<String> 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<String> datacenters = new HashSet<>();
-        datacenters.add("us-west-2");
-        Set<String> racks = new HashSet<>();
-        racks.add("us-west-2a");
-        Assert.assertFalse(Ec2Snitch.validate(datacenters, racks, true));
-    }
-
-    @AfterClass
-    public static void tearDown()
-    {
-        StorageService.instance.stopClient();
-    }
 }