This is an automated email from the ASF dual-hosted git repository.
xiangfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 0f93d52ef0b [audit] Add SPI-based token resolver for custom audit
identity resolution (#17658)
0f93d52ef0b is described below
commit 0f93d52ef0bfa1c2dc367b31f137a63bd25f9d2f
Author: Suvodeep Pyne <[email protected]>
AuthorDate: Tue Feb 10 00:17:37 2026 -0800
[audit] Add SPI-based token resolver for custom audit identity resolution
(#17658)
* [audit] Add support for custom token resolvers in audit identity
resolution
- Introduced `AuditTokenResolver` SPI for handling proprietary token
formats.
- Updated `AuditIdentityResolver` to integrate resolver-based identity
resolution.
- Enhanced `AuditConfig` to support token resolver class configuration.
- Added unit tests and mock resolver for validation.
* [audit] Refactor audit identity resolution to use structured
`AuditUserIdentity`
- Updated `AuditTokenResolver` to return `AuditUserIdentity` instead of
plain principal strings.
- Introduced `AuditUserIdentity` interface for extensible identity
representations.
- Refactored related classes and tests to use structured identity.
- Enhanced `AuditIdentityResolver` to handle the updated resolver contract.
* [audit] Simplify `MockAuditTokenResolver` implementation and enhance
error handling in `AuditIdentityResolver`
- Removed unnecessary `MockConfig` wrapper and replaced `AtomicReference`
with direct static fields.
- Streamlined logic in `MockAuditTokenResolver` for clearer behavior.
- Improved resilience in `AuditIdentityResolver` to handle null resolver
class gracefully.
* [audit] Add no-arg constructor to MockAuditTokenResolver for PluginManager
PluginManager.createInstance() requires a no-arg constructor to
instantiate classes. The missing constructor caused
testResolverLoadedViaPluginManager to fail in CI.
---
.../org/apache/pinot/common/audit/AuditConfig.java | 12 ++
.../org/apache/pinot/common/audit/AuditEvent.java | 3 +-
.../pinot/common/audit/AuditIdentityResolver.java | 105 ++++++++++++++-
.../common/audit/AuditIdentityResolverTest.java | 145 ++++++++++-----------
.../pinot/common/audit/MockAuditTokenResolver.java | 93 +++++++++++++
.../apache/pinot/spi/audit/AuditTokenResolver.java | 51 ++++++++
.../apache/pinot/spi/audit/AuditUserIdentity.java | 41 ++++++
7 files changed, 368 insertions(+), 82 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditConfig.java
b/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditConfig.java
index 9e542cf635c..e8ce1013202 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditConfig.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditConfig.java
@@ -60,6 +60,9 @@ public final class AuditConfig {
@JsonProperty("capture.response.enabled")
private boolean _captureResponseEnabled = false;
+ @JsonProperty("token.resolver.class")
+ private String _tokenResolverClass = "";
+
public boolean isEnabled() {
return _enabled;
}
@@ -132,6 +135,14 @@ public final class AuditConfig {
_captureResponseEnabled = captureResponseEnabled;
}
+ public String getTokenResolverClass() {
+ return _tokenResolverClass;
+ }
+
+ public void setTokenResolverClass(String tokenResolverClass) {
+ _tokenResolverClass = tokenResolverClass;
+ }
+
@Override
public String toString() {
return new StringJoiner(", ", AuditConfig.class.getSimpleName() + "[",
"]").add("_enabled=" + _enabled)
@@ -143,6 +154,7 @@ public final class AuditConfig {
.add("_useridHeader='" + _useridHeader + "'")
.add("_useridJwtClaimName='" + _useridJwtClaimName + "'")
.add("_captureResponseEnabled=" + _captureResponseEnabled)
+ .add("_tokenResolverClass='" + _tokenResolverClass + "'")
.toString();
}
}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditEvent.java
b/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditEvent.java
index 8f370eeb3db..7a2c2bd760d 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditEvent.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditEvent.java
@@ -22,6 +22,7 @@ import com.fasterxml.jackson.annotation.JsonAutoDetect;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import java.util.Map;
+import org.apache.pinot.spi.audit.AuditUserIdentity;
/**
@@ -209,7 +210,7 @@ public class AuditEvent {
}
@JsonInclude(JsonInclude.Include.NON_NULL)
- public static class UserIdentity {
+ public static class UserIdentity implements AuditUserIdentity {
@JsonProperty("principal")
private String _principal;
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditIdentityResolver.java
b/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditIdentityResolver.java
index 19fe2cd85fe..1522cc56214 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditIdentityResolver.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/audit/AuditIdentityResolver.java
@@ -18,14 +18,20 @@
*/
package org.apache.pinot.common.audit;
+import com.google.common.annotations.VisibleForTesting;
import com.nimbusds.jwt.JWT;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.JWTParser;
+import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nullable;
import javax.inject.Inject;
import javax.inject.Singleton;
import javax.ws.rs.container.ContainerRequestContext;
import javax.ws.rs.core.HttpHeaders;
import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.audit.AuditTokenResolver;
+import org.apache.pinot.spi.audit.AuditUserIdentity;
+import org.apache.pinot.spi.plugin.PluginManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -36,6 +42,7 @@ import org.slf4j.LoggerFactory;
* This resolver supports multiple identity resolution strategies in order of
priority:
* <ol>
* <li>Custom identity header - as configured in the audit configuration</li>
+ * <li>Custom token resolver (via SPI) - for proprietary token formats</li>
* <li>JWT token in Authorization header - extracting principal from JWT
claims</li>
* </ol>
* <p>
@@ -51,18 +58,26 @@ public class AuditIdentityResolver {
private static final String BEARER_PREFIX = "Bearer ";
private final AuditConfigManager _configManager;
+ private final AtomicReference<ResolverHolder> _resolverHolder = new
AtomicReference<>(new ResolverHolder());
@Inject
public AuditIdentityResolver(AuditConfigManager configManager) {
_configManager = configManager;
}
+ @VisibleForTesting
+ AuditIdentityResolver(AuditConfigManager configManager, @Nullable
AuditTokenResolver tokenResolver) {
+ _configManager = configManager;
+ _resolverHolder.set(new ResolverHolder(tokenResolver));
+ }
+
/**
* Resolves user identity from the given HTTP request context.
* <p>
* The resolution follows a priority order:
* <ol>
* <li>Check for a custom identity header as specified in the audit
configuration</li>
+ * <li>Use custom token resolver (if configured) to resolve from
Authorization header</li>
* <li>Extract principal from JWT token in the Authorization header</li>
* </ol>
* <p>
@@ -73,6 +88,7 @@ public class AuditIdentityResolver {
* @return a {@link AuditEvent.UserIdentity} containing the resolved
principal, or {@code null} if no identity
* could be resolved
*/
+ @Nullable
public AuditEvent.UserIdentity resolveIdentity(ContainerRequestContext
requestContext) {
AuditConfig config = _configManager.getCurrentConfig();
@@ -85,9 +101,23 @@ public class AuditIdentityResolver {
}
}
- // Priority 2: Check JWT in Authorization header
+ // Get Authorization header for subsequent checks
String authHeader =
requestContext.getHeaderString(HttpHeaders.AUTHORIZATION);
- if (StringUtils.isNotBlank(authHeader) &&
authHeader.startsWith(BEARER_PREFIX)) {
+ if (StringUtils.isBlank(authHeader)) {
+ return null;
+ }
+
+ // Priority 2: Try custom token resolver
+ AuditTokenResolver resolver = getTokenResolver(config);
+ if (resolver != null) {
+ AuditUserIdentity identity = resolver.resolve(authHeader);
+ if (identity != null && StringUtils.isNotBlank(identity.getPrincipal()))
{
+ return new
AuditEvent.UserIdentity().setPrincipal(identity.getPrincipal());
+ }
+ }
+
+ // Priority 3: Fallback to JWT parsing
+ if (authHeader.startsWith(BEARER_PREFIX)) {
String token = authHeader.substring(BEARER_PREFIX.length()).trim();
String principal = extractJwtPrincipal(token,
config.getUseridJwtClaimName());
if (StringUtils.isNotBlank(principal)) {
@@ -95,10 +125,45 @@ public class AuditIdentityResolver {
}
}
- // Return null instead of anonymous
return null;
}
+ @Nullable
+ private AuditTokenResolver getTokenResolver(AuditConfig config) {
+ String resolverClass = config.getTokenResolverClass();
+ ResolverHolder currentHolder = _resolverHolder.get();
+
+ // If no resolver class configured or already loaded, return current
resolver
+ if (StringUtils.isBlank(resolverClass) ||
currentHolder.isLoaded(resolverClass)) {
+ return currentHolder.getResolver();
+ }
+
+ // Need to load new resolver - use synchronized to prevent concurrent
loading
+ synchronized (this) {
+ currentHolder = _resolverHolder.get();
+ if (currentHolder.isLoaded(resolverClass)) {
+ return currentHolder.getResolver();
+ }
+
+ AuditTokenResolver newResolver = loadTokenResolver(resolverClass);
+ // Don't cache the instance if it does not exist. Occasionally, we can
run into loading failures.
+ _resolverHolder.set(new ResolverHolder(newResolver, newResolver != null
? resolverClass : null));
+ return newResolver;
+ }
+ }
+
+ @Nullable
+ private AuditTokenResolver loadTokenResolver(String className) {
+ try {
+ AuditTokenResolver resolver =
PluginManager.get().createInstance(className);
+ LOG.info("Successfully loaded AuditTokenResolver: {}", className);
+ return resolver;
+ } catch (Exception e) {
+ LOG.error("Failed to load AuditTokenResolver: {}", className, e);
+ return null;
+ }
+ }
+
private String extractJwtPrincipal(String token, String claimName) {
try {
JWT jwt = JWTParser.parse(token);
@@ -119,4 +184,38 @@ public class AuditIdentityResolver {
return null;
}
}
+
+ /**
+ * Immutable holder for resolver and its class name to enable atomic updates.
+ */
+ private static final class ResolverHolder {
+ @Nullable
+ private final AuditTokenResolver _resolver;
+ @Nullable
+ private final String _className;
+
+ ResolverHolder() {
+ _resolver = null;
+ _className = null;
+ }
+
+ ResolverHolder(@Nullable AuditTokenResolver resolver) {
+ _resolver = resolver;
+ _className = resolver != null ? resolver.getClass().getName() : null;
+ }
+
+ ResolverHolder(@Nullable AuditTokenResolver resolver, @Nullable String
className) {
+ _resolver = resolver;
+ _className = className;
+ }
+
+ @Nullable
+ AuditTokenResolver getResolver() {
+ return _resolver;
+ }
+
+ boolean isLoaded(@Nullable String className) {
+ return className != null && className.equals(_className);
+ }
+ }
}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/audit/AuditIdentityResolverTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/audit/AuditIdentityResolverTest.java
index d33bc9cf650..04e12eeb546 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/audit/AuditIdentityResolverTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/audit/AuditIdentityResolverTest.java
@@ -29,6 +29,7 @@ import java.util.HashMap;
import java.util.Map;
import javax.ws.rs.container.ContainerRequestContext;
import javax.ws.rs.core.HttpHeaders;
+import org.apache.pinot.spi.audit.AuditTokenResolver;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.testng.annotations.BeforeMethod;
@@ -53,14 +54,17 @@ public class AuditIdentityResolverTest {
public void setUp()
throws Exception {
MockitoAnnotations.openMocks(this);
+ MockAuditTokenResolver.reset();
_auditIdentityResolver = new AuditIdentityResolver(_auditConfigManager);
_auditConfig = new AuditConfig();
when(_auditConfigManager.getCurrentConfig()).thenReturn(_auditConfig);
}
+ // ==================== Custom Header Tests ====================
+
@Test
- public void testResolveIdentityFromCustomHeaderSuccess() {
+ public void testResolveIdentityFromCustomHeader() {
_auditConfig.setUseridHeader("X-User-Email");
when(_requestContext.getHeaderString("X-User-Email")).thenReturn("[email protected]");
@@ -70,13 +74,12 @@ public class AuditIdentityResolverTest {
assertThat(result.getPrincipal()).isEqualTo("[email protected]");
}
+ // ==================== JWT Tests ====================
+
@Test
- public void testResolveIdentityFromCustomHeaderEmptyHeaderValue()
+ public void testResolveIdentityFromJwtCustomClaim()
throws Exception {
- _auditConfig.setUseridHeader("X-User-Email");
_auditConfig.setUseridJwtClaimName("email");
-
- when(_requestContext.getHeaderString("X-User-Email")).thenReturn("");
String validJwt = createJwtToken("user123", Map.of("email",
"[email protected]"));
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
@@ -87,157 +90,143 @@ public class AuditIdentityResolverTest {
}
@Test
- public void testResolveIdentityFromCustomHeaderMissingHeader()
+ public void testResolveIdentityFromJwtSubjectWhenClaimMissing()
throws Exception {
- _auditConfig.setUseridHeader("X-User-Email");
_auditConfig.setUseridJwtClaimName("email");
-
- when(_requestContext.getHeaderString("X-User-Email")).thenReturn(null);
- String validJwt = createJwtToken("user123", Map.of("email",
"[email protected]"));
+ String validJwt = createJwtToken("user123", new HashMap<>());
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("[email protected]");
+ assertThat(result.getPrincipal()).isEqualTo("user123");
}
@Test
- public void testResolveIdentityFromJwtCustomClaimSuccess()
- throws Exception {
- _auditConfig.setUseridJwtClaimName("email");
-
- String validJwt = createJwtToken("user123", Map.of("email",
"[email protected]"));
-
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
+ public void testInvalidJwtTokenReturnsNull() {
+
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
invalid-token");
AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
- assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("[email protected]");
+ assertThat(result).isNull();
}
+ // ==================== Custom Resolver Tests ====================
+
@Test
- public void testResolveIdentityFromJwtSubjectWhenCustomClaimNotConfigured()
- throws Exception {
- String validJwt = createJwtToken("user123", new HashMap<>());
-
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
+ public void testCustomTokenResolverSuccess() {
+ AuditTokenResolver mockResolver = new
MockAuditTokenResolver("resolved-user");
+ AuditIdentityResolver resolver = new
AuditIdentityResolver(_auditConfigManager, mockResolver);
+
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
custom-token");
- AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
+ AuditEvent.UserIdentity result = resolver.resolveIdentity(_requestContext);
assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("user123");
+ assertThat(result.getPrincipal()).isEqualTo("resolved-user");
}
@Test
- public void testResolveIdentityFromJwtSubjectWhenCustomClaimMissing()
+ public void testCustomResolverReturnsNullFallsBackToJwt()
throws Exception {
- _auditConfig.setUseridJwtClaimName("email");
-
- String validJwt = createJwtToken("user123", new HashMap<>());
+ AuditTokenResolver mockResolver = new MockAuditTokenResolver(null);
+ AuditIdentityResolver resolver = new
AuditIdentityResolver(_auditConfigManager, mockResolver);
+ String validJwt = createJwtToken("jwt-user", new HashMap<>());
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
- AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
+ AuditEvent.UserIdentity result = resolver.resolveIdentity(_requestContext);
assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("user123");
+ assertThat(result.getPrincipal()).isEqualTo("jwt-user");
}
@Test
- public void testInvalidJwtToken() {
-
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
invalid-token");
+ public void testCustomResolverReceivesFullAuthHeader() {
+ String expectedAuthHeader = "Bearer custom-token-value";
+ MockAuditTokenResolver mockResolver = new MockAuditTokenResolver("user");
+ AuditIdentityResolver resolver = new
AuditIdentityResolver(_auditConfigManager, mockResolver);
+
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn(expectedAuthHeader);
- AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
+ resolver.resolveIdentity(_requestContext);
- assertThat(result).isNull();
+
assertThat(MockAuditTokenResolver.getLastAuthHeader()).isEqualTo(expectedAuthHeader);
}
- @Test
- public void testHeaderTakesPriorityOverJwt()
- throws Exception {
- _auditConfig.setUseridHeader("X-User-Email");
- _auditConfig.setUseridJwtClaimName("email");
+ // ==================== PluginManager Tests ====================
-
when(_requestContext.getHeaderString("X-User-Email")).thenReturn("[email protected]");
- String validJwt = createJwtToken("user123", Map.of("email",
"[email protected]"));
-
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
+ @Test
+ public void testResolverLoadedViaPluginManager() {
+ _auditConfig.setTokenResolverClass(MockAuditTokenResolver.class.getName());
+
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
test-token");
AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("[email protected]");
+ assertThat(result.getPrincipal()).isEqualTo("mock-resolved-user");
}
@Test
- public void testFallbackChainNoHeaderFallsBackToJwtClaim()
+ public void testInvalidResolverClassFallsBackToJwt()
throws Exception {
- _auditConfig.setUseridHeader("X-User-Email");
- _auditConfig.setUseridJwtClaimName("email");
-
- when(_requestContext.getHeaderString("X-User-Email")).thenReturn(null);
- String validJwt = createJwtToken("user123", Map.of("email",
"[email protected]"));
+ _auditConfig.setTokenResolverClass("com.invalid.NonExistentResolver");
+ String validJwt = createJwtToken("jwt-user", new HashMap<>());
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("[email protected]");
+ assertThat(result.getPrincipal()).isEqualTo("jwt-user");
}
+ // ==================== Priority Tests ====================
+
@Test
- public void testFallbackChainNoHeaderNoClaimFallsBackToSubject()
+ public void testPriorityHeaderOverJwt()
throws Exception {
_auditConfig.setUseridHeader("X-User-Email");
- _auditConfig.setUseridJwtClaimName("email");
-
- when(_requestContext.getHeaderString("X-User-Email")).thenReturn(null);
- String validJwt = createJwtToken("user123", new HashMap<>());
+
when(_requestContext.getHeaderString("X-User-Email")).thenReturn("header-user");
+ String validJwt = createJwtToken("jwt-user", new HashMap<>());
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("user123");
+ assertThat(result.getPrincipal()).isEqualTo("header-user");
}
@Test
- public void testNoAuthenticationPresentReturnsNull() {
- AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
-
- assertThat(result).isNull();
- }
-
- @Test
- public void testBlankConfigurationUsesDefaults()
- throws Exception {
- String validJwt = createJwtToken("user123", new HashMap<>());
-
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
" + validJwt);
+ public void testPriorityHeaderOverResolver() {
+ AuditTokenResolver mockResolver = new
MockAuditTokenResolver("resolver-user");
+ AuditIdentityResolver resolver = new
AuditIdentityResolver(_auditConfigManager, mockResolver);
+ _auditConfig.setUseridHeader("X-User-Email");
+
when(_requestContext.getHeaderString("X-User-Email")).thenReturn("header-user");
+
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Bearer
token");
- AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
+ AuditEvent.UserIdentity result = resolver.resolveIdentity(_requestContext);
assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo("user123");
+ assertThat(result.getPrincipal()).isEqualTo("header-user");
}
- @Test
- public void testNonBearerAuthorizationHeaderReturnsNull() {
-
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Basic
dXNlcjpwYXNz");
+ // ==================== Edge Cases ====================
+ @Test
+ public void testNoAuthenticationReturnsNull() {
AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
assertThat(result).isNull();
}
@Test
- public void testWhitespaceInHeaderValueTrimmed() {
- _auditConfig.setUseridHeader("X-User-Email");
- when(_requestContext.getHeaderString("X-User-Email")).thenReturn("
[email protected] ");
+ public void testNonBearerAuthorizationReturnsNull() {
+
when(_requestContext.getHeaderString(HttpHeaders.AUTHORIZATION)).thenReturn("Basic
dXNlcjpwYXNz");
AuditEvent.UserIdentity result =
_auditIdentityResolver.resolveIdentity(_requestContext);
- assertThat(result).isNotNull();
- assertThat(result.getPrincipal()).isEqualTo(" [email protected] ");
+ assertThat(result).isNull();
}
+ // ==================== Helper Methods ====================
+
private String createJwtToken(String subject, Map<String, Object>
customClaims)
throws Exception {
JWTClaimsSet.Builder claimsBuilder = new
JWTClaimsSet.Builder().subject(subject)
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/audit/MockAuditTokenResolver.java
b/pinot-common/src/test/java/org/apache/pinot/common/audit/MockAuditTokenResolver.java
new file mode 100644
index 00000000000..19e6943576b
--- /dev/null
+++
b/pinot-common/src/test/java/org/apache/pinot/common/audit/MockAuditTokenResolver.java
@@ -0,0 +1,93 @@
+/**
+ * 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.pinot.common.audit;
+
+import javax.annotation.Nullable;
+import org.apache.pinot.spi.audit.AuditTokenResolver;
+import org.apache.pinot.spi.audit.AuditUserIdentity;
+
+
+/**
+ * Mock implementation of AuditTokenResolver for unit testing.
+ * <p>
+ * Supports two modes:
+ * <ul>
+ * <li>Direct instantiation with configurable return value</li>
+ * <li>PluginManager loading (no-arg constructor) with static
configuration</li>
+ * </ul>
+ */
+public class MockAuditTokenResolver implements AuditTokenResolver {
+
+ private static final String TEST_PREFIX = "Bearer test-";
+ private static final String DEFAULT_PRINCIPAL = "mock-resolved-user";
+
+ private static String _staticReturnValue = DEFAULT_PRINCIPAL;
+ private static String _lastAuthHeader;
+
+ @Nullable
+ private final String _returnValue;
+
+ /**
+ * No-arg constructor for PluginManager loading.
+ */
+ public MockAuditTokenResolver() {
+ _returnValue = null;
+ }
+
+ /**
+ * Constructor for direct instantiation with configurable return value.
+ */
+ public MockAuditTokenResolver(@Nullable String returnValue) {
+ _returnValue = returnValue;
+ }
+
+ @Override
+ @Nullable
+ public AuditUserIdentity resolve(String authHeaderValue) {
+ _lastAuthHeader = authHeaderValue;
+
+ // If instantiated with a specific return value, use it
+ if (_returnValue != null) {
+ return () -> _returnValue;
+ }
+
+ // For PluginManager-loaded instances, check prefix and use static config
+ if (authHeaderValue.startsWith(TEST_PREFIX)) {
+ String principal = _staticReturnValue;
+ return () -> principal;
+ }
+ return null;
+ }
+
+ /**
+ * Resets static state to defaults.
+ */
+ public static void reset() {
+ _staticReturnValue = DEFAULT_PRINCIPAL;
+ _lastAuthHeader = null;
+ }
+
+ /**
+ * Returns the last auth header value passed to resolve().
+ */
+ @Nullable
+ public static String getLastAuthHeader() {
+ return _lastAuthHeader;
+ }
+}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/audit/AuditTokenResolver.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/audit/AuditTokenResolver.java
new file mode 100644
index 00000000000..0a3b788aa21
--- /dev/null
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/audit/AuditTokenResolver.java
@@ -0,0 +1,51 @@
+/**
+ * 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.pinot.spi.audit;
+
+import javax.annotation.Nullable;
+
+
+/**
+ * Service Provider Interface for resolving user identity from authentication
tokens
+ * for audit logging purposes.
+ * <p>
+ * Implementations can handle custom token formats (e.g., proprietary tokens,
API keys)
+ * that are not standard JWTs. The resolver receives the full Authorization
header value
+ * and returns the principal if it can handle the token format.
+ * <p>
+ * If the resolver cannot handle the token format, it should return null to
allow
+ * fallback to the default JWT-based resolution.
+ */
+public interface AuditTokenResolver {
+
+ /**
+ * Resolves the user identity from an authorization header value.
+ * <p>
+ * The implementation should:
+ * <ul>
+ * <li>Return an {@link AuditUserIdentity} if it can handle the token
format</li>
+ * <li>Return null if it cannot handle the token format (to allow
fallback)</li>
+ * </ul>
+ *
+ * @param authHeaderValue the full Authorization header value (e.g., "Bearer
<token>")
+ * @return the resolved identity, or null if this resolver cannot handle the
token
+ */
+ @Nullable
+ AuditUserIdentity resolve(String authHeaderValue);
+}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/audit/AuditUserIdentity.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/audit/AuditUserIdentity.java
new file mode 100644
index 00000000000..e4d6efd88ff
--- /dev/null
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/audit/AuditUserIdentity.java
@@ -0,0 +1,41 @@
+/**
+ * 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.pinot.spi.audit;
+
+import javax.annotation.Nullable;
+
+
+/**
+ * Represents a resolved user identity for audit logging purposes.
+ * <p>
+ * This interface allows {@link AuditTokenResolver} implementations to return
+ * structured identity information that can be extended in the future
+ * (e.g., roles, groups) without breaking the SPI contract.
+ */
+@FunctionalInterface
+public interface AuditUserIdentity {
+
+ /**
+ * Returns the principal (user identifier) for this identity.
+ *
+ * @return the principal, or {@code null} if not available
+ */
+ @Nullable
+ String getPrincipal();
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
