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;
  
  import com.google.common.util.concurrent.ThreadFactoryBuilder;
  import io.grpc.Server;
  import io.grpc.ServerInterceptors;
  import io.grpc.netty.GrpcSslContexts;
  import io.grpc.netty.NettyServerBuilder;
  import io.netty.channel.EventLoopGroup;
  import io.netty.channel.ServerChannel;
  import io.netty.handler.ssl.ClientAuth;
  import io.netty.handler.ssl.SslContext;
  import io.netty.handler.ssl.SslContextBuilder;
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.IOException;
  import java.io.InputStream;
  import java.lang.reflect.InvocationTargetException;
  import java.net.URI;
  import java.net.URISyntaxException;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.TimeUnit;
  import java.util.function.Consumer;
  import javax.net.ssl.SSLException;
  import org.apache.arrow.flight.auth.ServerAuthHandler;
  import org.apache.arrow.flight.auth.ServerAuthInterceptor;
  import org.apache.arrow.flight.auth2.Auth2Constants;
  import org.apache.arrow.flight.auth2.CallHeaderAuthenticator;
  import org.apache.arrow.flight.auth2.ServerCallHeaderAuthMiddleware;
  import org.apache.arrow.flight.grpc.ServerBackpressureThresholdInterceptor;
  import org.apache.arrow.flight.grpc.ServerInterceptorAdapter;
  import org.apache.arrow.flight.grpc.ServerInterceptorAdapter.KeyFactory;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.util.Preconditions;
  import org.apache.arrow.util.VisibleForTesting;
  
  /**
   * Generic server of flight data that is customized via construction with delegate classes for the
   * actual logic. The server currently uses GRPC as its transport mechanism.
   */
  public class FlightServer implements AutoCloseable {
  
    private static final org.slf4j.Logger logger =
        org.slf4j.LoggerFactory.getLogger(FlightServer.class);
  
    private final Location location;
    private final Server server;
    // The executor used by the gRPC server. We don't use it here, but we do need to clean it up with
    // the server.
    // May be null, if a user-supplied executor was provided (as we do not want to clean that up)
    @VisibleForTesting final ExecutorService grpcExecutor;
  
    /** The maximum size of an individual gRPC message. This effectively disables the limit. */
    static final int MAX_GRPC_MESSAGE_SIZE = Integer.MAX_VALUE;
  
    /** The default number of bytes that can be queued on an output stream before blocking. */
    public static final int DEFAULT_BACKPRESSURE_THRESHOLD = 10 * 1024 * 1024; // 10MB
  
    /** Create a new instance from a gRPC server. For internal use only. */
    private FlightServer(Location location, Server server, ExecutorService grpcExecutor) {
      this.location = location;
      this.server = server;
      this.grpcExecutor = grpcExecutor;
    }
  
    /** Start the server. */
    public FlightServer start() throws IOException {
      server.start();
      return this;
    }
  
    /** Get the port the server is running on (if applicable). */
    public int getPort() {
      return server.getPort();
    }
  
    /** Get the location for this server. */
    public Location getLocation() {
      if (location.getUri().getPort() == 0) {
        // If the server was bound to port 0, replace the port in the location with the real port.
        final URI uri = location.getUri();
        try {
          return new Location(
              new URI(
                  uri.getScheme(),
                  uri.getUserInfo(),
                  uri.getHost(),
                  getPort(),
                  uri.getPath(),
                  uri.getQuery(),
                  uri.getFragment()));
        } catch (URISyntaxException e) {
          // We don't expect this to happen
          throw new RuntimeException(e);
        }
      }
      return location;
    }
  
    /** Block until the server shuts down. */
    public void awaitTermination() throws InterruptedException {
      server.awaitTermination();
    }
  
    /** Request that the server shut down. */
    public void shutdown() {
      server.shutdown();
      if (grpcExecutor != null) {
        grpcExecutor.shutdown();
      }
    }
  
    /**
     * Wait for the server to shut down with a timeout.
     *
     * @return true if the server shut down successfully.
     */
    public boolean awaitTermination(final long timeout, final TimeUnit unit)
        throws InterruptedException {
      return server.awaitTermination(timeout, unit);
    }
  
    /** Shutdown the server, waits for up to 6 seconds for successful shutdown before returning. */
    @Override
    public void close() throws InterruptedException {
      shutdown();
      final boolean terminated = awaitTermination(3000, TimeUnit.MILLISECONDS);
      if (terminated) {
        logger.debug("Server was terminated within 3s");
        return;
      }
  
      // get more aggressive in termination.
      server.shutdownNow();
  
      int count = 0;
      while (!server.isTerminated() && count < 30) {
        count++;
        logger.debug("Waiting for termination");
        Thread.sleep(100);
      }
  
      if (!server.isTerminated()) {
        logger.warn("Couldn't shutdown server, resources likely will be leaked.");
      }
    }
  
    /** Create a builder for a Flight server. */
    public static Builder builder() {
      return new Builder();
    }
  
    /** Create a builder for a Flight server. */
    public static Builder builder(
        BufferAllocator allocator, Location location, FlightProducer producer) {
      return new Builder(allocator, location, producer);
    }
  
    /** A builder for Flight servers. */
    public static final class Builder {
      private BufferAllocator allocator;
      private Location location;
      private FlightProducer producer;
      private final Map<String, Object> builderOptions;
      private ServerAuthHandler authHandler = ServerAuthHandler.NO_OP;
      private CallHeaderAuthenticator headerAuthenticator = CallHeaderAuthenticator.NO_OP;
      private ExecutorService executor = null;
      private int maxInboundMessageSize = MAX_GRPC_MESSAGE_SIZE;
      private int maxHeaderListSize = MAX_GRPC_MESSAGE_SIZE;
      private int backpressureThreshold = DEFAULT_BACKPRESSURE_THRESHOLD;
      private InputStream certChain;
      private InputStream key;
      private InputStream mTlsCACert;
      private SslContext sslContext;
      private final List<KeyFactory<?>> interceptors;
      // Keep track of inserted interceptors
      private final Set<String> interceptorKeys;
  
      Builder() {
        builderOptions = new HashMap<>();
        interceptors = new ArrayList<>();
        interceptorKeys = new HashSet<>();
      }
  
      Builder(BufferAllocator allocator, Location location, FlightProducer producer) {
        this();
        this.allocator = Preconditions.checkNotNull(allocator);
        this.location = Preconditions.checkNotNull(location);
        this.producer = Preconditions.checkNotNull(producer);
      }
  
      /** Create the server for this builder. */
      public FlightServer build() {
        // Add the auth middleware if applicable.
        if (headerAuthenticator != CallHeaderAuthenticator.NO_OP) {
          this.middleware(
              FlightServerMiddleware.Key.of(Auth2Constants.AUTHORIZATION_HEADER),
              new ServerCallHeaderAuthMiddleware.Factory(headerAuthenticator));
        }
  
        this.middleware(FlightConstants.HEADER_KEY, new ServerHeaderMiddleware.Factory());
  
        final NettyServerBuilder builder;
        switch (location.getUri().getScheme()) {
          case LocationSchemes.GRPC_DOMAIN_SOCKET:
            {
              // The implementation is platform-specific, so we have to find the classes at runtime
              builder = NettyServerBuilder.forAddress(location.toSocketAddress());
              try {
                try {
                  // Linux
                  builder.channelType(
                      Class.forName("io.netty.channel.epoll.EpollServerDomainSocketChannel")
                          .asSubclass(ServerChannel.class));
                  final EventLoopGroup elg =
                      Class.forName("io.netty.channel.epoll.EpollEventLoopGroup")
                          .asSubclass(EventLoopGroup.class)
                          .getConstructor()
                          .newInstance();
                  builder.bossEventLoopGroup(elg).workerEventLoopGroup(elg);
                } catch (ClassNotFoundException e) {
                  // BSD
                  builder.channelType(
                      Class.forName("io.netty.channel.kqueue.KQueueServerDomainSocketChannel")
                          .asSubclass(ServerChannel.class));
                  final EventLoopGroup elg =
                      Class.forName("io.netty.channel.kqueue.KQueueEventLoopGroup")
                          .asSubclass(EventLoopGroup.class)
                          .getConstructor()
                          .newInstance();
                  builder.bossEventLoopGroup(elg).workerEventLoopGroup(elg);
                }
              } catch (ClassNotFoundException
                  | InstantiationException
                  | IllegalAccessException
                  | NoSuchMethodException
                  | InvocationTargetException e) {
                throw new UnsupportedOperationException(
                    "Could not find suitable Netty native transport implementation for domain socket address.");
              }
              break;
            }
          case LocationSchemes.GRPC:
          case LocationSchemes.GRPC_INSECURE:
            {
              builder = NettyServerBuilder.forAddress(location.toSocketAddress());
              break;
            }
          case LocationSchemes.GRPC_TLS:
            {
              if (certChain == null) {
                throw new IllegalArgumentException(
                    "Must provide a certificate and key to serve gRPC over TLS");
              }
              builder = NettyServerBuilder.forAddress(location.toSocketAddress());
              break;
            }
          default:
            throw new IllegalArgumentException(
                "Scheme is not supported: " + location.getUri().getScheme());
        }
  
        if (certChain != null) {
          SslContextBuilder sslContextBuilder = GrpcSslContexts.forServer(certChain, key);
  
          if (mTlsCACert != null) {
            sslContextBuilder.clientAuth(ClientAuth.REQUIRE).trustManager(mTlsCACert);
          }
          try {
            sslContext = sslContextBuilder.build();
          } catch (SSLException e) {
            throw new RuntimeException(e);
          } finally {
            closeMTlsCACert();
            closeCertChain();
            closeKey();
          }
  
          builder.sslContext(sslContext);
        }
  
        // Share one executor between the gRPC service, DoPut, and Handshake
        final ExecutorService exec;
        // We only want to have FlightServer close the gRPC executor if we created it here. We should
        // not close
        // user-supplied executors.
        final ExecutorService grpcExecutor;
        if (executor != null) {
          exec = executor;
          grpcExecutor = null;
        } else {
          exec =
              Executors.newCachedThreadPool(
                  // Name threads for better debuggability
                  new ThreadFactoryBuilder()
                      .setNameFormat("flight-server-default-executor-%d")
                      .build());
          grpcExecutor = exec;
        }
  
        final FlightBindingService flightService =
            new FlightBindingService(allocator, producer, authHandler, exec);
        builder
            .executor(exec)
            .maxInboundMessageSize(maxInboundMessageSize)
            .maxInboundMetadataSize(maxHeaderListSize)
            .addService(
                ServerInterceptors.intercept(
                    flightService,
                    new ServerBackpressureThresholdInterceptor(backpressureThreshold),
                    new ServerAuthInterceptor(authHandler)));
  
        // Allow hooking into the gRPC builder. This is not guaranteed to be available on all Arrow
        // versions or
        // Flight implementations.
        builderOptions.computeIfPresent(
            "grpc.builderConsumer",
            (key, builderConsumer) -> {
              final Consumer<NettyServerBuilder> consumer =
                  (Consumer<NettyServerBuilder>) builderConsumer;
              consumer.accept(builder);
              return null;
            });
  
        // Allow explicitly setting some Netty-specific options
        builderOptions.computeIfPresent(
            "netty.channelType",
            (key, channelType) -> {
              builder.channelType((Class<? extends ServerChannel>) channelType);
              return null;
            });
        builderOptions.computeIfPresent(
            "netty.bossEventLoopGroup",
            (key, elg) -> {
              builder.bossEventLoopGroup((EventLoopGroup) elg);
              return null;
            });
        builderOptions.computeIfPresent(
            "netty.workerEventLoopGroup",
            (key, elg) -> {
              builder.workerEventLoopGroup((EventLoopGroup) elg);
              return null;
            });
  
        builder.intercept(new ServerInterceptorAdapter(interceptors));
        return new FlightServer(location, builder.build(), grpcExecutor);
      }
  
      public Builder setMaxHeaderListSize(int maxHeaderListSize) {
        this.maxHeaderListSize = maxHeaderListSize;
        return this;
      }
  
      /**
       * Set the maximum size of a message. Defaults to "unlimited", depending on the underlying
       * transport.
       */
      public Builder maxInboundMessageSize(int maxMessageSize) {
        this.maxInboundMessageSize = maxMessageSize;
        return this;
      }
  
      /**
       * Set the number of bytes that may be queued on a server output stream before writes are
       * blocked.
       */
      public Builder backpressureThreshold(int backpressureThreshold) {
        Preconditions.checkArgument(backpressureThreshold > 0);
        this.backpressureThreshold = backpressureThreshold;
        return this;
      }
  
      /**
       * A small utility function to ensure that InputStream attributes. are closed if they are not
       * null
       *
       * @param stream The InputStream to close (if it is not null).
       */
      private void closeInputStreamIfNotNull(InputStream stream) {
        if (stream != null) {
          try {
            stream.close();
          } catch (IOException expected) {
            // stream closes gracefully, doesn't expect an exception.
          }
        }
      }
  
      /**
       * A small utility function to ensure that the certChain attribute is closed if it is not null.
       * It then sets the attribute to null.
       */
      private void closeCertChain() {
        closeInputStreamIfNotNull(certChain);
        certChain = null;
      }
  
      /**
       * A small utility function to ensure that the key attribute is closed if it is not null. It
       * then sets the attribute to null.
       */
      private void closeKey() {
        closeInputStreamIfNotNull(key);
        key = null;
      }
  
      /**
       * A small utility function to ensure that the mTlsCACert attribute is closed if it is not null.
       * It then sets the attribute to null.
       */
      private void closeMTlsCACert() {
        closeInputStreamIfNotNull(mTlsCACert);
        mTlsCACert = null;
      }
  
      /**
       * Enable TLS on the server.
       *
       * @param certChain The certificate chain to use.
       * @param key The private key to use.
       */
      public Builder useTls(final File certChain, final File key) throws IOException {
        closeCertChain();
        this.certChain = new FileInputStream(certChain);
  
        closeKey();
        this.key = new FileInputStream(key);
  
        return this;
      }
  
      /**
       * Enable Client Verification via mTLS on the server.
       *
       * @param mTlsCACert The CA certificate to use for verifying clients.
       */
      public Builder useMTlsClientVerification(final File mTlsCACert) throws IOException {
        closeMTlsCACert();
        this.mTlsCACert = new FileInputStream(mTlsCACert);
        return this;
      }
  
      /**
       * Enable TLS on the server.
       *
       * @param certChain The certificate chain to use.
       * @param key The private key to use.
       */
      public Builder useTls(final InputStream certChain, final InputStream key) throws IOException {
        closeCertChain();
        this.certChain = certChain;
  
        closeKey();
        this.key = key;
  
        return this;
      }
  
      /**
       * Enable mTLS on the server.
       *
       * @param mTlsCACert The CA certificate to use for verifying clients.
       */
      public Builder useMTlsClientVerification(final InputStream mTlsCACert) throws IOException {
        closeMTlsCACert();
        this.mTlsCACert = mTlsCACert;
        return this;
      }
  
      /**
       * Set the executor used by the server.
       *
       * <p>Flight will NOT take ownership of the executor. The application must clean it up if one is
       * provided. (If not provided, Flight will use a default executor which it will clean up.)
       */
      public Builder executor(ExecutorService executor) {
        this.executor = executor;
        return this;
      }
  
      /** Set the authentication handler. */
      public Builder authHandler(ServerAuthHandler authHandler) {
        this.authHandler = authHandler;
        return this;
      }
  
      /** Set the header-based authentication mechanism. */
      public Builder headerAuthenticator(CallHeaderAuthenticator headerAuthenticator) {
        this.headerAuthenticator = headerAuthenticator;
        return this;
      }
  
      /** Provide a transport-specific option. Not guaranteed to have any effect. */
      public Builder transportHint(final String key, Object option) {
        builderOptions.put(key, option);
        return this;
      }
  
      /**
       * Add a Flight middleware component to inspect and modify requests to this service.
       *
       * @param key An identifier for this middleware component. Service implementations can retrieve
       *     the middleware instance for the current call using {@link
       *     org.apache.arrow.flight.FlightProducer.CallContext}.
       * @param factory A factory for the middleware.
       * @param <T> The middleware type.
       * @throws IllegalArgumentException if the key already exists
       */
      public <T extends FlightServerMiddleware> Builder middleware(
          final FlightServerMiddleware.Key<T> key, final FlightServerMiddleware.Factory<T> factory) {
        if (interceptorKeys.contains(key.key)) {
          throw new IllegalArgumentException("Key already exists: " + key.key);
        }
        interceptors.add(new KeyFactory<>(key, factory));
        interceptorKeys.add(key.key);
        return this;
      }
  
      public Builder allocator(BufferAllocator allocator) {
        this.allocator = Preconditions.checkNotNull(allocator);
        return this;
      }
  
      public Builder location(Location location) {
        this.location = Preconditions.checkNotNull(location);
        return this;
      }
  
      public Builder producer(FlightProducer producer) {
        this.producer = Preconditions.checkNotNull(producer);
        return this;
      }
    }
  }
  