org.projectnessie.services.authz.AuthzPaginationIterator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of nessie-services Show documentation
The newest version!
/*
 * Copyright (C) 2023 Dremio
 *
 * Licensed 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.projectnessie.services.authz;

import static com.google.common.collect.Sets.newHashSetWithExpectedSize;
import static java.util.Collections.emptyIterator;
import static java.util.Collections.emptyMap;
import static org.projectnessie.services.authz.BatchAccessChecker.throwForFailedChecks;

import com.google.common.collect.AbstractIterator;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;
import org.projectnessie.versioned.paging.PaginationIterator;

public abstract class AuthzPaginationIterator extends AbstractIterator
    implements PaginationIterator {

  private final PaginationIterator source;

  private final Supplier checkerSupplier;
  private final int checkBatchSize;
  private final Set initialChecks;
  private final Set succeededChecks;
  private final Set failedChecks;
  private Iterator currentResult = emptyIterator();
  private E currentEntry;

  public AuthzPaginationIterator(
      PaginationIterator source,
      Supplier checkerSupplier,
      int checkBatchSize) {
    this.source = source;
    this.checkerSupplier = checkerSupplier;
    this.checkBatchSize = checkBatchSize;
    initialChecks = new HashSet<>();
    succeededChecks = new HashSet<>();
    failedChecks = new HashSet<>();
  }

  public AuthzPaginationIterator initialCheck(Check initialCheck) {
    initialChecks.add(initialCheck);
    return this;
  }

  protected abstract Set checksForEntry(E entry);

  @Override
  protected final E computeNext() {
    while (true) {
      // Return entries from current result batch
      if (currentResult.hasNext()) {
        E entry = currentResult.next();
        currentEntry = entry;
        return entry;
      }

      // collect batch from source
      List batchEntries = new ArrayList<>(checkBatchSize);
      List> checksForEntries = new ArrayList<>(checkBatchSize);
      Set batchChecks = newHashSetWithExpectedSize(checkBatchSize);
      batchChecks.addAll(initialChecks);
      while (source.hasNext()) {
        E entry = source.next();
        Set checksForEntry = checksForEntry(entry);
        if (checksForEntry.stream().anyMatch(failedChecks::contains)) {
          // At least one of the checks for the current entry failed in a previous
          // batch-access-check,
          // so skip this entry
          continue;
        }

        batchEntries.add(entry);
        checksForEntries.add(checksForEntry);

        checksForEntry.stream().filter(c -> !succeededChecks.contains(c)).forEach(batchChecks::add);

        if (batchChecks.size() >= checkBatchSize) {
          break;
        }
      }

      // Perform batch access check
      Map failed = emptyMap();
      if (!batchChecks.isEmpty()) {
        BatchAccessChecker checker = checkerSupplier.get();
        batchChecks.forEach(checker::can);
        failed = checker.check();
      }

      // Validate the "initial checks" first
      if (!initialChecks.isEmpty()) {
        Map failedInitial = new HashMap<>();
        for (Check initialCheck : initialChecks) {
          String failure = failed.get(initialCheck);
          if (failure != null) {
            failedInitial.put(initialCheck, failure);
          }
        }
        initialChecks.clear();
        if (!failedInitial.isEmpty()) {
          throwForFailedChecks(failedInitial);
        }
      }

      int batchEntryCount = batchEntries.size();
      if (batchEntryCount == 0) {
        // no more entries to check
        return endOfData();
      }

      // Construct the iterator containing the checked entries
      List current = new ArrayList<>(batchEntryCount);
      for (int i = 0; i < batchEntryCount; i++) {
        E entry = batchEntries.get(i);
        Set checksForEntry = checksForEntries.get(i);
        boolean anyFailed = false;
        for (Check check : checksForEntry) {
          if (failed.containsKey(check)) {
            anyFailed = true;
            failedChecks.add(check);
          } else {
            succeededChecks.add(check);
          }
        }
        if (!anyFailed) {
          current.add(entry);
        }
      }
      currentResult = current.iterator();
    }
  }

  @Override
  public final String tokenForCurrent() {
    return source.tokenForEntry(currentEntry);
  }

  @Override
  public final String tokenForEntry(E entry) {
    return source.tokenForEntry(entry);
  }

  @Override
  public final void close() {}
}