content: /*
   * 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.transport;
  
  import java.io.IOException;
  import java.util.Map;
  import java.util.concurrent.TimeoutException;
  import java.util.function.Consumer;
  
  import com.google.common.collect.ImmutableMap;
  import org.junit.Before;
  import org.junit.BeforeClass;
  import org.junit.Test;
  
  import org.apache.cassandra.auth.AuthCacheService;
  import org.apache.cassandra.auth.IAuthenticator;
  import org.apache.cassandra.auth.MutualTlsAuthenticator;
  import org.apache.cassandra.auth.SpiffeCertificateValidator;
  import org.apache.cassandra.config.EncryptionOptions;
  import org.apache.cassandra.config.ParameterizedClass;
  import org.apache.cassandra.cql3.CQLTester;
  import org.apache.cassandra.cql3.QueryProcessor;
  import org.apache.cassandra.exceptions.AuthenticationException;
  import org.apache.cassandra.transport.messages.ErrorMessage;
  import org.apache.cassandra.transport.messages.ReadyMessage;
  import org.apache.cassandra.transport.messages.StartupMessage;
  
  import static org.apache.cassandra.auth.AuthTestUtils.addIdentityToRole;
  import static org.apache.cassandra.auth.AuthTestUtils.truncateIdentityRolesTable;
  import static org.apache.cassandra.transport.messages.StartupMessage.CQL_VERSION;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.fail;
  import static org.psjava.util.AssertStatus.assertTrue;
  
  /**
   * Verifies the behavior of Cassandra given an authenticator ({@link MutualTlsAuthenticator}) returning
   * {@link IAuthenticator#supportsEarlyAuthentication()} as <code>true</code>.
   */
  public class EarlyAuthenticationTest extends CQLTester
  {
  
      // configures server with client encryption enabled.
      static final Consumer<Server.Builder> serverConfigurator = server -> server.withTlsEncryptionPolicy(EncryptionOptions.TlsEncryptionPolicy.ENCRYPTED);
  
      static final Map<String, String> authenticatorParams = ImmutableMap.of("validator_class_name", SpiffeCertificateValidator.class.getSimpleName());
  
      @BeforeClass
      public static void setup()
      {
          setupWithAuthenticator(new MutualTlsAuthenticator(authenticatorParams));
      }
  
      static void setupWithAuthenticator(IAuthenticator authenticator)
      {
          requireNativeProtocolClientEncryption();
          requireNetwork(serverConfigurator, cluster -> {
          });
          requireAuthentication(authenticator);
      }
  
      @Before
      public void initNetwork() throws IOException, TimeoutException
      {
          truncateIdentityRolesTable();
          AuthCacheService.instance.invalidateCaches();
          AuthCacheService.instance.warmCaches();
          reinitializeNetwork(serverConfigurator, cluster -> {
          });
      }
  
      private EncryptionOptions.ClientEncryptionOptions clientEncryptionOptions(boolean presentClientCertificate)
      {
          EncryptionOptions.ClientEncryptionOptions.Builder builder = new EncryptionOptions.ClientEncryptionOptions.Builder();
          builder.withEnabled(true)
                 .withRequireClientAuth(EncryptionOptions.ClientEncryptionOptions.ClientAuth.OPTIONAL)
                 .withTrustStore(TlsTestUtils.CLIENT_TRUSTSTORE_PATH)
                 .withTrustStorePassword(TlsTestUtils.CLIENT_TRUSTSTORE_PASSWORD)
                 .withSslContextFactory(new ParameterizedClass(SimpleClientSslContextFactory.class.getName()));
  
          if (presentClientCertificate)
          {
              builder.withKeyStore(TlsTestUtils.CLIENT_SPIFFE_KEYSTORE_PATH)
                                      .withStoreType("JKS")
                                      .withKeyStorePassword(TlsTestUtils.CLIENT_SPIFFE_KEYSTORE_PASSWORD);
          }
  
          return new EncryptionOptions.ClientEncryptionOptions(builder.build());
      }
  
      @Test
      public void testEarlyAuthSuccess()
      {
          // given server is configured with a Mutual TLS Authenticator and the identity in the client's keystore is bound to cassandra.
          addIdentityToRole(TlsTestUtils.CLIENT_SPIFFE_IDENTITY, "cassandra");
  
          // when connecting, we expect to get a 'READY' message after sending a 'STARTUP' as MutualTlsAuthenticator
          // supports early authentication and the client presented a cert with an identity that was bound to a role.
          testStartupResponse(true, startupResponse -> {
              if (!(startupResponse instanceof ReadyMessage))
              {
                  fail("Expected an READY in response to a STARTUP, got: " + startupResponse);
              }
          });
      }
  
      @Test
      public void testEarlyAuthFailureLoginDisallowed()
      {
          // given server is configured with a Mutual TLS Authenticator and the identity in the client's keystore is bound
          // to a role that is not permitted to log in.
          // when connecting, we expect an 'ERROR' message.
          addIdentityToRole(TlsTestUtils.CLIENT_SPIFFE_IDENTITY, "readonly_user");
          testStartupResponse(true, expectAuthenticationError("readonly_user is not permitted to log in"));
      }
  
      @Test
      public void testEarlyAuthFailureMissingIdentity()
      {
          // given server is configured with a Mutual TLS Authenticator, but no identities are bound to roles.
          // when connecting, we expect an 'ERROR' message.
          testStartupResponse(true, expectAuthenticationError(String.format("Certificate identity '%s' not authorized", TlsTestUtils.CLIENT_SPIFFE_IDENTITY)));
      }
  
      @Test
      public void testNoClientCertificatePresented()
      {
          // given server is configured with a Mutual TLS Authenticator, but no client certificate is presented.
          // when connecting, we expect an 'ERROR' message.
          testStartupResponse(false, expectAuthenticationError("No certificate present on connection"));
      }
  
      public void testStartupResponse(boolean presentClientCertificate, Consumer<Message.Response> testFn)
      {
          SimpleClient.Builder builder = SimpleClient.builder(nativeAddr.getHostAddress(), nativePort)
                                                     .encryption(clientEncryptionOptions(presentClientCertificate));
          try (SimpleClient client = builder.build())
          {
              client.establishConnection();
  
              // Send a StartupMessage and pass the response to the handling function
              StartupMessage startup = new StartupMessage(ImmutableMap.of(CQL_VERSION, QueryProcessor.CQL_VERSION.toString()));
              Message.Response startupResponse = client.execute(startup, false);
              testFn.accept(startupResponse);
          }
          catch (IOException e)
          {
              fail("Error establishing connection");
          }
      }
  
      public Consumer<Message.Response> expectAuthenticationError(final String expectedMessage)
      {
          return startupResponse -> {
              if (startupResponse instanceof ErrorMessage)
              {
                  ErrorMessage errorMessage = (ErrorMessage) startupResponse;
                  assertTrue(errorMessage.error instanceof AuthenticationException, "Expected an AuthenticationException, got: " + errorMessage.error);
                  AuthenticationException authException = (AuthenticationException) errorMessage.error;
                  assertEquals(expectedMessage, authException.getMessage());
              }
              else
              {
                  fail("Expected an ErrorMessage but got: " + startupResponse);
              }
          };
      }
  }
  
  