test/unit/org/apache/cassandra/service/PartitionDenylistTest.java - cassandra - Git at Google

 /*
  * 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.service;

 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;

 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.cql3.CQLTester;
 import org.apache.cassandra.cql3.UntypedResultSet;
 import org.apache.cassandra.cql3.statements.schema.CreateTableStatement;
 import org.apache.cassandra.db.ConsistencyLevel;
 import org.apache.cassandra.exceptions.InvalidRequestException;
 import org.apache.cassandra.schema.KeyspaceMetadata;
 import org.apache.cassandra.schema.KeyspaceParams;
 import org.apache.cassandra.schema.SchemaTestUtil;
 import org.apache.cassandra.schema.Tables;

 import static org.apache.cassandra.cql3.QueryProcessor.process;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;

 public class PartitionDenylistTest
 {
     private final static String ks_cql = "partition_denylist_keyspace";

     @BeforeClass
     public static void init()
     {
         CQLTester.prepareServer();

         KeyspaceMetadata schema = KeyspaceMetadata.create(ks_cql,
                                                           KeyspaceParams.simple(1),
                                                           Tables.of(
             CreateTableStatement.parse("CREATE TABLE table1 ("
                                        + "keyone text, "
                                        + "keytwo text, "
                                        + "qux text, "
                                        + "quz text, "
                                        + "foo text, "
                                        + "PRIMARY KEY((keyone, keytwo), qux, quz) ) ", ks_cql).build(),
             CreateTableStatement.parse("CREATE TABLE table2 ("
                                        + "keyone text, "
                                        + "keytwo text, "
                                        + "keythree text, "
                                        + "value text, "
                                        + "PRIMARY KEY((keyone, keytwo), keythree) ) ", ks_cql).build(),
             CreateTableStatement.parse("CREATE TABLE table3 ("
                                        + "keyone text, "
                                        + "keytwo text, "
                                        + "keythree text, "
                                        + "value text, "
                                        + "PRIMARY KEY((keyone, keytwo), keythree) ) ", ks_cql).build()
         ));
         SchemaTestUtil.addOrUpdateKeyspace(schema, false);
         DatabaseDescriptor.setPartitionDenylistEnabled(true);
         DatabaseDescriptor.setDenylistRangeReadsEnabled(true);
         DatabaseDescriptor.setDenylistConsistencyLevel(ConsistencyLevel.ONE);
         DatabaseDescriptor.setDenylistRefreshSeconds(1);
         StorageService.instance.initServer(0);
     }

     @Before
     public void setup()
     {
         DatabaseDescriptor.setPartitionDenylistEnabled(true);
         resetDenylist();

         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('aaa', 'bbb', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('ccc', 'ddd', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('ddd', 'eee', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('eee', 'fff', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('fff', 'ggg', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('ggg', 'hhh', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('hhh', 'iii', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('iii', 'jjj', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('jjj', 'kkk', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);


         for (int i = 0; i < 50; i++)
             process(String.format("INSERT INTO " + ks_cql + ".table2 (keyone, keytwo, keythree, value) VALUES ('%d', '%d', '%d', '%d')", i, i, i, i), ConsistencyLevel.ONE);

         for (int i = 0; i < 50; i++)
             process(String.format("INSERT INTO " + ks_cql + ".table3 (keyone, keytwo, keythree, value) VALUES ('%d', '%d', '%d', '%d')", i, i, i, i), ConsistencyLevel.ONE);

         denylist("table1", "bbb:ccc");
         refreshList();
     }


     private static void denylist(String table, final String key)
     {
         StorageProxy.instance.denylistKey(ks_cql, table, key);
     }

     private static void refreshList()
     {
         StorageProxy.instance.loadPartitionDenylist();
     }

     /**
      * @return Whether the *attempt* to remove the denylisted key and refresh succeeded. Doesn't necessarily indicate the key
      * was previously blocked and found.
      */
     private static boolean removeDenylist(final String ks, final String table, final String key)
     {
         return StorageProxy.instance.removeDenylistKey(ks, table, key);
     }

     @Test
     public void testRead()
     {
         process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='aaa' and keytwo='bbb'", ConsistencyLevel.ONE);
     }

     @Test
     public void testReadDenylisted()
     {
         assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='bbb' and keytwo='ccc'", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Unable to read denylisted partition");
     }

     @Test
     public void testIsKeyDenylistedAPI()
     {
         Assert.assertTrue(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));
         resetDenylist();
         Assert.assertFalse(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));

         // Confirm an add mutates cache state
         denylist("table1", "bbb:ccc");
         Assert.assertTrue(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));

         // Confirm removal then mutates cache w/out explicit reload
         StorageProxy.instance.removeDenylistKey(ks_cql, "table1", "bbb:ccc");
         Assert.assertFalse(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));
     }

     @Test
     public void testReadUndenylisted()
     {
         resetDenylist();
         process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='ccc' and keytwo='ddd'", ConsistencyLevel.ONE);
     }

     @Test
     public void testWrite()
     {
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('eee', 'fff', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
         process("DELETE FROM " + ks_cql + ".table1 WHERE keyone='eee' and keytwo='fff'", ConsistencyLevel.ONE);
     }

     @Test
     public void testWriteDenylisted()
     {
         assertThatThrownBy(() -> process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'eee', 'fff', 'w')", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Unable to write to denylisted partition");
     }

     @Test
     public void testCASWriteDenylisted()
     {
         assertThatThrownBy(() -> process("UPDATE " + ks_cql + ".table1 SET foo='w' WHERE keyone='bbb' AND keytwo='ccc' AND qux='eee' AND quz='fff' IF foo='v'", ConsistencyLevel.LOCAL_SERIAL))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Unable to CAS write to denylisted partition");
     }

     @Test
     public void testWriteUndenylisted()
     {
         resetDenylist();
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'eee', 'fff', 'w')", ConsistencyLevel.ONE);
     }

     @Test
     public void testRangeSlice()
     {
         UntypedResultSet rows;
         rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) < token('bbb', 'ccc')", ConsistencyLevel.ONE);
         Assert.assertEquals(1, rows.size());

         // 10 entries total in our table
         rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('bbb', 'ccc')", ConsistencyLevel.ONE);
         Assert.assertEquals(8, rows.size());

         rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) >= token('aaa', 'bbb') and token(keyone, keytwo) < token('bbb', 'ccc')", ConsistencyLevel.ONE);
         Assert.assertEquals(1, rows.size());

         rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('bbb', 'ccc') and token(keyone, keytwo) <= token('ddd', 'eee')", ConsistencyLevel.ONE);
         Assert.assertEquals(2, rows.size());
     }

     @Test
     public void testRangeDenylisted()
     {
         assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
     }

     @Test
     public void testRangeDenylisted2()
     {
         assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) >= token('aaa', 'bbb') and token (keyone, keytwo) <= token('bbb', 'ccc')", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
     }

     @Test
     public void testRangeDenylisted3()
     {
         assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) >= token('bbb', 'ccc') and token (keyone, keytwo) <= token('ccc', 'ddd')", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
     }

     @Test
     public void testRangeDenylisted4()
     {
         assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('aaa', 'bbb') and token (keyone, keytwo) < token('ccc', 'ddd')", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
     }

     @Test
     public void testRangeDenylisted5()
     {
         assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('aaa', 'bbb')", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
     }

     @Test
     public void testRangeDenylisted6()
     {
         assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) < token('ddd', 'eee')", ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
     }

     @Test
     public void testInsertUnknownPKIsGraceful()
     {
         Assert.assertTrue(StorageProxy.instance.denylistKey(ks_cql, "table1", "hohoho"));
     }

     @Test
     public void testInsertInvalidTableIsGraceful()
     {
         Assert.assertFalse(StorageProxy.instance.denylistKey(ks_cql, "asldkfjadlskjf", "alksdjfads"));
     }

     @Test
     public void testInsertInvalidKSIsGraceful()
     {
         Assert.assertFalse(StorageProxy.instance.denylistKey("asdklfjas", "asldkfjadlskjf", "alksdjfads"));
     }

     @Test
     public void testDisabledDenylistThrowsNoExceptions()
     {
         process(String.format("TRUNCATE TABLE %s.table2", ks_cql), ConsistencyLevel.ONE);
         process(String.format("TRUNCATE TABLE %s.table3", ks_cql), ConsistencyLevel.ONE);
         denyAllKeys();
         DatabaseDescriptor.setPartitionDenylistEnabled(false);
         process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'eee', 'fff', 'w')", ConsistencyLevel.ONE);
         process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='bbb' and keytwo='ccc'", ConsistencyLevel.ONE);
         process("SELECT * FROM " + ks_cql + ".table1", ConsistencyLevel.ONE);

         for (int i = 0; i < 50; i++)
         {
             process(String.format("INSERT INTO %s.table2 (keyone, keytwo, keythree, value) VALUES ('%s', '%s', '%s', '%s')", ks_cql, i, i, i, i), ConsistencyLevel.ONE);
             process(String.format("SELECT * FROM %s.table2 WHERE keyone='%s' and keytwo='%s'", ks_cql, i, i), ConsistencyLevel.ONE);
         }

         for (int i = 0; i < 50; i++)
         {
             process(String.format("INSERT INTO %s.table3 (keyone, keytwo, keythree, value) VALUES ('%s', '%s', '%s', '%s')", ks_cql, i, i, i, i), ConsistencyLevel.ONE);
             process(String.format("SELECT * FROM %s.table3 WHERE keyone='%s' and keytwo='%s'", ks_cql, i, i), ConsistencyLevel.ONE);
         }
     }

     /**
      * Want to make sure we don't throw anything or explode when people try to remove a key that's not there
      */
     @Test
     public void testRemoveMissingIsGraceful()
     {
         confirmDenied("table1", "bbb", "ccc");
         Assert.assertTrue(removeDenylist(ks_cql, "table1", "bbb:ccc"));

         // We expect this to silently not find and succeed at *trying* to remove it
         Assert.assertTrue(removeDenylist(ks_cql, "table1", "bbb:ccc"));
         refreshList();

         confirmAllowed("table1", "bbb", "ccc");
     }

     /**
      * We need to confirm that the entire cache is reloaded rather than being an additive change; we don't want keys to
      * persist after their removal and reload from CQL.
      */
     @Test
     public void testRemoveWorksOnReload()
     {
         denyAllKeys();
         refreshList();

         confirmDenied("table1", "aaa", "bbb");
         confirmDenied("table1", "eee", "fff");
         confirmDenied("table1", "iii", "jjj");

         // poke a hole in the middle and reload
         removeDenylist(ks_cql, "table1", "eee:fff");
         refreshList();

         confirmAllowed("table1", "eee", "fff");
         confirmDenied("table1", "aaa", "bbb");
         confirmDenied("table1", "iii", "jjj");
     }

     /**
      * We go through a few steps here:
      *  1) Add more keys than we're allowed
      *  2) Confirm that the overflow keys are *not* denied
      *  3) Raise the allowable limit
      *  4) Confirm that the overflow keys are now denied (and no longer really "overflow" for that matter)
      */
     @Test
     public void testShrinkAndGrow()
     {
         denyAllKeys();
         refreshList();

         // Initial control; check denial of both initial and final keys
         confirmDenied("table1", "aaa", "bbb");
         confirmDenied("table1", "iii", "jjj");

         // Lower our limit to 5 allowable denies and then check and see that things past the limit are ignored
         StorageProxy.instance.setDenylistMaxKeysPerTable(5);
         StorageProxy.instance.setDenylistMaxKeysTotal(5);
         refreshList();

         // Confirm overflowed keys are allowed; first come first served
         confirmDenied("table1", "aaa", "bbb");
         confirmAllowed("table1", "iii", "jjj");

         // Now we raise the limit back up and do nothing else and confirm it's blocked
         StorageProxy.instance.setDenylistMaxKeysPerTable(1000);
         StorageProxy.instance.setDenylistMaxKeysTotal(1000);
         refreshList();
         confirmDenied("table1", "aaa", "bbb");
         confirmDenied("table1", "iii", "jjj");

         // Unblock via overflow the table 1 sentinel we'll check in a second
         StorageProxy.instance.setDenylistMaxKeysPerTable(5);
         StorageProxy.instance.setDenylistMaxKeysTotal(5);
         refreshList();
         confirmAllowed("table1", "iii", "jjj");

         // Now, we remove the denylist entries for our first 5, drop the limit back down, and confirm those overflowed keys now block
         removeDenylist(ks_cql, "table1", "aaa:bbb");
         removeDenylist(ks_cql, "table1", "bbb:ccc");
         removeDenylist(ks_cql, "table1", "ccc:ddd");
         removeDenylist(ks_cql, "table1", "ddd:eee");
         removeDenylist(ks_cql, "table1", "eee:fff");
         refreshList();
         confirmDenied("table1", "iii", "jjj");
     }

     /*
     We need to test that, during a violation of our global allowable limit, we still enforce our limit of keys queried
     on a per-table basis.
      */
     @Test
     public void testTableLimitRespected()
     {
         StorageProxy.instance.setDenylistMaxKeysPerTable(5);
         StorageProxy.instance.setDenylistMaxKeysTotal(12);
         denyAllKeys();
         refreshList();

         // Table 1: expect first 5 denied
         confirmDenied("table1", "aaa", "bbb");
         confirmDenied("table1", "eee", "fff");
         confirmAllowed("table1", "fff", "ggg");

         // Table 2: expect first 5 denied
         for (int i = 0; i < 5; i++)
             confirmDenied("table2", Integer.toString(i), Integer.toString(i));

         // Confirm remainder are allowed because we hit our table limit at 5
         for (int i = 5; i < 50; i++)
             confirmAllowed("table2", Integer.toString(i), Integer.toString(i));

         // Table 3: expect only first 2 denied; global limit enforcement
         confirmDenied("table3", "0", "0");
         confirmDenied("table3", "1", "1");

         // And our final 48 should be allowed
         for (int i = 2; i < 50; i++)
             confirmAllowed("table3", Integer.toString(i), Integer.toString(i));
     }

     /**
      * Test that we don't allow overflowing global limit due to accumulation of allowable table queries
      */
     @Test
     public void testGlobalLimitRespected()
     {
         StorageProxy.instance.setDenylistMaxKeysPerTable(50);
         StorageProxy.instance.setDenylistMaxKeysTotal(15);
         denyAllKeys();
         refreshList();

         // Table 1: expect all 10 denied
         confirmDenied("table1", "aaa", "bbb");
         confirmDenied("table1", "jjj", "kkk");

         // Table 2: expect only 5 denied up to global limit trigger
         for (int i = 0; i < 5; i++)
             confirmDenied("table2", Integer.toString(i), Integer.toString(i));

         // Remainder of Table 2 should be allowed; testing overflow boundary logic
         for (int i = 5; i < 50; i++)
             confirmAllowed("table2", Integer.toString(i), Integer.toString(i));

         // Table 3: expect all allowed as we're past global limit by the time we got to this table load. This confirms that
         // we bypass load completely on tables once we're at our global limit.
         for (int i = 0; i < 50; i++)
             confirmAllowed("table3", Integer.toString(i), Integer.toString(i));
     }

     private void confirmDenied(String table, String keyOne, String keyTwo)
     {
         String query = String.format("SELECT * FROM " + ks_cql + "." + table + " WHERE keyone='%s' and keytwo='%s'", keyOne, keyTwo);
         assertThatThrownBy(() -> process(query, ConsistencyLevel.ONE))
                            .isInstanceOf(InvalidRequestException.class)
                            .hasMessageContaining("Unable to read denylisted partition");
     }

     private void confirmAllowed(String table, String keyOne, String keyTwo)
     {
         process(String.format("SELECT * FROM %s.%s WHERE keyone='%s' and keytwo='%s'", ks_cql, table, keyOne, keyTwo), ConsistencyLevel.ONE);
     }

     private void resetDenylist()
     {
         process("TRUNCATE system_distributed.partition_denylist", ConsistencyLevel.ONE);
         StorageProxy.instance.setDenylistMaxKeysTotal(1000);
         StorageProxy.instance.setDenylistMaxKeysPerTable(1000);
         StorageProxy.instance.loadPartitionDenylist();
     }

     private void denyAllKeys()
     {
         denylist("table1", "aaa:bbb");
         denylist("table1", "bbb:ccc");
         denylist("table1", "ccc:ddd");
         denylist("table1", "ddd:eee");
         denylist("table1", "eee:fff");
         denylist("table1", "fff:ggg");
         denylist("table1", "ggg:hhh");
         denylist("table1", "hhh:iii");
         denylist("table1", "iii:jjj");
         denylist("table1", "jjj:kkk");

         for (int i = 0; i < 50; i++)
         {
             denylist("table2", String.format("%d:%d", i, i));
             denylist("table3", String.format("%d:%d", i, i));
         }

     }
 }
	/*
	* 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.service;

	import org.junit.Assert;
	import org.junit.Before;
	import org.junit.BeforeClass;
	import org.junit.Test;

	import org.apache.cassandra.config.DatabaseDescriptor;
	import org.apache.cassandra.cql3.CQLTester;
	import org.apache.cassandra.cql3.UntypedResultSet;
	import org.apache.cassandra.cql3.statements.schema.CreateTableStatement;
	import org.apache.cassandra.db.ConsistencyLevel;
	import org.apache.cassandra.exceptions.InvalidRequestException;
	import org.apache.cassandra.schema.KeyspaceMetadata;
	import org.apache.cassandra.schema.KeyspaceParams;
	import org.apache.cassandra.schema.SchemaTestUtil;
	import org.apache.cassandra.schema.Tables;

	import static org.apache.cassandra.cql3.QueryProcessor.process;
	import static org.assertj.core.api.Assertions.assertThatThrownBy;

	public class PartitionDenylistTest
	{
	private final static String ks_cql = "partition_denylist_keyspace";

	@BeforeClass
	public static void init()
	{
	CQLTester.prepareServer();

	KeyspaceMetadata schema = KeyspaceMetadata.create(ks_cql,
	KeyspaceParams.simple(1),
	Tables.of(
	CreateTableStatement.parse("CREATE TABLE table1 ("
	+ "keyone text, "
	+ "keytwo text, "
	+ "qux text, "
	+ "quz text, "
	+ "foo text, "
	+ "PRIMARY KEY((keyone, keytwo), qux, quz) ) ", ks_cql).build(),
	CreateTableStatement.parse("CREATE TABLE table2 ("
	+ "keyone text, "
	+ "keytwo text, "
	+ "keythree text, "
	+ "value text, "
	+ "PRIMARY KEY((keyone, keytwo), keythree) ) ", ks_cql).build(),
	CreateTableStatement.parse("CREATE TABLE table3 ("
	+ "keyone text, "
	+ "keytwo text, "
	+ "keythree text, "
	+ "value text, "
	+ "PRIMARY KEY((keyone, keytwo), keythree) ) ", ks_cql).build()
	));
	SchemaTestUtil.addOrUpdateKeyspace(schema, false);
	DatabaseDescriptor.setPartitionDenylistEnabled(true);
	DatabaseDescriptor.setDenylistRangeReadsEnabled(true);
	DatabaseDescriptor.setDenylistConsistencyLevel(ConsistencyLevel.ONE);
	DatabaseDescriptor.setDenylistRefreshSeconds(1);
	StorageService.instance.initServer(0);
	}

	@Before
	public void setup()
	{
	DatabaseDescriptor.setPartitionDenylistEnabled(true);
	resetDenylist();

	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('aaa', 'bbb', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('ccc', 'ddd', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('ddd', 'eee', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('eee', 'fff', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('fff', 'ggg', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('ggg', 'hhh', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('hhh', 'iii', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('iii', 'jjj', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('jjj', 'kkk', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);


	for (int i = 0; i < 50; i++)
	process(String.format("INSERT INTO " + ks_cql + ".table2 (keyone, keytwo, keythree, value) VALUES ('%d', '%d', '%d', '%d')", i, i, i, i), ConsistencyLevel.ONE);

	for (int i = 0; i < 50; i++)
	process(String.format("INSERT INTO " + ks_cql + ".table3 (keyone, keytwo, keythree, value) VALUES ('%d', '%d', '%d', '%d')", i, i, i, i), ConsistencyLevel.ONE);

	denylist("table1", "bbb:ccc");
	refreshList();
	}


	private static void denylist(String table, final String key)
	{
	StorageProxy.instance.denylistKey(ks_cql, table, key);
	}

	private static void refreshList()
	{
	StorageProxy.instance.loadPartitionDenylist();
	}

	/**
	* @return Whether the attempt to remove the denylisted key and refresh succeeded. Doesn't necessarily indicate the key
	* was previously blocked and found.
	*/
	private static boolean removeDenylist(final String ks, final String table, final String key)
	{
	return StorageProxy.instance.removeDenylistKey(ks, table, key);
	}

	@Test
	public void testRead()
	{
	process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='aaa' and keytwo='bbb'", ConsistencyLevel.ONE);
	}

	@Test
	public void testReadDenylisted()
	{
	assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='bbb' and keytwo='ccc'", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Unable to read denylisted partition");
	}

	@Test
	public void testIsKeyDenylistedAPI()
	{
	Assert.assertTrue(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));
	resetDenylist();
	Assert.assertFalse(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));

	// Confirm an add mutates cache state
	denylist("table1", "bbb:ccc");
	Assert.assertTrue(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));

	// Confirm removal then mutates cache w/out explicit reload
	StorageProxy.instance.removeDenylistKey(ks_cql, "table1", "bbb:ccc");
	Assert.assertFalse(StorageProxy.instance.isKeyDenylisted(ks_cql, "table1", "bbb:ccc"));
	}

	@Test
	public void testReadUndenylisted()
	{
	resetDenylist();
	process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='ccc' and keytwo='ddd'", ConsistencyLevel.ONE);
	}

	@Test
	public void testWrite()
	{
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('eee', 'fff', 'ccc', 'ddd', 'v')", ConsistencyLevel.ONE);
	process("DELETE FROM " + ks_cql + ".table1 WHERE keyone='eee' and keytwo='fff'", ConsistencyLevel.ONE);
	}

	@Test
	public void testWriteDenylisted()
	{
	assertThatThrownBy(() -> process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'eee', 'fff', 'w')", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Unable to write to denylisted partition");
	}

	@Test
	public void testCASWriteDenylisted()
	{
	assertThatThrownBy(() -> process("UPDATE " + ks_cql + ".table1 SET foo='w' WHERE keyone='bbb' AND keytwo='ccc' AND qux='eee' AND quz='fff' IF foo='v'", ConsistencyLevel.LOCAL_SERIAL))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Unable to CAS write to denylisted partition");
	}

	@Test
	public void testWriteUndenylisted()
	{
	resetDenylist();
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'eee', 'fff', 'w')", ConsistencyLevel.ONE);
	}

	@Test
	public void testRangeSlice()
	{
	UntypedResultSet rows;
	rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) < token('bbb', 'ccc')", ConsistencyLevel.ONE);
	Assert.assertEquals(1, rows.size());

	// 10 entries total in our table
	rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('bbb', 'ccc')", ConsistencyLevel.ONE);
	Assert.assertEquals(8, rows.size());

	rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) >= token('aaa', 'bbb') and token(keyone, keytwo) < token('bbb', 'ccc')", ConsistencyLevel.ONE);
	Assert.assertEquals(1, rows.size());

	rows = process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('bbb', 'ccc') and token(keyone, keytwo) <= token('ddd', 'eee')", ConsistencyLevel.ONE);
	Assert.assertEquals(2, rows.size());
	}

	@Test
	public void testRangeDenylisted()
	{
	assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
	}

	@Test
	public void testRangeDenylisted2()
	{
	assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) >= token('aaa', 'bbb') and token (keyone, keytwo) <= token('bbb', 'ccc')", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
	}

	@Test
	public void testRangeDenylisted3()
	{
	assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) >= token('bbb', 'ccc') and token (keyone, keytwo) <= token('ccc', 'ddd')", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
	}

	@Test
	public void testRangeDenylisted4()
	{
	assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('aaa', 'bbb') and token (keyone, keytwo) < token('ccc', 'ddd')", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
	}

	@Test
	public void testRangeDenylisted5()
	{
	assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) > token('aaa', 'bbb')", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
	}

	@Test
	public void testRangeDenylisted6()
	{
	assertThatThrownBy(() -> process("SELECT * FROM " + ks_cql + ".table1 WHERE token(keyone, keytwo) < token('ddd', 'eee')", ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Attempted to read a range containing 1 denylisted keys");
	}

	@Test
	public void testInsertUnknownPKIsGraceful()
	{
	Assert.assertTrue(StorageProxy.instance.denylistKey(ks_cql, "table1", "hohoho"));
	}

	@Test
	public void testInsertInvalidTableIsGraceful()
	{
	Assert.assertFalse(StorageProxy.instance.denylistKey(ks_cql, "asldkfjadlskjf", "alksdjfads"));
	}

	@Test
	public void testInsertInvalidKSIsGraceful()
	{
	Assert.assertFalse(StorageProxy.instance.denylistKey("asdklfjas", "asldkfjadlskjf", "alksdjfads"));
	}

	@Test
	public void testDisabledDenylistThrowsNoExceptions()
	{
	process(String.format("TRUNCATE TABLE %s.table2", ks_cql), ConsistencyLevel.ONE);
	process(String.format("TRUNCATE TABLE %s.table3", ks_cql), ConsistencyLevel.ONE);
	denyAllKeys();
	DatabaseDescriptor.setPartitionDenylistEnabled(false);
	process("INSERT INTO " + ks_cql + ".table1 (keyone, keytwo, qux, quz, foo) VALUES ('bbb', 'ccc', 'eee', 'fff', 'w')", ConsistencyLevel.ONE);
	process("SELECT * FROM " + ks_cql + ".table1 WHERE keyone='bbb' and keytwo='ccc'", ConsistencyLevel.ONE);
	process("SELECT * FROM " + ks_cql + ".table1", ConsistencyLevel.ONE);

	for (int i = 0; i < 50; i++)
	{
	process(String.format("INSERT INTO %s.table2 (keyone, keytwo, keythree, value) VALUES ('%s', '%s', '%s', '%s')", ks_cql, i, i, i, i), ConsistencyLevel.ONE);
	process(String.format("SELECT * FROM %s.table2 WHERE keyone='%s' and keytwo='%s'", ks_cql, i, i), ConsistencyLevel.ONE);
	}

	for (int i = 0; i < 50; i++)
	{
	process(String.format("INSERT INTO %s.table3 (keyone, keytwo, keythree, value) VALUES ('%s', '%s', '%s', '%s')", ks_cql, i, i, i, i), ConsistencyLevel.ONE);
	process(String.format("SELECT * FROM %s.table3 WHERE keyone='%s' and keytwo='%s'", ks_cql, i, i), ConsistencyLevel.ONE);
	}
	}

	/**
	* Want to make sure we don't throw anything or explode when people try to remove a key that's not there
	*/
	@Test
	public void testRemoveMissingIsGraceful()
	{
	confirmDenied("table1", "bbb", "ccc");
	Assert.assertTrue(removeDenylist(ks_cql, "table1", "bbb:ccc"));

	// We expect this to silently not find and succeed at trying to remove it
	Assert.assertTrue(removeDenylist(ks_cql, "table1", "bbb:ccc"));
	refreshList();

	confirmAllowed("table1", "bbb", "ccc");
	}

	/**
	* We need to confirm that the entire cache is reloaded rather than being an additive change; we don't want keys to
	* persist after their removal and reload from CQL.
	*/
	@Test
	public void testRemoveWorksOnReload()
	{
	denyAllKeys();
	refreshList();

	confirmDenied("table1", "aaa", "bbb");
	confirmDenied("table1", "eee", "fff");
	confirmDenied("table1", "iii", "jjj");

	// poke a hole in the middle and reload
	removeDenylist(ks_cql, "table1", "eee:fff");
	refreshList();

	confirmAllowed("table1", "eee", "fff");
	confirmDenied("table1", "aaa", "bbb");
	confirmDenied("table1", "iii", "jjj");
	}

	/**
	* We go through a few steps here:
	* 1) Add more keys than we're allowed
	* 2) Confirm that the overflow keys are not denied
	* 3) Raise the allowable limit
	* 4) Confirm that the overflow keys are now denied (and no longer really "overflow" for that matter)
	*/
	@Test
	public void testShrinkAndGrow()
	{
	denyAllKeys();
	refreshList();

	// Initial control; check denial of both initial and final keys
	confirmDenied("table1", "aaa", "bbb");
	confirmDenied("table1", "iii", "jjj");

	// Lower our limit to 5 allowable denies and then check and see that things past the limit are ignored
	StorageProxy.instance.setDenylistMaxKeysPerTable(5);
	StorageProxy.instance.setDenylistMaxKeysTotal(5);
	refreshList();

	// Confirm overflowed keys are allowed; first come first served
	confirmDenied("table1", "aaa", "bbb");
	confirmAllowed("table1", "iii", "jjj");

	// Now we raise the limit back up and do nothing else and confirm it's blocked
	StorageProxy.instance.setDenylistMaxKeysPerTable(1000);
	StorageProxy.instance.setDenylistMaxKeysTotal(1000);
	refreshList();
	confirmDenied("table1", "aaa", "bbb");
	confirmDenied("table1", "iii", "jjj");

	// Unblock via overflow the table 1 sentinel we'll check in a second
	StorageProxy.instance.setDenylistMaxKeysPerTable(5);
	StorageProxy.instance.setDenylistMaxKeysTotal(5);
	refreshList();
	confirmAllowed("table1", "iii", "jjj");

	// Now, we remove the denylist entries for our first 5, drop the limit back down, and confirm those overflowed keys now block
	removeDenylist(ks_cql, "table1", "aaa:bbb");
	removeDenylist(ks_cql, "table1", "bbb:ccc");
	removeDenylist(ks_cql, "table1", "ccc:ddd");
	removeDenylist(ks_cql, "table1", "ddd:eee");
	removeDenylist(ks_cql, "table1", "eee:fff");
	refreshList();
	confirmDenied("table1", "iii", "jjj");
	}

	/*
	We need to test that, during a violation of our global allowable limit, we still enforce our limit of keys queried
	on a per-table basis.
	*/
	@Test
	public void testTableLimitRespected()
	{
	StorageProxy.instance.setDenylistMaxKeysPerTable(5);
	StorageProxy.instance.setDenylistMaxKeysTotal(12);
	denyAllKeys();
	refreshList();

	// Table 1: expect first 5 denied
	confirmDenied("table1", "aaa", "bbb");
	confirmDenied("table1", "eee", "fff");
	confirmAllowed("table1", "fff", "ggg");

	// Table 2: expect first 5 denied
	for (int i = 0; i < 5; i++)
	confirmDenied("table2", Integer.toString(i), Integer.toString(i));

	// Confirm remainder are allowed because we hit our table limit at 5
	for (int i = 5; i < 50; i++)
	confirmAllowed("table2", Integer.toString(i), Integer.toString(i));

	// Table 3: expect only first 2 denied; global limit enforcement
	confirmDenied("table3", "0", "0");
	confirmDenied("table3", "1", "1");

	// And our final 48 should be allowed
	for (int i = 2; i < 50; i++)
	confirmAllowed("table3", Integer.toString(i), Integer.toString(i));
	}

	/**
	* Test that we don't allow overflowing global limit due to accumulation of allowable table queries
	*/
	@Test
	public void testGlobalLimitRespected()
	{
	StorageProxy.instance.setDenylistMaxKeysPerTable(50);
	StorageProxy.instance.setDenylistMaxKeysTotal(15);
	denyAllKeys();
	refreshList();

	// Table 1: expect all 10 denied
	confirmDenied("table1", "aaa", "bbb");
	confirmDenied("table1", "jjj", "kkk");

	// Table 2: expect only 5 denied up to global limit trigger
	for (int i = 0; i < 5; i++)
	confirmDenied("table2", Integer.toString(i), Integer.toString(i));

	// Remainder of Table 2 should be allowed; testing overflow boundary logic
	for (int i = 5; i < 50; i++)
	confirmAllowed("table2", Integer.toString(i), Integer.toString(i));

	// Table 3: expect all allowed as we're past global limit by the time we got to this table load. This confirms that
	// we bypass load completely on tables once we're at our global limit.
	for (int i = 0; i < 50; i++)
	confirmAllowed("table3", Integer.toString(i), Integer.toString(i));
	}

	private void confirmDenied(String table, String keyOne, String keyTwo)
	{
	String query = String.format("SELECT * FROM " + ks_cql + "." + table + " WHERE keyone='%s' and keytwo='%s'", keyOne, keyTwo);
	assertThatThrownBy(() -> process(query, ConsistencyLevel.ONE))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining("Unable to read denylisted partition");
	}

	private void confirmAllowed(String table, String keyOne, String keyTwo)
	{
	process(String.format("SELECT * FROM %s.%s WHERE keyone='%s' and keytwo='%s'", ks_cql, table, keyOne, keyTwo), ConsistencyLevel.ONE);
	}

	private void resetDenylist()
	{
	process("TRUNCATE system_distributed.partition_denylist", ConsistencyLevel.ONE);
	StorageProxy.instance.setDenylistMaxKeysTotal(1000);
	StorageProxy.instance.setDenylistMaxKeysPerTable(1000);
	StorageProxy.instance.loadPartitionDenylist();
	}

	private void denyAllKeys()
	{
	denylist("table1", "aaa:bbb");
	denylist("table1", "bbb:ccc");
	denylist("table1", "ccc:ddd");
	denylist("table1", "ddd:eee");
	denylist("table1", "eee:fff");
	denylist("table1", "fff:ggg");
	denylist("table1", "ggg:hhh");
	denylist("table1", "hhh:iii");
	denylist("table1", "iii:jjj");
	denylist("table1", "jjj:kkk");

	for (int i = 0; i < 50; i++)
	{
	denylist("table2", String.format("%d:%d", i, i));
	denylist("table3", String.format("%d:%d", i, i));
	}

	}
	}