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.auth2;
  
  import io.grpc.StatusRuntimeException;
  import io.grpc.stub.StreamObserver;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutionException;
  import org.apache.arrow.flight.CallStatus;
  import org.apache.arrow.flight.FlightRuntimeException;
  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 executing a handshake with a FlightServer. */
  public class ClientHandshakeWrapper {
    private static final org.slf4j.Logger logger =
        org.slf4j.LoggerFactory.getLogger(ClientHandshakeWrapper.class);
  
    /**
     * Do handshake for a client. The stub will be authenticated after this method returns.
     *
     * @param stub The service stub.
     */
    public static void doClientHandshake(FlightServiceStub stub) {
      final HandshakeObserver observer = new HandshakeObserver();
      try {
        observer.requestObserver = stub.handshake(observer);
        observer.requestObserver.onNext(HandshakeRequest.newBuilder().build());
        observer.requestObserver.onCompleted();
        try {
          if (!observer.completed.get()) {
            // TODO: ARROW-5681
            throw CallStatus.UNAUTHENTICATED.toRuntimeException();
          }
        } catch (InterruptedException ex) {
          Thread.currentThread().interrupt();
          throw ex;
        } catch (ExecutionException ex) {
          final FlightRuntimeException wrappedException = StatusUtils.fromThrowable(ex.getCause());
          logger.error("Failed on completing future", wrappedException);
          throw wrappedException;
        }
      } catch (StatusRuntimeException sre) {
        logger.error("Failed with SRE", sre);
        throw StatusUtils.fromGrpcRuntimeException(sre);
      } catch (Throwable ex) {
        logger.error("Failed with unknown", ex);
        if (ex instanceof FlightRuntimeException) {
          throw (FlightRuntimeException) ex;
        }
        throw StatusUtils.fromThrowable(ex);
      }
    }
  
    private static class HandshakeObserver implements StreamObserver<HandshakeResponse> {
  
      private volatile StreamObserver<HandshakeRequest> requestObserver;
      private final CompletableFuture<Boolean> completed;
  
      public HandshakeObserver() {
        super();
        completed = new CompletableFuture<>();
      }
  
      @Override
      public void onNext(HandshakeResponse value) {}
  
      @Override
      public void onError(Throwable t) {
        completed.completeExceptionally(t);
      }
  
      @Override
      public void onCompleted() {
        completed.complete(true);
      }
    }
  }
  