HBASE-29666:Additional key management APIs

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

@@ -2669,5 +2669,23 @@ List<LogEntry> getLogEntries(Set<ServerName> serverNames, String logType, Server
    * Refresh the system key cache on all specified region servers.
    * @param regionServers the list of region servers to refresh the system key cache on
    */
-  void refreshSystemKeyCacheOnServers(Set<ServerName> regionServers) throws IOException;
+  void refreshSystemKeyCacheOnServers(List<ServerName> regionServers) throws IOException;
+
+  /**
+   * Eject a specific managed key entry from the managed key data cache on all specified region
+   * servers.
+   * @param regionServers the list of region servers to eject the managed key entry from
+   * @param keyCustodian  the key custodian
+   * @param keyNamespace  the key namespace
+   * @param keyMetadata   the key metadata
+   */
+  void ejectManagedKeyDataCacheEntryOnServers(List<ServerName> regionServers, byte[] keyCustodian,
+    String keyNamespace, String keyMetadata) throws IOException;
+
+  /**
+   * Clear all entries in the managed key data cache on all specified region servers without having
+   * to restart the process.
+   * @param regionServers the list of region servers to clear the managed key data cache on
+   */
+  void clearManagedKeyDataCacheOnServers(List<ServerName> regionServers) throws IOException;
 }

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

@@ -1148,7 +1148,19 @@ public void restoreBackupSystemTable(String snapshotName) throws IOException {
   }
 
   @Override
-  public void refreshSystemKeyCacheOnServers(Set<ServerName> regionServers) throws IOException {
+  public void refreshSystemKeyCacheOnServers(List<ServerName> regionServers) throws IOException {
     get(admin.refreshSystemKeyCacheOnServers(regionServers));
   }
+
+  @Override
+  public void ejectManagedKeyDataCacheEntryOnServers(List<ServerName> regionServers,
+    byte[] keyCustodian, String keyNamespace, String keyMetadata) throws IOException {
+    get(admin.ejectManagedKeyDataCacheEntryOnServers(regionServers, keyCustodian, keyNamespace,
+      keyMetadata));
+  }
+
+  @Override
+  public void clearManagedKeyDataCacheOnServers(List<ServerName> regionServers) throws IOException {
+    get(admin.clearManagedKeyDataCacheOnServers(regionServers));
+  }
 }

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

@@ -1879,5 +1879,22 @@ CompletableFuture<List<LogEntry>> getLogEntries(Set<ServerName> serverNames, Str
    * Refresh the system key cache on all specified region servers.
    * @param regionServers the list of region servers to refresh the system key cache on
    */
-  CompletableFuture<Void> refreshSystemKeyCacheOnServers(Set<ServerName> regionServers);
+  CompletableFuture<Void> refreshSystemKeyCacheOnServers(List<ServerName> regionServers);
+
+  /**
+   * Eject a specific managed key entry from the managed key data cache on all specified region
+   * servers.
+   * @param regionServers the list of region servers to eject the managed key entry from
+   * @param keyCustodian  the key custodian
+   * @param keyNamespace  the key namespace
+   * @param keyMetadata   the key metadata
+   */
+  CompletableFuture<Void> ejectManagedKeyDataCacheEntryOnServers(List<ServerName> regionServers,
+    byte[] keyCustodian, String keyNamespace, String keyMetadata);
+
+  /**
+   * Clear all entries in the managed key data cache on all specified region servers.
+   * @param regionServers the list of region servers to clear the managed key data cache on
+   */
+  CompletableFuture<Void> clearManagedKeyDataCacheOnServers(List<ServerName> regionServers);
 }

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

@@ -687,10 +687,22 @@ public CompletableFuture<Void> updateConfiguration(String groupName) {
   }
 
   @Override
-  public CompletableFuture<Void> refreshSystemKeyCacheOnServers(Set<ServerName> regionServers) {
+  public CompletableFuture<Void> refreshSystemKeyCacheOnServers(List<ServerName> regionServers) {
     return wrap(rawAdmin.refreshSystemKeyCacheOnServers(regionServers));
   }
 
+  @Override
+  public CompletableFuture<Void> ejectManagedKeyDataCacheEntryOnServers(
+    List<ServerName> regionServers, byte[] keyCustodian, String keyNamespace, String keyMetadata) {
+    return wrap(rawAdmin.ejectManagedKeyDataCacheEntryOnServers(regionServers, keyCustodian,
+      keyNamespace, keyMetadata));
+  }
+
+  @Override
+  public CompletableFuture<Void> clearManagedKeyDataCacheOnServers(List<ServerName> regionServers) {
+    return wrap(rawAdmin.clearManagedKeyDataCacheOnServers(regionServers));
+  }
+
   @Override
   public CompletableFuture<Void> rollWALWriter(ServerName serverName) {
     return wrap(rawAdmin.rollWALWriter(serverName));

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

@@ -76,6 +76,7 @@
 import org.apache.hadoop.hbase.client.replication.TableCFs;
 import org.apache.hadoop.hbase.client.security.SecurityCapability;
 import org.apache.hadoop.hbase.exceptions.DeserializationException;
+import org.apache.hadoop.hbase.io.crypto.ManagedKeyData;
 import org.apache.hadoop.hbase.ipc.HBaseRpcController;
 import org.apache.hadoop.hbase.net.Address;
 import org.apache.hadoop.hbase.quotas.QuotaFilter;
@@ -152,6 +153,8 @@
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.EmptyMsg;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.LastHighestWalFilenum;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.ManagedKeyEntryRequest;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.ManagedKeyRequest;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.NameStringPair;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.ProcedureDescription;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.RegionSpecifier.RegionSpecifierType;
@@ -4665,7 +4668,7 @@ MasterProtos.RestoreBackupSystemTableResponse> procedureCall(request,
   }
 
   @Override
-  public CompletableFuture<Void> refreshSystemKeyCacheOnServers(Set<ServerName> regionServers) {
+  public CompletableFuture<Void> refreshSystemKeyCacheOnServers(List<ServerName> regionServers) {
     CompletableFuture<Void> future = new CompletableFuture<>();
     List<CompletableFuture<Void>> futures =
       regionServers.stream().map(this::refreshSystemKeyCache).collect(Collectors.toList());
@@ -4687,4 +4690,62 @@ private CompletableFuture<Void> refreshSystemKeyCache(ServerName serverName) {
         (s, c, req, done) -> s.refreshSystemKeyCache(controller, req, done), resp -> null))
       .serverName(serverName).call();
   }
+
+  @Override
+  public CompletableFuture<Void> ejectManagedKeyDataCacheEntryOnServers(
+    List<ServerName> regionServers, byte[] keyCustodian, String keyNamespace, String keyMetadata) {
+    CompletableFuture<Void> future = new CompletableFuture<>();
+    // Create the request once instead of repeatedly for each server
+    byte[] keyMetadataHash = ManagedKeyData.constructMetadataHash(keyMetadata);
+    ManagedKeyEntryRequest request = ManagedKeyEntryRequest.newBuilder()
+      .setKeyCustNs(ManagedKeyRequest.newBuilder().setKeyCust(ByteString.copyFrom(keyCustodian))
+        .setKeyNamespace(keyNamespace).build())
+      .setKeyMetadataHash(ByteString.copyFrom(keyMetadataHash)).build();
+    List<CompletableFuture<Void>> futures =
+      regionServers.stream().map(serverName -> ejectManagedKeyDataCacheEntry(serverName, request))
+        .collect(Collectors.toList());
+    addListener(CompletableFuture.allOf(futures.toArray(new CompletableFuture<?>[0])),
+      (result, err) -> {
+        if (err != null) {
+          future.completeExceptionally(err);
+        } else {
+          future.complete(result);
+        }
+      });
+    return future;
+  }
+
+  private CompletableFuture<Void> ejectManagedKeyDataCacheEntry(ServerName serverName,
+    ManagedKeyEntryRequest request) {
+    return this.<Void> newAdminCaller()
+      .action((controller, stub) -> this.<ManagedKeyEntryRequest, HBaseProtos.BooleanMsg,
+        Void> adminCall(controller, stub, request,
+          (s, c, req, done) -> s.ejectManagedKeyDataCacheEntry(controller, req, done),
+          resp -> null))
+      .serverName(serverName).call();
+  }
+
+  @Override
+  public CompletableFuture<Void> clearManagedKeyDataCacheOnServers(List<ServerName> regionServers) {
+    CompletableFuture<Void> future = new CompletableFuture<>();
+    List<CompletableFuture<Void>> futures =
+      regionServers.stream().map(this::clearManagedKeyDataCache).collect(Collectors.toList());
+    addListener(CompletableFuture.allOf(futures.toArray(new CompletableFuture<?>[0])),
+      (result, err) -> {
+        if (err != null) {
+          future.completeExceptionally(err);
+        } else {
+          future.complete(result);
+        }
+      });
+    return future;
+  }
+
+  private CompletableFuture<Void> clearManagedKeyDataCache(ServerName serverName) {
+    return this.<Void> newAdminCaller()
+      .action((controller, stub) -> this.<EmptyMsg, EmptyMsg, Void> adminCall(controller, stub,
+        EmptyMsg.getDefaultInstance(),
+        (s, c, req, done) -> s.clearManagedKeyDataCache(controller, req, done), resp -> null))
+      .serverName(serverName).call();
+  }
 }

hbase-client/src/main/java/org/apache/hadoop/hbase/keymeta/KeymetaAdminClient.java

@@ -21,19 +21,23 @@
 import java.security.KeyException;
 import java.util.ArrayList;
 import java.util.List;
+import org.apache.commons.lang3.NotImplementedException;
 import org.apache.hadoop.hbase.client.Connection;
 import org.apache.hadoop.hbase.io.crypto.ManagedKeyData;
 import org.apache.hadoop.hbase.io.crypto.ManagedKeyState;
-import org.apache.hadoop.hbase.protobuf.generated.ManagedKeysProtos;
-import org.apache.hadoop.hbase.protobuf.generated.ManagedKeysProtos.ManagedKeyRequest;
-import org.apache.hadoop.hbase.protobuf.generated.ManagedKeysProtos.ManagedKeyResponse;
 import org.apache.yetus.audience.InterfaceAudience;
 
 import org.apache.hbase.thirdparty.com.google.protobuf.ByteString;
 import org.apache.hbase.thirdparty.com.google.protobuf.ServiceException;
 
 import org.apache.hadoop.hbase.shaded.protobuf.ProtobufUtil;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.BooleanMsg;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.EmptyMsg;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.GetManagedKeysResponse;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.ManagedKeyEntryRequest;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.ManagedKeyRequest;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.ManagedKeyResponse;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.ManagedKeysProtos;
 
 @InterfaceAudience.Public
 public class KeymetaAdminClient implements KeymetaAdmin {
@@ -48,9 +52,8 @@ public KeymetaAdminClient(Connection conn) throws IOException {
   public ManagedKeyData enableKeyManagement(byte[] keyCust, String keyNamespace)
     throws IOException {
     try {
-      ManagedKeysProtos.ManagedKeyResponse response =
-        stub.enableKeyManagement(null, ManagedKeyRequest.newBuilder()
-          .setKeyCust(ByteString.copyFrom(keyCust)).setKeyNamespace(keyNamespace).build());
+      ManagedKeyResponse response = stub.enableKeyManagement(null, ManagedKeyRequest.newBuilder()
+        .setKeyCust(ByteString.copyFrom(keyCust)).setKeyNamespace(keyNamespace).build());
       return generateKeyData(response);
     } catch (ServiceException e) {
       throw ProtobufUtil.handleRemoteException(e);
@@ -61,7 +64,7 @@ public ManagedKeyData enableKeyManagement(byte[] keyCust, String keyNamespace)
   public List<ManagedKeyData> getManagedKeys(byte[] keyCust, String keyNamespace)
     throws IOException, KeyException {
     try {
-      ManagedKeysProtos.GetManagedKeysResponse statusResponse =
+      GetManagedKeysResponse statusResponse =
         stub.getManagedKeys(null, ManagedKeyRequest.newBuilder()
           .setKeyCust(ByteString.copyFrom(keyCust)).setKeyNamespace(keyNamespace).build());
       return generateKeyDataList(statusResponse);
@@ -73,26 +76,95 @@ public List<ManagedKeyData> getManagedKeys(byte[] keyCust, String keyNamespace)
   @Override
   public boolean rotateSTK() throws IOException {
     try {
-      ManagedKeysProtos.RotateSTKResponse response =
-        stub.rotateSTK(null, EmptyMsg.getDefaultInstance());
-      return response.getRotated();
+      BooleanMsg response = stub.rotateSTK(null, EmptyMsg.getDefaultInstance());
+      return response.getBoolMsg();
+    } catch (ServiceException e) {
+      throw ProtobufUtil.handleRemoteException(e);
+    }
+  }
+
+  @Override
+  public void ejectManagedKeyDataCacheEntry(byte[] keyCustodian, String keyNamespace,
+    String keyMetadata) throws IOException {
+    throw new NotImplementedException(
+      "ejectManagedKeyDataCacheEntry not supported in KeymetaAdminClient");
+  }
+
+  @Override
+  public void clearManagedKeyDataCache() throws IOException {
+    throw new NotImplementedException(
+      "clearManagedKeyDataCache not supported in KeymetaAdminClient");
+  }
+
+  @Override
+  public ManagedKeyData disableKeyManagement(byte[] keyCust, String keyNamespace)
+    throws IOException, KeyException {
+    try {
+      ManagedKeyResponse response = stub.disableKeyManagement(null, ManagedKeyRequest.newBuilder()
+        .setKeyCust(ByteString.copyFrom(keyCust)).setKeyNamespace(keyNamespace).build());
+      return generateKeyData(response);
+    } catch (ServiceException e) {
+      throw ProtobufUtil.handleRemoteException(e);
+    }
+  }
+
+  @Override
+  public ManagedKeyData disableManagedKey(byte[] keyCust, String keyNamespace,
+    byte[] keyMetadataHash) throws IOException, KeyException {
+    try {
+      ManagedKeyResponse response = stub.disableManagedKey(null,
+        ManagedKeyEntryRequest.newBuilder()
+          .setKeyCustNs(ManagedKeyRequest.newBuilder().setKeyCust(ByteString.copyFrom(keyCust))
+            .setKeyNamespace(keyNamespace).build())
+          .setKeyMetadataHash(ByteString.copyFrom(keyMetadataHash)).build());
+      return generateKeyData(response);
     } catch (ServiceException e) {
       throw ProtobufUtil.handleRemoteException(e);
     }
   }
 
-  private static List<ManagedKeyData>
-    generateKeyDataList(ManagedKeysProtos.GetManagedKeysResponse stateResponse) {
+  @Override
+  public ManagedKeyData rotateManagedKey(byte[] keyCust, String keyNamespace)
+    throws IOException, KeyException {
+    try {
+      ManagedKeyResponse response = stub.rotateManagedKey(null, ManagedKeyRequest.newBuilder()
+        .setKeyCust(ByteString.copyFrom(keyCust)).setKeyNamespace(keyNamespace).build());
+      return generateKeyData(response);
+    } catch (ServiceException e) {
+      throw ProtobufUtil.handleRemoteException(e);
+    }
+  }
+
+  @Override
+  public void refreshManagedKeys(byte[] keyCust, String keyNamespace)
+    throws IOException, KeyException {
+    try {
+      stub.refreshManagedKeys(null, ManagedKeyRequest.newBuilder()
+        .setKeyCust(ByteString.copyFrom(keyCust)).setKeyNamespace(keyNamespace).build());
+    } catch (ServiceException e) {
+      throw ProtobufUtil.handleRemoteException(e);
+    }
+  }
+
+  private static List<ManagedKeyData> generateKeyDataList(GetManagedKeysResponse stateResponse) {
     List<ManagedKeyData> keyStates = new ArrayList<>();
     for (ManagedKeyResponse state : stateResponse.getStateList()) {
       keyStates.add(generateKeyData(state));
     }
     return keyStates;
   }
 
-  private static ManagedKeyData generateKeyData(ManagedKeysProtos.ManagedKeyResponse response) {
-    return new ManagedKeyData(response.getKeyCust().toByteArray(), response.getKeyNamespace(), null,
-      ManagedKeyState.forValue((byte) response.getKeyState().getNumber()),
-      response.getKeyMetadata(), response.getRefreshTimestamp());
+  private static ManagedKeyData generateKeyData(ManagedKeyResponse response) {
+    // Use hash-only constructor for client-side ManagedKeyData
+    byte[] keyMetadataHash =
+      response.hasKeyMetadataHash() ? response.getKeyMetadataHash().toByteArray() : null;
+    if (keyMetadataHash == null) {
+      return new ManagedKeyData(response.getKeyCust().toByteArray(), response.getKeyNamespace(),
+        ManagedKeyState.forValue((byte) response.getKeyState().getNumber()));
+    } else {
+      return new ManagedKeyData(response.getKeyCust().toByteArray(), response.getKeyNamespace(),
+        ManagedKeyState.forValue((byte) response.getKeyState().getNumber()), keyMetadataHash,
+        response.getRefreshTimestamp());
+    }
   }
 }