content: /*
   * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
   * (the "License"). You may not use this work except in compliance with the License, which is
   * available at www.apache.org/licenses/LICENSE-2.0
   *
   * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
   * either express or implied, as more fully set forth in the License.
   *
   * See the NOTICE file distributed with this work for information regarding copyright ownership.
   */
  
  package alluxio.grpc;
  
  import alluxio.conf.AlluxioConfiguration;
  import alluxio.conf.PropertyKey;
  import alluxio.security.authentication.AuthenticatedUserInjector;
  import alluxio.security.authentication.AuthenticationServer;
  import alluxio.security.authentication.DefaultAuthenticationServer;
  import alluxio.util.SecurityUtils;
  import alluxio.util.network.tls.SslContextProvider;
  
  import com.google.common.io.Closer;
  import io.grpc.ServerInterceptor;
  import io.grpc.ServerInterceptors;
  import io.grpc.ServerServiceDefinition;
  import io.grpc.netty.NettyServerBuilder;
  import io.grpc.protobuf.services.ProtoReflectionService;
  import io.netty.channel.ChannelOption;
  import io.netty.channel.EventLoopGroup;
  import io.netty.channel.ServerChannel;
  import io.netty.handler.ssl.SslContext;
  
  import java.util.HashSet;
  import java.util.Set;
  import java.util.concurrent.Executor;
  import java.util.concurrent.TimeUnit;
  import java.util.function.Supplier;
  import javax.annotation.Nullable;
  
  /**
   * Provides authenticated gRPC server creation.
   */
  public final class GrpcServerBuilder {
  
    /** Internal netty builder. */
    private final NettyServerBuilder mNettyServerBuilder;
    /** Set of services that this server has. */
    private final Set<ServiceType> mServices = new HashSet<>();
    /** Authentication server instance that will be used by this server. */
    private final AuthenticationServer mAuthenticationServer;
    /** Used to register closers that needs to be called during server shut-down. */
    private final Closer mCloser = Closer.create();
    /** Alluxio configuration.  */
    private final AlluxioConfiguration mConfiguration;
  
    private final boolean mGrpcReflectionEnabled;
  
    private GrpcServerBuilder(GrpcServerAddress serverAddress,
        AuthenticationServer authenticationServer, AlluxioConfiguration conf) {
      mNettyServerBuilder = NettyServerBuilder.forAddress(serverAddress.getSocketAddress());
      mConfiguration = conf;
      mGrpcReflectionEnabled = conf.getBoolean(PropertyKey.GRPC_REFLECTION_ENABLED);
  
      if (conf.getBoolean(alluxio.conf.PropertyKey.NETWORK_TLS_ENABLED)) {
        sslContext(SslContextProvider.Factory.create(mConfiguration).getServerSSLContext());
      }
  
      if (SecurityUtils.isAuthenticationEnabled(mConfiguration)) {
        if (authenticationServer == null) {
          authenticationServer =
              new DefaultAuthenticationServer(serverAddress.getHostName(), mConfiguration);
        }
        addService(new GrpcService(authenticationServer).disableAuthentication()
            .withCloseable(authenticationServer));
      }
      mAuthenticationServer = authenticationServer;
    }
  
    /**
     * Create an new instance of {@link GrpcServerBuilder} with authentication support.
     *
     * @param serverAddress server address
     * @param conf Alluxio configuration
     * @return a new instance of {@link GrpcServerBuilder}
     */
    public static GrpcServerBuilder forAddress(GrpcServerAddress serverAddress,
        AlluxioConfiguration conf) {
      return new GrpcServerBuilder(serverAddress, null, conf);
    }
  
    /**
     * Set the executor for this server.
     *
     * @param executor the executor
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder executor(@Nullable Executor executor) {
      mNettyServerBuilder.executor(executor);
      return this;
    }
  
    /**
     * Sets flow control window.
     *
     * @param flowControlWindow the HTTP2 flow control window
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder flowControlWindow(int flowControlWindow) {
      mNettyServerBuilder.flowControlWindow(flowControlWindow);
      return this;
    }
  
    /**
     * Sets the server keep-alive time.
     *
     * @param keepAliveTime the time to wait after idle before pinging client
     * @param timeUnit unit of the time
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder keepAliveTime(long keepAliveTime, TimeUnit timeUnit) {
      mNettyServerBuilder.keepAliveTime(keepAliveTime, timeUnit);
      return this;
    }
  
    /**
     * Sets the server keep-alive timeout.
     *
     * @param keepAliveTimeout time to wait after pinging client before closing the connection
     * @param timeUnit unit of the timeout
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder keepAliveTimeout(long keepAliveTimeout, TimeUnit timeUnit) {
      mNettyServerBuilder.keepAliveTimeout(keepAliveTimeout, timeUnit);
      return this;
    }
  
    /**
     * Sets the high-bar for client-side keep-alive frequency.
     * Clients pushing this bar will be held back by closing their connections.
     *
     * @param permitKeepAlive permitted client-side keep-alive time
     * @param timeUnit unit of the timeout
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder permitKeepAlive(long permitKeepAlive, TimeUnit timeUnit) {
      mNettyServerBuilder.permitKeepAliveTime(permitKeepAlive, timeUnit);
      return this;
    }
  
    /**
     * Sets the netty channel type.
     *
     * @param channelType the netty channel type for the server
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder channelType(Class<? extends ServerChannel> channelType) {
      mNettyServerBuilder.channelType(channelType);
      return this;
    }
  
    /**
     * Sets a netty channel option.
     *
     * @param <T> channel option type
     * @param option the option to be set
     * @param value the new value
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public <T> GrpcServerBuilder withChildOption(ChannelOption<T> option, T value) {
      mNettyServerBuilder.withChildOption(option, value);
      return this;
    }
  
    /**
     * Sets the boss {@link EventLoopGroup}.
     *
     * @param bossGroup the boss event loop group
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder bossEventLoopGroup(EventLoopGroup bossGroup) {
      mNettyServerBuilder.bossEventLoopGroup(bossGroup);
      return this;
    }
  
    /**
     * Sets the worker {@link EventLoopGroup}.
     *
     * @param workerGroup the worker event loop group
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder workerEventLoopGroup(EventLoopGroup workerGroup) {
      mNettyServerBuilder.workerEventLoopGroup(workerGroup);
      return this;
    }
  
    /**
     * Sets the maximum size of inbound messages.
     * @param messageSize maximum size of the message
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder maxInboundMessageSize(int messageSize) {
      mNettyServerBuilder.maxInboundMessageSize(messageSize);
      return this;
    }
  
    /**
     * Add a service to this server.
     *
     * @param serviceType the type of service
     * @param serviceDefinition the service definition of new service
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder addService(ServiceType serviceType, GrpcService serviceDefinition) {
      mServices.add(serviceType);
      return addService(serviceDefinition);
    }
  
    /**
     * Add a service to this server.
     *
     * @param serviceDefinition the service definition of new service
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder addService(GrpcService serviceDefinition) {
      ServerServiceDefinition service = serviceDefinition.getServiceDefinition();
      if (SecurityUtils.isAuthenticationEnabled(mConfiguration)
          && serviceDefinition.isAuthenticated()) {
        // Intercept the service with authenticated user injector.
        service = ServerInterceptors.intercept(service,
            new AuthenticatedUserInjector(mAuthenticationServer));
      }
      mNettyServerBuilder.addService(service);
      mCloser.register(serviceDefinition.getCloser());
      return this;
    }
  
    /**
     * Adds an interceptor for this server.
     *
     * @param interceptor server interceptor
     * @return an updates instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder intercept(ServerInterceptor interceptor) {
      mNettyServerBuilder.intercept(interceptor);
      return this;
    }
  
    /**
     * Sets TLS context.
     *
     * @param sslContext TLS context
     * @return an updated instance of this {@link GrpcServerBuilder}
     */
    public GrpcServerBuilder sslContext(SslContext sslContext) {
      mNettyServerBuilder.sslContext(sslContext);
      return this;
    }
  
    /**
     * Build the server.
     * It attaches required services and interceptors for authentication.
     *
     * @return the built {@link GrpcServer}
     */
    public GrpcServer build() {
      return build(null);
    }
  
    /**
     * Build the server.
     * It attaches required services and interceptors for authentication.
     *
     * @param nodeStateSupplier a supplier to provide the node state (PRIMARY/STANDBY)
     * @return the built {@link GrpcServer}
     */
    public GrpcServer build(@Nullable Supplier<NodeState> nodeStateSupplier) {
      addService(new GrpcService(new ServiceVersionClientServiceHandler(mServices, nodeStateSupplier))
          .disableAuthentication());
      if (mGrpcReflectionEnabled) {
        // authentication needs to be disabled so that the grpc command line tools can call
        // this reflection endpoint and get the current grpc services and their interfaces.
        addService(new GrpcService(ProtoReflectionService.newInstance()).disableAuthentication());
      }
      return new GrpcServer(mNettyServerBuilder.build(), mAuthenticationServer, mCloser,
          mConfiguration.getMs(PropertyKey.NETWORK_CONNECTION_SERVER_SHUTDOWN_TIMEOUT));
    }
  }
  