Skip to content

Commit

Permalink
rollback cql-session check (#1138)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Yuqi-Du
Yuqi-Du authored Jun 3, 2024
1 parent 4d84b17 commit e42550a
Show file tree
Hide file tree
Showing 2 changed files with 108 additions and 104 deletions.
70 changes: 39 additions & 31 deletions src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/CQLSessionCache.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,42 +174,50 @@ private CqlSession getNewSession(SessionCacheKey cacheKey) {
.withConfigLoader(loader)
.addSchemaChangeListener(new SchemaChangeListener(schemaCache, cacheKey.tenantId))
.build();
if (!isAstraSessionValid(cqlSession, cacheKey.tenantId())) {
throw new UnauthorizedException("Provided username token and/or password are incorrect");
}
// if (!isAstraSessionValid(cqlSession, cacheKey.tenantId())) {
// throw new UnauthorizedException("Provided username token and/or password are
// incorrect");
// }
return cqlSession;
}
throw new RuntimeException("Unsupported database type: " + databaseConfig.type());
}

/**
* This method checks if the session is valid for the tenant. If a token is generated for tenant A
* and if it is used to access tenant B's data, the cqlsession object still gets created without
* any error but it has no metadata or keyspaces information. So, this situation leads to return
* misleading no keyspace found error, instead of authorization error.
*
* <p>This method checks if the session is valid, first by checking if there are any keyspaces and
* returns true if there are any keyspaces. If there are no keyspaces, then it tries to execute a
* query on system_virtual_schema.tables and returns true if the query is successful. Failure to
* execute the query with an UnauthorizedException means the session is invalid i.e. not meant for
* the tenant in the request.
*
* @param cqlSession CqlSession
* @param tenantId tenant id
* @return true if the session is valid, false otherwise
*/
private boolean isAstraSessionValid(CqlSession cqlSession, String tenantId) {
if (!cqlSession.getMetadata().getKeyspaces().isEmpty()) {
return true;
}
try {
cqlSession.execute("SELECT * FROM system_virtual_schema.tables");
return true;
} catch (com.datastax.oss.driver.api.core.servererrors.UnauthorizedException e) {
LOGGER.error("Unauthorized to access tenant %s's data".formatted(tenantId), e);
return false;
}
}
// /**
// * This method checks if the session is valid for the tenant. If a token is generated for
// tenant A
// * and if it is used to access tenant B's data, the cqlsession object still gets created
// without
// * any error but it has no metadata or keyspaces information. So, this situation leads to
// return
// * misleading no keyspace found error, instead of authorization error.
// *
// * <p>This method checks if the session is valid, first by checking if there are any keyspaces
// and
// * returns true if there are any keyspaces. If there are no keyspaces, then it tries to
// execute a
// * query on system_virtual_schema.tables and returns true if the query is successful. Failure
// to
// * execute the query with an UnauthorizedException means the session is invalid i.e. not meant
// for
// * the tenant in the request.
// *
// * @param cqlSession CqlSession
// * @param tenantId tenant id
// * @return true if the session is valid, false otherwise
// */
// private boolean isAstraSessionValid(CqlSession cqlSession, String tenantId) {
// if (!cqlSession.getMetadata().getKeyspaces().isEmpty()) {
// return true;
// }
// try {
// cqlSession.execute("SELECT * FROM system_virtual_schema.tables");
// return true;
// } catch (com.datastax.oss.driver.api.core.servererrors.UnauthorizedException e) {
// LOGGER.error("Unauthorized to access tenant %s's data".formatted(tenantId), e);
// return false;
// }
// }

/**
* Get CQLSession from cache.
Expand Down
142 changes: 69 additions & 73 deletions src/test/java/io/stargate/sgv2/jsonapi/service/cqldriver/CqlSessionCacheTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,7 @@
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.metadata.Metadata;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
import com.datastax.oss.driver.internal.core.context.DefaultDriverContext;
import io.micrometer.core.instrument.FunctionCounter;
import io.micrometer.core.instrument.Gauge;
Expand All @@ -23,11 +19,8 @@
import io.stargate.sgv2.jsonapi.config.OperationsConfig;
import jakarta.inject.Inject;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
Expand Down Expand Up @@ -240,70 +233,73 @@ public void testOSSCxCQLSessionCacheSizeEviction()
assertThat(cacheLoadMetric).isNotNull();
assertThat(cacheLoadMetric.count()).isEqualTo(sessionsToCreate);
}

@Test
public void testInvalidSession()
throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
CQLSessionCache cqlSessionCacheForTest = new CQLSessionCache(operationsConfig, meterRegistry);
CqlSession cqlSession = mock(CqlSession.class);
Metadata metadata = mock(Metadata.class);
when(cqlSession.getMetadata()).thenReturn(metadata);
when(metadata.getKeyspaces()).thenReturn(Map.of());
Node node = mock(Node.class);
com.datastax.oss.driver.api.core.servererrors.UnauthorizedException unauthorizedException =
new com.datastax.oss.driver.api.core.servererrors.UnauthorizedException(
node, "Unauthorized");
when(cqlSession.execute("SELECT * FROM system_virtual_schema.tables"))
.thenThrow(unauthorizedException);
Method isValidMethod =
cqlSessionCacheForTest
.getClass()
.getDeclaredMethod("isAstraSessionValid", CqlSession.class, String.class);
isValidMethod.setAccessible(true);
boolean isValid =
(boolean) isValidMethod.invoke(cqlSessionCacheForTest, cqlSession, TENANT_ID_FOR_TEST);
assertThat(isValid).isFalse();
}

@Test
public void testAValidSessionWithKeyspaces()
throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
CQLSessionCache cqlSessionCacheForTest = new CQLSessionCache(operationsConfig, meterRegistry);
CqlSession cqlSession = mock(CqlSession.class);
Metadata metadata = mock(Metadata.class);
when(cqlSession.getMetadata()).thenReturn(metadata);
when(metadata.getKeyspaces())
.thenReturn(Map.of(CqlIdentifier.fromCql("ks1"), mock(KeyspaceMetadata.class)));
Method isValidMethod =
cqlSessionCacheForTest
.getClass()
.getDeclaredMethod("isAstraSessionValid", CqlSession.class, String.class);
isValidMethod.setAccessible(true);
boolean isValid =
(boolean) isValidMethod.invoke(cqlSessionCacheForTest, cqlSession, TENANT_ID_FOR_TEST);
assertThat(isValid).isTrue();
}

@Test
public void testAValidSessionNoKeyspaces()
throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
CQLSessionCache cqlSessionCacheForTest = new CQLSessionCache(operationsConfig, meterRegistry);
CqlSession cqlSession = mock(CqlSession.class);
Metadata metadata = mock(Metadata.class);
when(cqlSession.getMetadata()).thenReturn(metadata);
when(metadata.getKeyspaces()).thenReturn(Map.of());
Node node = mock(Node.class);
com.datastax.oss.driver.api.core.servererrors.UnauthorizedException unauthorizedException =
new com.datastax.oss.driver.api.core.servererrors.UnauthorizedException(
node, "Unauthorized");
when(cqlSession.execute("SELECT * FROM system_virtual_schema.tables")).thenReturn(null);
Method isValidMethod =
cqlSessionCacheForTest
.getClass()
.getDeclaredMethod("isAstraSessionValid", CqlSession.class, String.class);
isValidMethod.setAccessible(true);
boolean isValid =
(boolean) isValidMethod.invoke(cqlSessionCacheForTest, cqlSession, TENANT_ID_FOR_TEST);
assertThat(isValid).isTrue();
}
//
// @Test
// public void testInvalidSession()
// throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
// CQLSessionCache cqlSessionCacheForTest = new CQLSessionCache(operationsConfig,
// meterRegistry);
// CqlSession cqlSession = mock(CqlSession.class);
// Metadata metadata = mock(Metadata.class);
// when(cqlSession.getMetadata()).thenReturn(metadata);
// when(metadata.getKeyspaces()).thenReturn(Map.of());
// Node node = mock(Node.class);
// com.datastax.oss.driver.api.core.servererrors.UnauthorizedException unauthorizedException =
// new com.datastax.oss.driver.api.core.servererrors.UnauthorizedException(
// node, "Unauthorized");
// when(cqlSession.execute("SELECT * FROM system_virtual_schema.tables"))
// .thenThrow(unauthorizedException);
// Method isValidMethod =
// cqlSessionCacheForTest
// .getClass()
// .getDeclaredMethod("isAstraSessionValid", CqlSession.class, String.class);
// isValidMethod.setAccessible(true);
// boolean isValid =
// (boolean) isValidMethod.invoke(cqlSessionCacheForTest, cqlSession, TENANT_ID_FOR_TEST);
// assertThat(isValid).isFalse();
// }
//
// @Test
// public void testAValidSessionWithKeyspaces()
// throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
// CQLSessionCache cqlSessionCacheForTest = new CQLSessionCache(operationsConfig,
// meterRegistry);
// CqlSession cqlSession = mock(CqlSession.class);
// Metadata metadata = mock(Metadata.class);
// when(cqlSession.getMetadata()).thenReturn(metadata);
// when(metadata.getKeyspaces())
// .thenReturn(Map.of(CqlIdentifier.fromCql("ks1"), mock(KeyspaceMetadata.class)));
// Method isValidMethod =
// cqlSessionCacheForTest
// .getClass()
// .getDeclaredMethod("isAstraSessionValid", CqlSession.class, String.class);
// isValidMethod.setAccessible(true);
// boolean isValid =
// (boolean) isValidMethod.invoke(cqlSessionCacheForTest, cqlSession, TENANT_ID_FOR_TEST);
// assertThat(isValid).isTrue();
// }
//
// @Test
// public void testAValidSessionNoKeyspaces()
// throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
// CQLSessionCache cqlSessionCacheForTest = new CQLSessionCache(operationsConfig,
// meterRegistry);
// CqlSession cqlSession = mock(CqlSession.class);
// Metadata metadata = mock(Metadata.class);
// when(cqlSession.getMetadata()).thenReturn(metadata);
// when(metadata.getKeyspaces()).thenReturn(Map.of());
// Node node = mock(Node.class);
// com.datastax.oss.driver.api.core.servererrors.UnauthorizedException unauthorizedException =
// new com.datastax.oss.driver.api.core.servererrors.UnauthorizedException(
// node, "Unauthorized");
// when(cqlSession.execute("SELECT * FROM system_virtual_schema.tables")).thenReturn(null);
// Method isValidMethod =
// cqlSessionCacheForTest
// .getClass()
// .getDeclaredMethod("isAstraSessionValid", CqlSession.class, String.class);
// isValidMethod.setAccessible(true);
// boolean isValid =
// (boolean) isValidMethod.invoke(cqlSessionCacheForTest, cqlSession, TENANT_ID_FOR_TEST);
// assertThat(isValid).isTrue();
// }
}

0 comments on commit e42550a

Please sign in to comment.