HDDS-14373. [STS] Revoked STS token logic tweaks (#9604)

Author: fmorg-git

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/s3/security/S3RevokeSTSTokenRequest.java

@@ -22,6 +22,8 @@
 import java.time.ZoneOffset;
 import java.util.HashMap;
 import java.util.Map;
+import org.apache.hadoop.hdds.utils.db.cache.CacheKey;
+import org.apache.hadoop.hdds.utils.db.cache.CacheValue;
 import org.apache.hadoop.ozone.OzoneConsts;
 import org.apache.hadoop.ozone.audit.OMAction;
 import org.apache.hadoop.ozone.om.OzoneManager;
@@ -106,6 +108,10 @@ public OMClientResponse validateAndUpdateCache(OzoneManager ozoneManager, Execut
     markForAudit(ozoneManager.getAuditLogger(), buildAuditMessage(
         OMAction.REVOKE_STS_TOKEN, auditMap, null, userInfo));
 
+    // Update the cache immediately so subsequent validation checks see the revocation
+    ozoneManager.getMetadataManager().getS3RevokedStsTokenTable().addCacheEntry(
+        new CacheKey<>(sessionToken), CacheValue.get(context.getIndex(), CLOCK.millis()));
+
     LOG.info("Marked STS session token '{}' as revoked.", sessionToken);
     return omClientResponse;
   }

hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/s3/security/TestS3RevokeSTSTokenRequest.java

@@ -21,19 +21,28 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import java.io.IOException;
 import java.util.Optional;
 import java.util.UUID;
 import org.apache.hadoop.hdds.security.symmetric.SecretKeyClient;
+import org.apache.hadoop.hdds.utils.db.Table;
+import org.apache.hadoop.hdds.utils.db.cache.CacheKey;
+import org.apache.hadoop.hdds.utils.db.cache.CacheValue;
 import org.apache.hadoop.ipc.ExternalCall;
 import org.apache.hadoop.ipc.Server;
+import org.apache.hadoop.ozone.audit.AuditLogger;
+import org.apache.hadoop.ozone.om.OMMetadataManager;
 import org.apache.hadoop.ozone.om.OMMultiTenantManager;
 import org.apache.hadoop.ozone.om.OzoneManager;
 import org.apache.hadoop.ozone.om.exceptions.OMException;
+import org.apache.hadoop.ozone.om.execution.flowcontrol.ExecutionContext;
 import org.apache.hadoop.ozone.om.request.OMClientRequest;
+import org.apache.hadoop.ozone.om.response.OMClientResponse;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Type;
@@ -283,6 +292,41 @@ public void testPreExecuteFailsForNonOwnerNonAdminInTenant() throws Exception {
     assertEquals(OMException.ResultCodes.USER_MISMATCH, ex.getResult());
   }
 
+  @Test
+  public void testValidateAndUpdateCacheUpdatesCacheImmediately() throws Exception {
+    final String tempAccessKeyId = "ASIA4567891230";
+    final String originalAccessKeyId = "original-access-key-id";
+    final String sessionToken = createSessionToken(tempAccessKeyId, originalAccessKeyId);
+
+    final OzoneManager ozoneManager = mock(OzoneManager.class);
+    final OMMetadataManager omMetadataManager = mock(OMMetadataManager.class);
+    @SuppressWarnings("unchecked")
+    final Table<String, Long> s3RevokedStsTokenTable = mock(Table.class);
+    final ExecutionContext context = mock(ExecutionContext.class);
+    final AuditLogger auditLogger = mock(AuditLogger.class);
+
+    when(ozoneManager.getMetadataManager()).thenReturn(omMetadataManager);
+    when(omMetadataManager.getS3RevokedStsTokenTable()).thenReturn(s3RevokedStsTokenTable);
+    when(ozoneManager.getAuditLogger()).thenReturn(auditLogger);
+
+    final OzoneManagerProtocolProtos.RevokeSTSTokenRequest revokeRequest =
+        OzoneManagerProtocolProtos.RevokeSTSTokenRequest.newBuilder()
+            .setSessionToken(sessionToken)
+            .build();
+
+    final OMRequest omRequest = OMRequest.newBuilder()
+        .setClientId(UUID.randomUUID().toString())
+        .setCmdType(Type.RevokeSTSToken)
+        .setRevokeSTSTokenRequest(revokeRequest)
+        .build();
+
+    final S3RevokeSTSTokenRequest s3RevokeSTSTokenRequest = new S3RevokeSTSTokenRequest(omRequest);
+    final OMClientResponse omClientResponse = s3RevokeSTSTokenRequest.validateAndUpdateCache(ozoneManager, context);
+
+    assertEquals(OzoneManagerProtocolProtos.Status.OK, omClientResponse.getOMResponse().getStatus());
+    verify(s3RevokedStsTokenTable).addCacheEntry(eq(new CacheKey<>(sessionToken)), any(CacheValue.class));
+  }
+
   /**
    * Stub used to inject a remote user into the ProtobufRpcEngine.Server.getRemoteUser() thread-local.
    */

hadoop-ozone/s3gateway/src/main/java/org/apache/hadoop/ozone/s3/endpoint/EndpointBase.java

@@ -537,7 +537,8 @@ protected void auditReadFailure(AuditAction action, Exception ex) {
   protected boolean isAccessDenied(OMException ex) {
     ResultCodes result = ex.getResult();
     return result == ResultCodes.PERMISSION_DENIED
-        || result == ResultCodes.INVALID_TOKEN;
+        || result == ResultCodes.INVALID_TOKEN
+        || result == ResultCodes.REVOKED_TOKEN;
   }
 
 }

hadoop-ozone/s3gateway/src/test/java/org/apache/hadoop/ozone/s3/endpoint/TestEndpointBase.java

@@ -17,17 +17,21 @@
 
 package org.apache.hadoop.ozone.s3.endpoint;
 
+import static org.apache.hadoop.ozone.om.exceptions.OMException.ResultCodes;
 import static org.apache.hadoop.ozone.s3.util.S3Consts.CUSTOM_METADATA_HEADER_PREFIX;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.nio.charset.StandardCharsets;
 import java.util.Locale;
 import java.util.Map;
 import javax.ws.rs.core.MultivaluedHashMap;
 import javax.ws.rs.core.MultivaluedMap;
 import org.apache.hadoop.ozone.OzoneConsts;
+import org.apache.hadoop.ozone.om.exceptions.OMException;
 import org.apache.hadoop.ozone.s3.exception.OS3Exception;
 import org.junit.jupiter.api.Test;
 
@@ -114,4 +118,18 @@ public void init() { }
     assertEquals(value, customMetadata.get(key));
   }
 
+  @Test
+  public void testAccessDeniedResultCodes() {
+    final EndpointBase endpointBase = new EndpointBase() {
+      @Override
+      public void init() { }
+    };
+
+    assertTrue(endpointBase.isAccessDenied(new OMException(ResultCodes.PERMISSION_DENIED)));
+    assertTrue(endpointBase.isAccessDenied(new OMException(ResultCodes.INVALID_TOKEN)));
+    assertTrue(endpointBase.isAccessDenied(new OMException(ResultCodes.REVOKED_TOKEN)));
+    assertFalse(endpointBase.isAccessDenied(new OMException(ResultCodes.INTERNAL_ERROR)));
+    assertFalse(endpointBase.isAccessDenied(new OMException(ResultCodes.BUCKET_NOT_FOUND)));
+  }
+
 }