RATIS-2331. Reuse SslContext in gRPC. (#1288)

Author: szetszwo

ratis-common/src/main/java/org/apache/ratis/util/LifeCycle.java

@@ -117,6 +117,9 @@ static void validate(Object name, State from, State to) {
       if (LOG.isTraceEnabled()) {
         LOG.trace("TRACE", new Throwable());
       }
+      if (to == EXCEPTION) {
+        LOG.error("{} has failed ({} -> {})", name, from, to, new Throwable("TRACE"));
+      }
 
       Preconditions.assertTrue(isValid(from, to),
           "ILLEGAL TRANSITION: In %s, %s -> %s", name, from, to);

ratis-grpc/src/main/java/org/apache/ratis/grpc/GrpcFactory.java

@@ -32,11 +32,15 @@
 import org.apache.ratis.server.leader.FollowerInfo;
 import org.apache.ratis.server.leader.LeaderState;
 import org.apache.ratis.thirdparty.io.netty.buffer.PooledByteBufAllocator;
+import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContext;
 import org.apache.ratis.util.JavaUtils;
+import org.apache.ratis.util.MemoizedSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 public class GrpcFactory implements ServerFactory, ClientFactory {
 
@@ -65,19 +69,32 @@ static boolean checkPooledByteBufAllocatorUseCacheForAllThreads(Consumer<String>
     return value;
   }
 
-  private final GrpcServices.Customizer servicesCustomizer;
+  static final BiFunction<GrpcTlsConfig, SslContext, SslContext> BUILD_SSL_CONTEXT_FOR_SERVER
+      = (tlsConf, defaultContext) -> tlsConf == null ? defaultContext : GrpcUtil.buildSslContextForServer(tlsConf);
+
+  static final BiFunction<GrpcTlsConfig, SslContext, SslContext> BUILD_SSL_CONTEXT_FOR_CLIENT
+      = (tlsConf, defaultContext) -> tlsConf == null ? defaultContext : GrpcUtil.buildSslContextForClient(tlsConf);
 
-  private final GrpcTlsConfig tlsConfig;
-  private final GrpcTlsConfig adminTlsConfig;
-  private final GrpcTlsConfig clientTlsConfig;
-  private final GrpcTlsConfig serverTlsConfig;
+  static final class SslContexts {
+    private final SslContext adminSslContext;
+    private final SslContext clientSslContext;
+    private final SslContext serverSslContext;
 
-  public static Parameters newRaftParameters(GrpcTlsConfig conf) {
-    final Parameters p = new Parameters();
-    GrpcConfigKeys.TLS.setConf(p, conf);
-    return p;
+    private SslContexts(GrpcTlsConfig tlsConfig, GrpcTlsConfig adminTlsConfig,
+        GrpcTlsConfig clientTlsConfig, GrpcTlsConfig serverTlsConfig,
+        BiFunction<GrpcTlsConfig, SslContext, SslContext> buildMethod) {
+      final SslContext defaultSslContext = buildMethod.apply(tlsConfig, null);
+      this.adminSslContext = buildMethod.apply(adminTlsConfig, defaultSslContext);
+      this.clientSslContext = buildMethod.apply(clientTlsConfig, defaultSslContext);
+      this.serverSslContext = buildMethod.apply(serverTlsConfig, defaultSslContext);
+    }
   }
 
+  private final GrpcServices.Customizer servicesCustomizer;
+
+  private final Supplier<SslContexts> forServerSupplier;
+  private final Supplier<SslContexts> forClientSupplier;
+
   public GrpcFactory(Parameters parameters) {
     this(GrpcConfigKeys.Server.servicesCustomizer(parameters),
         GrpcConfigKeys.TLS.conf(parameters),
@@ -87,35 +104,15 @@ public GrpcFactory(Parameters parameters) {
     );
   }
 
-  public GrpcFactory(GrpcTlsConfig tlsConfig) {
-    this(null, tlsConfig, null, null, null);
-  }
-
   private GrpcFactory(GrpcServices.Customizer servicesCustomizer,
       GrpcTlsConfig tlsConfig, GrpcTlsConfig adminTlsConfig,
       GrpcTlsConfig clientTlsConfig, GrpcTlsConfig serverTlsConfig) {
     this.servicesCustomizer = servicesCustomizer;
 
-    this.tlsConfig = tlsConfig;
-    this.adminTlsConfig = adminTlsConfig;
-    this.clientTlsConfig = clientTlsConfig;
-    this.serverTlsConfig = serverTlsConfig;
-  }
-
-  public GrpcTlsConfig getTlsConfig() {
-    return tlsConfig;
-  }
-
-  public GrpcTlsConfig getAdminTlsConfig() {
-    return adminTlsConfig != null ? adminTlsConfig : tlsConfig;
-  }
-
-  public GrpcTlsConfig getClientTlsConfig() {
-    return clientTlsConfig != null ? clientTlsConfig : tlsConfig;
-  }
-
-  public GrpcTlsConfig getServerTlsConfig() {
-    return serverTlsConfig != null ? serverTlsConfig : tlsConfig;
+    this.forServerSupplier = MemoizedSupplier.valueOf(() -> new SslContexts(
+        tlsConfig, adminTlsConfig, clientTlsConfig, serverTlsConfig, BUILD_SSL_CONTEXT_FOR_SERVER));
+    this.forClientSupplier = MemoizedSupplier.valueOf(() -> new SslContexts(
+        tlsConfig, adminTlsConfig, clientTlsConfig, serverTlsConfig, BUILD_SSL_CONTEXT_FOR_CLIENT));
   }
 
   @Override
@@ -131,19 +128,24 @@ public LogAppender newLogAppender(RaftServer.Division server, LeaderState state,
   @Override
   public GrpcServices newRaftServerRpc(RaftServer server) {
     checkPooledByteBufAllocatorUseCacheForAllThreads(LOG::info);
+
+    final SslContexts forServer = forServerSupplier.get();
+    final SslContexts forClient = forClientSupplier.get();
     return GrpcServicesImpl.newBuilder()
         .setServer(server)
         .setCustomizer(servicesCustomizer)
-        .setAdminTlsConfig(getAdminTlsConfig())
-        .setServerTlsConfig(getServerTlsConfig())
-        .setClientTlsConfig(getClientTlsConfig())
+        .setAdminSslContext(forServer.adminSslContext)
+        .setServerSslContextForServer(forServer.serverSslContext)
+        .setServerSslContextForClient(forClient.serverSslContext)
+        .setClientSslContext(forServer.clientSslContext)
         .build();
   }
 
   @Override
   public GrpcClientRpc newRaftClientRpc(ClientId clientId, RaftProperties properties) {
     checkPooledByteBufAllocatorUseCacheForAllThreads(LOG::debug);
-    return new GrpcClientRpc(clientId, properties,
-        getAdminTlsConfig(), getClientTlsConfig());
+
+    final SslContexts forClient = forClientSupplier.get();
+    return new GrpcClientRpc(clientId, properties, forClient.adminSslContext, forClient.clientSslContext);
   }
 }

ratis-grpc/src/main/java/org/apache/ratis/grpc/GrpcUtil.java

@@ -28,7 +28,10 @@
 import org.apache.ratis.thirdparty.io.grpc.Metadata;
 import org.apache.ratis.thirdparty.io.grpc.Status;
 import org.apache.ratis.thirdparty.io.grpc.StatusRuntimeException;
+import org.apache.ratis.thirdparty.io.grpc.netty.GrpcSslContexts;
 import org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver;
+import org.apache.ratis.thirdparty.io.netty.handler.ssl.ClientAuth;
+import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContext;
 import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContextBuilder;
 import org.apache.ratis.util.IOUtils;
 import org.apache.ratis.util.JavaUtils;
@@ -39,13 +42,16 @@
 import org.slf4j.LoggerFactory;
 
 import javax.net.ssl.KeyManager;
+import javax.net.ssl.SSLException;
 import javax.net.ssl.TrustManager;
 import java.io.IOException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
+import static org.apache.ratis.thirdparty.io.netty.handler.ssl.SslProvider.OPENSSL;
+
 public interface GrpcUtil {
   Logger LOG = LoggerFactory.getLogger(GrpcUtil.class);
 
@@ -299,4 +305,38 @@ static void setKeyManager(SslContextBuilder b, KeyManagerConf keyManagerConfig)
       b.keyManager(privateKey.get(), certificates.get());
     }
   }
+
+  static SslContext buildSslContextForServer(GrpcTlsConfig tlsConf) {
+    if (tlsConf == null) {
+      return null;
+    }
+    SslContextBuilder b = initSslContextBuilderForServer(tlsConf.getKeyManager());
+    if (tlsConf.getMtlsEnabled()) {
+      b.clientAuth(ClientAuth.REQUIRE);
+      setTrustManager(b, tlsConf.getTrustManager());
+    }
+    b = GrpcSslContexts.configure(b, OPENSSL);
+    try {
+      return b.build();
+    } catch (Exception e) {
+      throw new IllegalArgumentException("Failed to buildSslContextForServer from tlsConfig " + tlsConf, e);
+    }
+  }
+
+  static SslContext buildSslContextForClient(GrpcTlsConfig tlsConf) {
+    if (tlsConf == null) {
+      return null;
+    }
+
+    final SslContextBuilder b = GrpcSslContexts.forClient();
+    setTrustManager(b, tlsConf.getTrustManager());
+    if (tlsConf.getMtlsEnabled()) {
+      setKeyManager(b, tlsConf.getKeyManager());
+    }
+    try {
+      return b.build();
+    } catch (SSLException e) {
+      throw new IllegalArgumentException("Failed to buildSslContextForClient from tlsConfig " + tlsConf, e);
+    }
+  }
 }

ratis-grpc/src/main/java/org/apache/ratis/grpc/client/GrpcClientProtocolClient.java

@@ -21,7 +21,6 @@
 import org.apache.ratis.client.impl.ClientProtoUtils;
 import org.apache.ratis.conf.RaftProperties;
 import org.apache.ratis.grpc.GrpcConfigKeys;
-import org.apache.ratis.grpc.GrpcTlsConfig;
 import org.apache.ratis.grpc.GrpcUtil;
 import org.apache.ratis.grpc.metrics.intercept.client.MetricClientInterceptor;
 import org.apache.ratis.proto.RaftProtos.GroupInfoReplyProto;
@@ -49,11 +48,10 @@
 import org.apache.ratis.protocol.exceptions.TimeoutIOException;
 import org.apache.ratis.thirdparty.io.grpc.ManagedChannel;
 import org.apache.ratis.thirdparty.io.grpc.StatusRuntimeException;
-import org.apache.ratis.thirdparty.io.grpc.netty.GrpcSslContexts;
 import org.apache.ratis.thirdparty.io.grpc.netty.NegotiationType;
 import org.apache.ratis.thirdparty.io.grpc.netty.NettyChannelBuilder;
 import org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver;
-import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContextBuilder;
+import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContext;
 import org.apache.ratis.util.CollectionUtils;
 import org.apache.ratis.util.JavaUtils;
 import org.apache.ratis.util.SizeInBytes;
@@ -97,7 +95,7 @@ public class GrpcClientProtocolClient implements Closeable {
   private final MetricClientInterceptor metricClientInterceptor;
 
   GrpcClientProtocolClient(ClientId id, RaftPeer target, RaftProperties properties,
-      GrpcTlsConfig adminTlsConfig, GrpcTlsConfig clientTlsConfig) {
+      SslContext adminSslContext, SslContext clientSslContext) {
     this.name = JavaUtils.memoize(() -> id + "->" + target.getId());
     this.target = target;
     final SizeInBytes flowControlWindow = GrpcConfigKeys.flowControlWindow(properties, LOG::debug);
@@ -110,11 +108,9 @@ public class GrpcClientProtocolClient implements Closeable {
         .filter(x -> !x.isEmpty()).orElse(target.getAddress());
     final boolean separateAdminChannel = !Objects.equals(clientAddress, adminAddress);
 
-    clientChannel = buildChannel(clientAddress, clientTlsConfig,
-        flowControlWindow, maxMessageSize);
+    clientChannel = buildChannel(clientAddress, clientSslContext, flowControlWindow, maxMessageSize);
     adminChannel = separateAdminChannel
-        ? buildChannel(adminAddress, adminTlsConfig,
-            flowControlWindow, maxMessageSize)
+        ? buildChannel(adminAddress, adminSslContext, flowControlWindow, maxMessageSize)
         : clientChannel;
 
     asyncStub = RaftClientProtocolServiceGrpc.newStub(clientChannel);
@@ -124,26 +120,16 @@ public class GrpcClientProtocolClient implements Closeable {
         RaftClientConfigKeys.Rpc.watchRequestTimeout(properties);
   }
 
-  private ManagedChannel buildChannel(String address, GrpcTlsConfig tlsConf,
+  private ManagedChannel buildChannel(String address, SslContext sslContext,
       SizeInBytes flowControlWindow, SizeInBytes maxMessageSize) {
     NettyChannelBuilder channelBuilder =
         NettyChannelBuilder.forTarget(address);
     // ignore any http proxy for grpc
     channelBuilder.proxyDetector(uri -> null);
 
-    if (tlsConf != null) {
+    if (sslContext != null) {
       LOG.debug("Setting TLS for {}", address);
-      SslContextBuilder sslContextBuilder = GrpcSslContexts.forClient();
-      GrpcUtil.setTrustManager(sslContextBuilder, tlsConf.getTrustManager());
-      if (tlsConf.getMtlsEnabled()) {
-        GrpcUtil.setKeyManager(sslContextBuilder, tlsConf.getKeyManager());
-      }
-      try {
-        channelBuilder.useTransportSecurity().sslContext(
-            sslContextBuilder.build());
-      } catch (Exception ex) {
-        throw new RuntimeException(ex);
-      }
+      channelBuilder.useTransportSecurity().sslContext(sslContext);
     } else {
       channelBuilder.negotiationType(NegotiationType.PLAINTEXT);
     }