From b8e21fb80afa6696374292dc99a480fac9788e64 Mon Sep 17 00:00:00 2001
From: Stefan Miklosovic <smiklosovic@apache.org>
Date: Wed, 17 May 2023 12:00:19 +0200
Subject: [PATCH] Validate the existence of a datacenter in nodetool rebuild

patch by Stefan Miklosovic; reviewed by Brandon Williams for CASSANDRA-14319
---
 CHANGES.txt                                   |  1 +
 .../cassandra/service/StorageService.java     | 25 +++++++++--
 .../service/StorageServiceServerTest.java     | 42 +++++++++++++++++++
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 2711ff24f5..a9fd724d6a 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,4 +1,5 @@
 3.0.30
+ * Validate the existence of a datacenter in nodetool rebuild (CASSANDRA-14319)
 
 3.0.29
  * Suppress CVE-2023-2251 (CASSANDRA-18497)
diff --git a/src/java/org/apache/cassandra/service/StorageService.java b/src/java/org/apache/cassandra/service/StorageService.java
index 90c1fde3af..9a37593de0 100644
--- a/src/java/org/apache/cassandra/service/StorageService.java
+++ b/src/java/org/apache/cassandra/service/StorageService.java
@@ -1308,10 +1308,29 @@ public class StorageService extends NotificationBroadcasterSupport implements IE
 
     public void rebuild(String sourceDc)
     {
-        // check ongoing rebuild
-        if (!isRebuilding.compareAndSet(false, true))
+        try
         {
-            throw new IllegalStateException("Node is still rebuilding. Check nodetool netstats.");
+            // check ongoing rebuild
+            if (!isRebuilding.compareAndSet(false, true))
+            {
+                throw new IllegalStateException("Node is still rebuilding. Check nodetool netstats.");
+            }
+
+            if (sourceDc != null)
+            {
+                TokenMetadata.Topology topology = getTokenMetadata().cloneOnlyTokenMap().getTopology();
+                Set<String> availableDCs = topology.getDatacenterEndpoints().keySet();
+                if (!availableDCs.contains(sourceDc))
+                {
+                    throw new IllegalArgumentException(String.format("Provided datacenter '%s' is not a valid datacenter, available datacenters are: %s",
+                                                                     sourceDc, String.join(",", availableDCs)));
+                }
+            }
+        }
+        catch (Throwable ex)
+        {
+            isRebuilding.set(false);
+            throw ex;
         }
 
         logger.info("rebuild from dc: {}", sourceDc == null ? "(any dc)" : sourceDc);
diff --git a/test/unit/org/apache/cassandra/service/StorageServiceServerTest.java b/test/unit/org/apache/cassandra/service/StorageServiceServerTest.java
index 95c2291354..a939ef79d2 100644
--- a/test/unit/org/apache/cassandra/service/StorageServiceServerTest.java
+++ b/test/unit/org/apache/cassandra/service/StorageServiceServerTest.java
@@ -27,6 +27,7 @@ import java.net.InetAddress;
 import java.util.*;
 
 import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.Multimap;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -52,9 +53,11 @@ import org.apache.cassandra.schema.KeyspaceParams;
 import org.apache.cassandra.schema.ReplicationParams;
 import org.apache.cassandra.schema.SchemaKeyspace;
 import org.apache.cassandra.utils.FBUtilities;
+import org.mockito.Mockito;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
 
 @RunWith(OrderedJUnit4ClassRunner.class)
@@ -590,4 +593,43 @@ public class StorageServiceServerTest
     {
         assert StorageService.instance.getConcurrentCompactors() == 4;
     }
+
+    @Test
+    public void testRebuildFailOnNonExistingDatacenter() throws Exception
+    {
+        String nonExistentDC = "NON_EXISTENT_DC";
+
+        try
+        {
+            getStorageService().rebuild(nonExistentDC);
+            fail();
+        }
+        catch (IllegalArgumentException ex)
+        {
+            assertEquals(String.format("Provided datacenter '%s' is not a valid datacenter, available datacenters are: %s",
+                                       nonExistentDC,
+                                       "datacenter1"),
+                         ex.getMessage());
+        }
+    }
+
+    private StorageService getStorageService() throws Exception
+    {
+        ImmutableMultimap.Builder<String, InetAddress> builder = ImmutableMultimap.builder();
+        builder.put("datacenter1", InetAddress.getByName("127.0.0.1"));
+
+        TokenMetadata.Topology tokenMetadataTopology = Mockito.mock(TokenMetadata.Topology.class);
+        Mockito.when(tokenMetadataTopology.getDatacenterEndpoints()).thenReturn(builder.build());
+
+        TokenMetadata metadata = new TokenMetadata();
+        TokenMetadata spiedMetadata = Mockito.spy(metadata);
+
+        Mockito.when(spiedMetadata.getTopology()).thenReturn(tokenMetadataTopology);
+
+        StorageService spiedStorageService = Mockito.spy(StorageService.instance);
+        Mockito.when(spiedStorageService.getTokenMetadata()).thenReturn(spiedMetadata);
+        Mockito.when(spiedMetadata.cloneOnlyTokenMap()).thenReturn(spiedMetadata);
+
+        return spiedStorageService;
+    }
 }