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.arrow.flight.auth;
  
  import com.google.protobuf.ByteString;
  import io.grpc.StatusRuntimeException;
  import io.grpc.stub.StreamObserver;
  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;
  
  /** 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();
        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);
      }
    }
  }
  