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.auth;
  
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Set;
  import java.util.function.Supplier;
  import java.util.stream.Collectors;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.Lists;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.cql3.QueryOptions;
  import org.apache.cassandra.cql3.QueryProcessor;
  import org.apache.cassandra.cql3.UntypedResultSet;
  import org.apache.cassandra.cql3.statements.SelectStatement;
  import org.apache.cassandra.db.ConsistencyLevel;
  import org.apache.cassandra.db.marshal.UTF8Type;
  import org.apache.cassandra.exceptions.RequestExecutionException;
  import org.apache.cassandra.schema.SchemaConstants;
  import org.apache.cassandra.service.ClientState;
  import org.apache.cassandra.service.reads.range.RangeCommands;
  import org.apache.cassandra.transport.Dispatcher;
  import org.apache.cassandra.transport.messages.ResultMessage;
  import org.apache.cassandra.utils.ByteBufferUtil;
  import org.apache.cassandra.utils.MBeanWrapper;
  
  import static org.apache.cassandra.service.QueryState.forInternalCalls;
  
  /**
   * Provides functionality to list/update/drop CIDR permissions of a role
   * Backend to build Role to CIDR permissions cache
   */
  public class CIDRPermissionsManager implements CIDRPermissionsManagerMBean, AuthCache.BulkLoader<RoleResource, CIDRPermissions>
  {
      public static final String MBEAN_NAME = "org.apache.cassandra.db:type=CIDRPermissionsManager";
  
      private static final Logger logger = LoggerFactory.getLogger(CIDRPermissionsManager.class);
      private SelectStatement getCidrPermissionsOfUserStatement = null;
  
      public void setup()
      {
          if (!MBeanWrapper.instance.isRegistered(MBEAN_NAME))
              MBeanWrapper.instance.registerMBean(this, MBEAN_NAME);
  
          String getCidrPermissionsOfUserQuery = String.format("SELECT cidr_groups FROM %s.%s WHERE role = ?",
                                                               SchemaConstants.AUTH_KEYSPACE_NAME,
                                                               AuthKeyspace.CIDR_PERMISSIONS);
          getCidrPermissionsOfUserStatement = (SelectStatement) QueryProcessor.getStatement(getCidrPermissionsOfUserQuery,
                                                                                            ClientState.forInternalCalls());
      }
  
      @VisibleForTesting
      ResultMessage.Rows select(SelectStatement statement, QueryOptions options)
      {
          return statement.execute(forInternalCalls(), options, Dispatcher.RequestTime.forImmediateExecution());
      }
  
      @VisibleForTesting
      UntypedResultSet process(String query, ConsistencyLevel cl) throws RequestExecutionException
      {
          return QueryProcessor.process(query, cl);
      }
  
      private Set<String> getAuthorizedCIDRGroups(String name)
      {
          QueryOptions options = QueryOptions.forInternalCalls(CassandraAuthorizer.authReadConsistencyLevel(),
                                                               Lists.newArrayList(ByteBufferUtil.bytes(name)));
  
          ResultMessage.Rows rows = select(getCidrPermissionsOfUserStatement, options);
          UntypedResultSet result = UntypedResultSet.create(rows.result);
          if (!result.isEmpty() && result.one().has("cidr_groups"))
          {
              return result.one().getFrozenSet("cidr_groups", UTF8Type.instance);
          }
  
          return Collections.emptySet();
      }
  
      private static String getCidrPermissionsSetString(CIDRPermissions permissions)
      {
          String inner = "";
  
          if (permissions.restrictsAccess())
              inner = permissions.allowedCIDRGroups().stream().map(s -> '\'' + s + '\'')
                                 .collect(Collectors.joining(", "));
  
          return '{' + inner + '}';
      }
  
      /**
       * Get CIDR permissions of a role
       * @param role role for which to get CIDR permissions
       * @return returns CIDR permissions of a role
       */
      public CIDRPermissions getCidrPermissionsForRole(RoleResource role)
      {
          if (!Roles.canLogin(role))
          {
              return CIDRPermissions.none();
          }
          if (Roles.hasSuperuserStatus(role) && !DatabaseDescriptor.getCidrChecksForSuperusers())
          {
              return CIDRPermissions.all();
          }
  
          Set<String> cidrGroups = getAuthorizedCIDRGroups(role.getRoleName());
          // User don't have CIDR permissions explicitly enabled, i.e, allow from all
          if (cidrGroups == null || cidrGroups.isEmpty())
          {
              // No explicit CIDR groups set for the role, allow from all
              return CIDRPermissions.all();
          }
  
          return CIDRPermissions.subset(cidrGroups);
      }
  
      /**
       * Set CIDR permissions for a given role
       * @param role role for which to set CIDR permissions
       * @param cidrPermissions CIR permissions to set for the role
       */
      public void setCidrGroupsForRole(RoleResource role, CIDRPermissions cidrPermissions)
      {
          String query = String.format("UPDATE %s.%s SET cidr_groups = %s WHERE role = '%s'",
                                       SchemaConstants.AUTH_KEYSPACE_NAME,
                                       AuthKeyspace.CIDR_PERMISSIONS,
                                       getCidrPermissionsSetString(cidrPermissions),
                                       role.getRoleName());
  
          process(query, CassandraAuthorizer.authWriteConsistencyLevel());
      }
  
      /**
       * Drop CIDR permissions of a role, i.e, delete corresponding row from the table
       * @param role for which to drop cidr permissions
       */
      public void drop(RoleResource role)
      {
          String query = String.format("DELETE FROM %s.%s WHERE role = '%s'",
                                       SchemaConstants.AUTH_KEYSPACE_NAME,
                                       AuthKeyspace.CIDR_PERMISSIONS,
                                       role.getRoleName());
  
          process(query, CassandraAuthorizer.authWriteConsistencyLevel());
      }
  
      /**
       * Function to bulk load role to cidr permissions cache
       * This gets called when cache warmup is enabled in the config
       * @return returns nothing
       */
      @Override
      public Supplier<Map<RoleResource, CIDRPermissions>> bulkLoader()
      {
          return () -> {
              if (!RangeCommands.sufficientLiveNodesForSelectStar(AuthKeyspace.metadata().tables.getNullable(AuthKeyspace.CIDR_PERMISSIONS),
                                                                  AuthProperties.instance.getReadConsistencyLevel()))
              {
                  // Prevent running the query we know will fail so as not to increment unavailable stats for a performance
                  // optimization
                  throw new RuntimeException("insufficient live nodes for " +
                                             AuthProperties.instance.getReadConsistencyLevel() +
                                             "pre-warm query again system_auth." + AuthKeyspace.CIDR_PERMISSIONS);
              }
  
              logger.info("Pre-warming CIDR permissions cache from cidr_permissions table");
              Map<RoleResource, CIDRPermissions> entries = new HashMap<>();
              UntypedResultSet rows = process(String.format("SELECT role, cidr_groups FROM %s.%s",
                                                            SchemaConstants.AUTH_KEYSPACE_NAME,
                                                            AuthKeyspace.CIDR_PERMISSIONS),
                                              CassandraAuthorizer.authReadConsistencyLevel());
  
              for (UntypedResultSet.Row row : rows)
              {
                  RoleResource role = RoleResource.role(row.getString("role"));
                  CIDRPermissions.Builder builder = new CIDRPermissions.Builder();
                  Set<String> cidrGroups = row.getFrozenSet("cidr_groups",
                                                            UTF8Type.instance);
                  for (String cidrGroup : cidrGroups)
                      builder.add(cidrGroup);
                  entries.put(role, builder.build());
              }
  
              return entries;
          };
      }
  
      public boolean invalidateCidrPermissionsCache(String roleName)
      {
          return DatabaseDescriptor.getCIDRAuthorizer().invalidateCidrPermissionsCache(roleName);
      }
  }
  