HBASE-25051 DIGEST based auth broken for rpc based ConnectionRegistry (#5631)

Signed-off-by: Bryan Beaudreault <bbeaudreault@apache.org>
(cherry picked from commit 16de74c)

Author: Apache9

hbase-client/src/main/java/org/apache/hadoop/hbase/client/AbstractRpcBasedConnectionRegistry.java

@@ -33,22 +33,17 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Predicate;
 import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.hbase.HConstants;
 import org.apache.hadoop.hbase.HRegionLocation;
 import org.apache.hadoop.hbase.RegionLocations;
 import org.apache.hadoop.hbase.ServerName;
 import org.apache.hadoop.hbase.exceptions.ClientExceptionsUtil;
 import org.apache.hadoop.hbase.exceptions.MasterRegistryFetchException;
 import org.apache.hadoop.hbase.ipc.HBaseRpcController;
-import org.apache.hadoop.hbase.ipc.RpcClient;
-import org.apache.hadoop.hbase.ipc.RpcClientFactory;
 import org.apache.hadoop.hbase.ipc.RpcControllerFactory;
 import org.apache.hadoop.hbase.security.User;
 import org.apache.hadoop.hbase.util.FutureUtils;
 import org.apache.yetus.audience.InterfaceAudience;
 
-import org.apache.hbase.thirdparty.com.google.common.base.Preconditions;
-import org.apache.hbase.thirdparty.com.google.common.collect.ImmutableMap;
 import org.apache.hbase.thirdparty.com.google.protobuf.Message;
 import org.apache.hbase.thirdparty.com.google.protobuf.RpcCallback;
 
@@ -79,30 +74,21 @@ abstract class AbstractRpcBasedConnectionRegistry implements ConnectionRegistry
 
   private final int hedgedReadFanOut;
 
-  // Configured list of end points to probe the meta information from.
-  private volatile ImmutableMap<ServerName, ClientMetaService.Interface> addr2Stub;
-
   // RPC client used to talk to the masters.
-  private final RpcClient rpcClient;
+  private final ConnectionRegistryRpcStubHolder rpcStubHolder;
   private final RpcControllerFactory rpcControllerFactory;
-  private final int rpcTimeoutMs;
 
   private final RegistryEndpointsRefresher registryEndpointRefresher;
 
-  protected AbstractRpcBasedConnectionRegistry(Configuration conf,
+  protected AbstractRpcBasedConnectionRegistry(Configuration conf, User user,
     String hedgedReqsFanoutConfigName, String initialRefreshDelaySecsConfigName,
     String refreshIntervalSecsConfigName, String minRefreshIntervalSecsConfigName)
     throws IOException {
     this.hedgedReadFanOut =
       Math.max(1, conf.getInt(hedgedReqsFanoutConfigName, HEDGED_REQS_FANOUT_DEFAULT));
-    rpcTimeoutMs = (int) Math.min(Integer.MAX_VALUE,
-      conf.getLong(HConstants.HBASE_RPC_TIMEOUT_KEY, HConstants.DEFAULT_HBASE_RPC_TIMEOUT));
-    // XXX: we pass cluster id as null here since we do not have a cluster id yet, we have to fetch
-    // this through the master registry...
-    // This is a problem as we will use the cluster id to determine the authentication method
-    rpcClient = RpcClientFactory.createClient(conf, null);
     rpcControllerFactory = RpcControllerFactory.instantiate(conf);
-    populateStubs(getBootstrapNodes(conf));
+    rpcStubHolder = new ConnectionRegistryRpcStubHolder(conf, user, rpcControllerFactory,
+      getBootstrapNodes(conf));
     // could return null here is refresh interval is less than zero
     registryEndpointRefresher =
       RegistryEndpointsRefresher.create(conf, initialRefreshDelaySecsConfigName,
@@ -114,19 +100,7 @@ protected AbstractRpcBasedConnectionRegistry(Configuration conf,
   protected abstract CompletableFuture<Set<ServerName>> fetchEndpoints();
 
   private void refreshStubs() throws IOException {
-    populateStubs(FutureUtils.get(fetchEndpoints()));
-  }
-
-  private void populateStubs(Set<ServerName> addrs) throws IOException {
-    Preconditions.checkNotNull(addrs);
-    ImmutableMap.Builder<ServerName, ClientMetaService.Interface> builder =
-      ImmutableMap.builderWithExpectedSize(addrs.size());
-    User user = User.getCurrent();
-    for (ServerName masterAddr : addrs) {
-      builder.put(masterAddr,
-        ClientMetaService.newStub(rpcClient.createRpcChannel(masterAddr, user, rpcTimeoutMs)));
-    }
-    addr2Stub = builder.build();
+    rpcStubHolder.refreshStubs(() -> FutureUtils.get(fetchEndpoints()));
   }
 
   /**
@@ -211,20 +185,25 @@ private <T extends Message> void groupCall(CompletableFuture<T> future, Set<Serv
 
   protected final <T extends Message> CompletableFuture<T> call(Callable<T> callable,
     Predicate<T> isValidResp, String debug) {
-    ImmutableMap<ServerName, ClientMetaService.Interface> addr2StubRef = addr2Stub;
-    Set<ServerName> servers = addr2StubRef.keySet();
-    List<ClientMetaService.Interface> stubs = new ArrayList<>(addr2StubRef.values());
-    Collections.shuffle(stubs, ThreadLocalRandom.current());
     CompletableFuture<T> future = new CompletableFuture<>();
-    groupCall(future, servers, stubs, 0, callable, isValidResp, debug,
-      new ConcurrentLinkedQueue<>());
+    FutureUtils.addListener(rpcStubHolder.getStubs(), (addr2Stub, error) -> {
+      if (error != null) {
+        future.completeExceptionally(error);
+        return;
+      }
+      Set<ServerName> servers = addr2Stub.keySet();
+      List<ClientMetaService.Interface> stubs = new ArrayList<>(addr2Stub.values());
+      Collections.shuffle(stubs, ThreadLocalRandom.current());
+      groupCall(future, servers, stubs, 0, callable, isValidResp, debug,
+        new ConcurrentLinkedQueue<>());
+    });
     return future;
   }
 
   @RestrictedApi(explanation = "Should only be called in tests", link = "",
       allowedOnPath = ".*/src/test/.*")
-  Set<ServerName> getParsedServers() {
-    return addr2Stub.keySet();
+  Set<ServerName> getParsedServers() throws IOException {
+    return FutureUtils.get(rpcStubHolder.getStubs()).keySet();
   }
 
   /**
@@ -277,8 +256,8 @@ public void close() {
       if (registryEndpointRefresher != null) {
         registryEndpointRefresher.stop();
       }
-      if (rpcClient != null) {
-        rpcClient.close();
+      if (rpcStubHolder != null) {
+        rpcStubHolder.close();
       }
     }, getClass().getSimpleName() + ".close");
   }

hbase-client/src/main/java/org/apache/hadoop/hbase/client/ClusterIdFetcher.java

@@ -0,0 +1,134 @@
+/*
+ * 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.hadoop.hbase.client;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.ServerName;
+import org.apache.hadoop.hbase.ipc.HBaseRpcController;
+import org.apache.hadoop.hbase.ipc.RpcClient;
+import org.apache.hadoop.hbase.ipc.RpcClientFactory;
+import org.apache.hadoop.hbase.ipc.RpcControllerFactory;
+import org.apache.hadoop.hbase.security.User;
+import org.apache.hadoop.hbase.util.FutureUtils;
+import org.apache.yetus.audience.InterfaceAudience;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.hbase.thirdparty.com.google.protobuf.RpcCallback;
+import org.apache.hbase.thirdparty.com.google.protobuf.RpcChannel;
+
+import org.apache.hadoop.hbase.shaded.protobuf.generated.RegistryProtos.ConnectionRegistryService;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.RegistryProtos.GetConnectionRegistryRequest;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.RegistryProtos.GetConnectionRegistryResponse;
+
+/**
+ * Fetch cluster id through special preamble header.
+ * <p>
+ * An instance of this class should only be used once, like:
+ *
+ * <pre>
+ * new ClusterIdFetcher().fetchClusterId()
+ * </pre>
+ *
+ * Calling the fetchClusterId multiple times will lead unexpected behavior.
+ * <p>
+ * See HBASE-25051 for more details.
+ */
+@InterfaceAudience.Private
+class ClusterIdFetcher {
+
+  private static final Logger LOG = LoggerFactory.getLogger(ClusterIdFetcher.class);
+
+  private final List<ServerName> bootstrapServers;
+
+  private final User user;
+
+  private final RpcClient rpcClient;
+
+  private final RpcControllerFactory rpcControllerFactory;
+
+  private final CompletableFuture<String> future;
+
+  ClusterIdFetcher(Configuration conf, User user, RpcControllerFactory rpcControllerFactory,
+    Set<ServerName> bootstrapServers) {
+    this.user = user;
+    // use null cluster id here as we do not know the cluster id yet, we will fetch it through this
+    // rpc client
+    this.rpcClient = RpcClientFactory.createClient(conf, null);
+    this.rpcControllerFactory = rpcControllerFactory;
+    this.bootstrapServers = new ArrayList<ServerName>(bootstrapServers);
+    // shuffle the bootstrap servers so we will not always fetch from the same one
+    Collections.shuffle(this.bootstrapServers);
+    future = new CompletableFuture<String>();
+  }
+
+  /**
+   * Try get cluster id from the server with the given {@code index} in {@link #bootstrapServers}.
+   */
+  private void getClusterId(int index) {
+    ServerName server = bootstrapServers.get(index);
+    LOG.debug("Going to request {} for getting cluster id", server);
+    // user and rpcTimeout are both not important here, as we will not actually send any rpc calls
+    // out, only a preamble connection header, but if we pass null as user, there will be NPE in
+    // some code paths...
+    RpcChannel channel = rpcClient.createRpcChannel(server, user, 0);
+    ConnectionRegistryService.Interface stub = ConnectionRegistryService.newStub(channel);
+    HBaseRpcController controller = rpcControllerFactory.newController();
+    stub.getConnectionRegistry(controller, GetConnectionRegistryRequest.getDefaultInstance(),
+      new RpcCallback<GetConnectionRegistryResponse>() {
+
+        @Override
+        public void run(GetConnectionRegistryResponse resp) {
+          if (!controller.failed()) {
+            LOG.debug("Got connection registry info: {}", resp);
+            future.complete(resp.getClusterId());
+            return;
+          }
+          if (ConnectionUtils.isUnexpectedPreambleHeaderException(controller.getFailed())) {
+            // this means we have connected to an old server where it does not support passing
+            // cluster id through preamble connnection header, so we fallback to use null
+            // cluster id, which is the old behavior
+            LOG.debug("Failed to get connection registry info, should be an old server,"
+              + " fallback to use null cluster id", controller.getFailed());
+            future.complete(null);
+          } else {
+            LOG.debug("Failed to get connection registry info", controller.getFailed());
+            if (index == bootstrapServers.size() - 1) {
+              future.completeExceptionally(controller.getFailed());
+            } else {
+              // try next bootstrap server
+              getClusterId(index + 1);
+            }
+          }
+        }
+      });
+
+  }
+
+  CompletableFuture<String> fetchClusterId() {
+    getClusterId(0);
+    // close the rpc client after we finish the request
+    FutureUtils.addListener(future, (r, e) -> rpcClient.close());
+    return future;
+  }
+}

hbase-client/src/main/java/org/apache/hadoop/hbase/client/ConnectionFactory.java

@@ -338,7 +338,7 @@ public static CompletableFuture<AsyncConnection> createAsyncConnection(Configura
     final User user, Map<String, byte[]> connectionAttributes) {
     return TraceUtil.tracedFuture(() -> {
       CompletableFuture<AsyncConnection> future = new CompletableFuture<>();
-      ConnectionRegistry registry = ConnectionRegistryFactory.getRegistry(conf);
+      ConnectionRegistry registry = ConnectionRegistryFactory.getRegistry(conf, user);
       addListener(registry.getClusterId(), (clusterId, error) -> {
         if (error != null) {
           registry.close();

hbase-client/src/main/java/org/apache/hadoop/hbase/client/ConnectionOverAsyncConnection.java

@@ -29,8 +29,8 @@
 import org.apache.hadoop.hbase.ServerName;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.log.HBaseMarkers;
-import org.apache.hadoop.hbase.util.ConcurrentMapUtils.IOExceptionSupplier;
 import org.apache.hadoop.hbase.util.FutureUtils;
+import org.apache.hadoop.hbase.util.IOExceptionSupplier;
 import org.apache.yetus.audience.InterfaceAudience;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

hbase-client/src/main/java/org/apache/hadoop/hbase/client/ConnectionRegistryFactory.java

@@ -20,6 +20,7 @@
 import static org.apache.hadoop.hbase.HConstants.CLIENT_CONNECTION_REGISTRY_IMPL_CONF_KEY;
 
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.security.User;
 import org.apache.hadoop.hbase.util.ReflectionUtils;
 import org.apache.yetus.audience.InterfaceAudience;
 
@@ -33,10 +34,10 @@ private ConnectionRegistryFactory() {
   }
 
   /** Returns The connection registry implementation to use. */
-  static ConnectionRegistry getRegistry(Configuration conf) {
+  static ConnectionRegistry getRegistry(Configuration conf, User user) {
     Class<? extends ConnectionRegistry> clazz =
       conf.getClass(CLIENT_CONNECTION_REGISTRY_IMPL_CONF_KEY, RpcConnectionRegistry.class,
         ConnectionRegistry.class);
-    return ReflectionUtils.newInstance(clazz, conf);
+    return ReflectionUtils.newInstance(clazz, conf, user);
   }
 }