import quincy 17.2.0

[ceph.git] / ceph / src / arrow / java / flight / flight-core / src / main / java / org / apache / arrow / flight / auth / ClientAuthWrapper.java

/*
 * 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.arrow.flight.auth;

import java.util.Iterator;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;

import org.apache.arrow.flight.auth.ClientAuthHandler.ClientAuthSender;
import org.apache.arrow.flight.grpc.StatusUtils;
import org.apache.arrow.flight.impl.Flight.HandshakeRequest;
import org.apache.arrow.flight.impl.Flight.HandshakeResponse;
import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceStub;

import com.google.protobuf.ByteString;

import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;

/**
 * Utility class for performing authorization over using a GRPC stub.
 */
public class ClientAuthWrapper {

  /**
   * Do client auth for a client.  The stub will be authenticated after this method returns.
   *
   * @param authHandler The handler to use.
   * @param stub The service stub.
   */
  public static void doClientAuth(ClientAuthHandler authHandler, FlightServiceStub stub) {
    AuthObserver observer = new AuthObserver();
    try {
      observer.responseObserver = stub.handshake(observer);
      authHandler.authenticate(observer.sender, observer.iter);
      if (!observer.sender.errored) {
        observer.responseObserver.onCompleted();
      }
    } catch (StatusRuntimeException sre) {
      throw StatusUtils.fromGrpcRuntimeException(sre);
    }
    try {
      if (!observer.completed.get()) {
        // TODO: ARROW-5681
        throw new RuntimeException("Unauthenticated");
      }
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    } catch (ExecutionException e) {
      throw StatusUtils.fromThrowable(e.getCause());
    }
  }

  private static class AuthObserver implements StreamObserver<HandshakeResponse> {

    private volatile StreamObserver<HandshakeRequest> responseObserver;
    private final LinkedBlockingQueue<byte[]> messages = new LinkedBlockingQueue<>();
    private final AuthSender sender = new AuthSender();
    private CompletableFuture<Boolean> completed;

    public AuthObserver() {
      super();
      completed = new CompletableFuture<>();
    }

    @Override
    public void onNext(HandshakeResponse value) {
      ByteString payload = value.getPayload();
      if (payload != null) {
        messages.add(payload.toByteArray());
      }
    }

    private Iterator<byte[]> iter = new Iterator<byte[]>() {

      @Override
      public byte[] next() {
        while (!completed.isDone() || !messages.isEmpty()) {
          byte[] bytes = messages.poll();
          if (bytes == null) {
            // busy wait.
            continue;
          } else {
            return bytes;
          }
        }

        if (completed.isCompletedExceptionally()) {
          // Preserve prior exception behavior
          // TODO: with ARROW-5681, throw an appropriate Flight exception if gRPC raised an exception
          try {
            completed.get();
          } catch (InterruptedException e) {
            throw new RuntimeException(e);
          } catch (ExecutionException e) {
            if (e.getCause() instanceof StatusRuntimeException) {
              throw (StatusRuntimeException) e.getCause();
            }
            throw new RuntimeException(e);
          }
        }

        throw new IllegalStateException("You attempted to retrieve messages after there were none.");
      }

      @Override
      public boolean hasNext() {
        return !messages.isEmpty();
      }
    };

    @Override
    public void onError(Throwable t) {
      completed.completeExceptionally(t);
    }

    private class AuthSender implements ClientAuthSender {

      private boolean errored = false;

      @Override
      public void send(byte[] payload) {
        try {
          responseObserver.onNext(HandshakeRequest.newBuilder()
              .setPayload(ByteString.copyFrom(payload))
              .build());
        } catch (StatusRuntimeException sre) {
          throw StatusUtils.fromGrpcRuntimeException(sre);
        }
      }

      @Override
      public void onError(Throwable cause) {
        this.errored = true;
        responseObserver.onError(StatusUtils.toGrpcException(cause));
      }

    }

    @Override
    public void onCompleted() {
      completed.complete(true);
    }
  }

}