Skip to content

Commit 28e0de8

Browse files
author
Fabian Morgan
committed
pr review updates for Tejaskriya
1 parent 65b57e9 commit 28e0de8

1 file changed

Lines changed: 41 additions & 84 deletions

File tree

‎hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/security/TestS3SecurityUtil.java‎

Lines changed: 41 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import static org.mockito.Mockito.mockStatic;
3030
import static org.mockito.Mockito.when;
3131

32-
import java.io.IOException;
3332
import java.time.Clock;
3433
import java.time.Instant;
3534
import java.util.UUID;
@@ -60,63 +59,48 @@ public class TestS3SecurityUtil {
6059
public void testValidateS3CredentialFailsWhenTokenRevoked() throws Exception {
6160
// If the revoked STS token table contains an entry for the temporary access key id extracted from the session
6261
// token, validateS3Credential should reject the request with REVOKED_TOKEN
63-
final String sessionToken = "session-token-a";
64-
final String tempAccessKeyId = "ASIA123456789";
62+
validateS3CredentialHelper("session-token-a", true, true, REVOKED_TOKEN);
63+
}
64+
65+
@Test
66+
public void testValidateS3CredentialWhenMetadataUnavailable() throws Exception {
67+
// If the metadata manager is not available, the revocation check should not cause the request to be rejected.
68+
validateS3CredentialHelper("session-token-b", false, false, null);
69+
}
70+
71+
@Test
72+
public void testValidateS3CredentialSuccessWhenNotRevoked() throws Exception {
73+
// Normal case: token is NOT revoked and request is accepted
74+
validateS3CredentialHelper("session-token-c", true, false, null);
75+
}
76+
77+
private void validateS3CredentialHelper(String sessionToken, boolean metadataAvailable, boolean isRevoked,
78+
OMException.ResultCodes expectedResult) throws Exception {
6579

6680
try (OzoneManager ozoneManager = mock(OzoneManager.class)) {
6781
when(ozoneManager.isSecurityEnabled()).thenReturn(true);
6882
when(ozoneManager.getSecretKeyClient()).thenReturn(mock(SecretKeyClient.class));
6983

70-
final OMMetadataManager metadataManager = mock(OMMetadataManager.class);
71-
when(ozoneManager.getMetadataManager()).thenReturn(metadataManager);
72-
7384
final Table<String, String> revokedSTSTokenTable = new InMemoryTestTable<>();
74-
when(metadataManager.getS3RevokedStsTokenTable()).thenReturn(revokedSTSTokenTable);
75-
76-
// Mock STSSecurityUtil to return a token whose tempAccessKeyId matches the one that's revoked.
77-
final STSTokenIdentifier stsTokenIdentifier = new STSTokenIdentifier(
78-
tempAccessKeyId, "original-access-key-id", "arn:aws:iam::123456789012:role/test-role",
79-
Instant.now().plusSeconds(3600), "secret-access-key", "session-policy",
80-
ENCRYPTION_KEY);
81-
82-
try (MockedStatic<STSSecurityUtil> stsSecurityUtilMock = mockStatic(STSSecurityUtil.class, CALLS_REAL_METHODS)) {
83-
stsSecurityUtilMock.when(
84-
() -> STSSecurityUtil.constructValidateAndDecryptSTSToken(
85-
eq(sessionToken), any(SecretKeyClient.class), any(Clock.class)))
86-
.thenReturn(stsTokenIdentifier);
85+
if (metadataAvailable) {
86+
final OMMetadataManager metadataManager = mock(OMMetadataManager.class);
87+
when(ozoneManager.getMetadataManager()).thenReturn(metadataManager);
88+
when(metadataManager.getS3RevokedStsTokenTable()).thenReturn(revokedSTSTokenTable);
89+
} else {
90+
when(ozoneManager.getMetadataManager()).thenReturn(null);
91+
}
8792

88-
// Revoke the tempAccessKeyId
93+
final String tempAccessKeyId = "temp-access-key-id";
94+
if (isRevoked) {
8995
revokedSTSTokenTable.put(tempAccessKeyId, sessionToken);
90-
91-
final OMRequest omRequest = createRequestWithSessionToken(sessionToken);
92-
final OMException ex = assertThrows(
93-
OMException.class, () -> S3SecurityUtil.validateS3Credential(omRequest, ozoneManager));
94-
assertEquals(REVOKED_TOKEN, ex.getResult());
9596
}
96-
}
97-
}
98-
99-
@Test
100-
public void testValidateS3CredentialWhenMetadataUnavailable() {
101-
// If the metadata manager is not available, the revocation check should not cause the request to be rejected.
102-
final String sessionToken = "session-token-b";
103-
104-
try (OzoneManager ozoneManager = mock(OzoneManager.class)) {
105-
when(ozoneManager.isSecurityEnabled()).thenReturn(true);
106-
when(ozoneManager.getMetadataManager()).thenReturn(null);
107-
when(ozoneManager.getSecretKeyClient()).thenReturn(mock(SecretKeyClient.class));
10897

109-
final OMRequest omRequest = createRequestWithSessionToken(sessionToken);
98+
final STSTokenIdentifier stsTokenIdentifier = createSTSTokenIdentifier();
11099

111100
try (MockedStatic<STSSecurityUtil> stsSecurityUtilMock = mockStatic(STSSecurityUtil.class, CALLS_REAL_METHODS);
112101
MockedStatic<AWSV4AuthValidator> awsV4AuthValidatorMock = mockStatic(
113102
AWSV4AuthValidator.class, CALLS_REAL_METHODS)) {
114103

115-
final STSTokenIdentifier stsTokenIdentifier = new STSTokenIdentifier(
116-
"temp-access-key-id", "original-access-key-id", "arn:aws:iam::123456789012:role/test-role",
117-
Instant.now().plusSeconds(3600), "secret-access-key", "session-policy",
118-
ENCRYPTION_KEY);
119-
120104
stsSecurityUtilMock.when(
121105
() -> STSSecurityUtil.constructValidateAndDecryptSTSToken(
122106
eq(sessionToken), any(SecretKeyClient.class), any(Clock.class)))
@@ -126,51 +110,24 @@ public void testValidateS3CredentialWhenMetadataUnavailable() {
126110
awsV4AuthValidatorMock.when(() -> AWSV4AuthValidator.validateRequest(anyString(), anyString(), anyString()))
127111
.thenReturn(true);
128112

129-
assertDoesNotThrow(() -> S3SecurityUtil.validateS3Credential(omRequest, ozoneManager));
113+
final OMRequest omRequest = createRequestWithSessionToken(sessionToken);
114+
115+
if (expectedResult != null) {
116+
final OMException ex = assertThrows(OMException.class,
117+
() -> S3SecurityUtil.validateS3Credential(omRequest, ozoneManager));
118+
assertEquals(expectedResult, ex.getResult());
119+
} else {
120+
assertDoesNotThrow(() -> S3SecurityUtil.validateS3Credential(omRequest, ozoneManager));
121+
}
130122
}
131-
} catch (IOException e) {
132-
throw new RuntimeException(e);
133123
}
134124
}
135125

136-
@Test
137-
public void testValidateS3CredentialSuccessWhenNotRevoked() {
138-
// Normal case: token is NOT revoked and request is accepted
139-
final String sessionToken = "session-token-c";
140-
141-
try (OzoneManager ozoneManager = mock(OzoneManager.class)) {
142-
when(ozoneManager.isSecurityEnabled()).thenReturn(true);
143-
when(ozoneManager.getSecretKeyClient()).thenReturn(mock(SecretKeyClient.class));
144-
145-
final OMMetadataManager metadataManager = mock(OMMetadataManager.class);
146-
when(ozoneManager.getMetadataManager()).thenReturn(metadataManager);
147-
148-
final Table<String, String> revokedSTSTokenTable = new InMemoryTestTable<>();
149-
when(metadataManager.getS3RevokedStsTokenTable()).thenReturn(revokedSTSTokenTable);
150-
151-
// Not revoked -> getIfExist returns null by default in InMemoryTestTable
152-
final OMRequest omRequest = createRequestWithSessionToken(sessionToken);
153-
final STSTokenIdentifier stsTokenIdentifier = new STSTokenIdentifier(
154-
"temp-access-key-id", "original-access-key-id", "arn:aws:iam::123456789012:role/test-role",
155-
Instant.now().plusSeconds(3600), "secret-access-key", "session-policy",
156-
ENCRYPTION_KEY);
157-
158-
try (MockedStatic<STSSecurityUtil> stsSecurityUtilMock = mockStatic(STSSecurityUtil.class, CALLS_REAL_METHODS);
159-
MockedStatic<AWSV4AuthValidator> awsV4AuthValidatorMock = mockStatic(
160-
AWSV4AuthValidator.class, CALLS_REAL_METHODS)) {
161-
162-
stsSecurityUtilMock.when(
163-
() -> STSSecurityUtil.constructValidateAndDecryptSTSToken(
164-
eq(sessionToken), any(SecretKeyClient.class), any(Clock.class)))
165-
.thenReturn(stsTokenIdentifier);
166-
awsV4AuthValidatorMock.when(() -> AWSV4AuthValidator.validateRequest(anyString(), anyString(), anyString()))
167-
.thenReturn(true);
168-
169-
assertDoesNotThrow(() -> S3SecurityUtil.validateS3Credential(omRequest, ozoneManager));
170-
}
171-
} catch (IOException e) {
172-
throw new RuntimeException(e);
173-
}
126+
private STSTokenIdentifier createSTSTokenIdentifier() {
127+
return new STSTokenIdentifier(
128+
"temp-access-key-id", "original-access-key-id", "arn:aws:iam::123456789012:role/test-role",
129+
Instant.now().plusSeconds(3600), "secret-access-key", "session-policy",
130+
ENCRYPTION_KEY);
174131
}
175132

176133
private static OMRequest createRequestWithSessionToken(String sessionToken) {

0 commit comments

Comments
 (0)