From 41e81044adf3fa2facc56ee8d8bfa6a2835f3afe Mon Sep 17 00:00:00 2001
From: Adrika Gupta <adrikagupta@google.com>
Date: Tue, 27 Aug 2024 18:50:59 +0530
Subject: [PATCH] added propagating internal credential capability in CDF
 router

---
 .../auth/ProxyUserIdentityExtractor.java      | 14 ++++++++++-
 .../auth/ProxyUserIdentityExtractorTest.java  | 25 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/cdap-security/src/main/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractor.java b/cdap-security/src/main/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractor.java
index b350da5849cf..805c6569deaa 100644
--- a/cdap-security/src/main/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractor.java
+++ b/cdap-security/src/main/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractor.java
@@ -19,6 +19,8 @@
 import com.google.inject.Inject;
 import io.cdap.cdap.common.conf.CConfiguration;
 import io.cdap.cdap.common.conf.Constants;
+import io.cdap.cdap.proto.security.Credential;
+import io.cdap.cdap.security.auth.UserIdentity.IdentifierType;
 import io.netty.handler.codec.http.HttpHeaderNames;
 import io.netty.handler.codec.http.HttpRequest;
 import java.util.LinkedHashSet;
@@ -68,19 +70,29 @@ public UserIdentityExtractionResponse extract(HttpRequest request)
     UserIdentity identity = new UserIdentity(userIdentity, UserIdentity.IdentifierType.EXTERNAL,
         new LinkedHashSet<>(), now, now + EXPIRATION_SECS);
 
-    // Parse the access token from authorization header. The header will be in "Bearer" form.
+    // Parse the access token from authorization header. The header will be in "Bearer" form, if external
+    // or "CDAP-Internal" if internal
     String auth = request.headers().get(HttpHeaderNames.AUTHORIZATION);
     LOG.trace("Extracted user identity header '{}' and authorization header length '{}'",
         userIdentity,
         auth == null ? "NULL" : String.valueOf(auth.length()));
     String userCredential = null;
+    String prefix = "";
     if (auth != null) {
       int idx = auth.trim().indexOf(' ');
       if (idx < 0) {
         return new UserIdentityExtractionResponse(new UserIdentityPair(null, identity));
       }
+      prefix = auth.substring(0, idx).trim();
       userCredential = auth.substring(idx + 1).trim();
     }
+
+    if (Credential.CREDENTIAL_TYPE_INTERNAL.equalsIgnoreCase(prefix)) {
+      UserIdentity internalIdentity = new UserIdentity(userIdentity, IdentifierType.INTERNAL,
+          new LinkedHashSet<>(), now, now + EXPIRATION_SECS);
+      return new UserIdentityExtractionResponse(new UserIdentityPair(userCredential, internalIdentity));
+    }
+
     return new UserIdentityExtractionResponse(new UserIdentityPair(userCredential, identity));
   }
 }
diff --git a/cdap-security/src/test/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractorTest.java b/cdap-security/src/test/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractorTest.java
index c0e5d6406973..7113e8c87022 100644
--- a/cdap-security/src/test/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractorTest.java
+++ b/cdap-security/src/test/java/io/cdap/cdap/security/auth/ProxyUserIdentityExtractorTest.java
@@ -20,6 +20,7 @@
 
 import io.cdap.cdap.common.conf.CConfiguration;
 import io.cdap.cdap.common.conf.Constants;
+import io.cdap.cdap.security.auth.UserIdentity.IdentifierType;
 import io.netty.handler.codec.http.DefaultHttpHeaders;
 import io.netty.handler.codec.http.DefaultHttpRequest;
 import io.netty.handler.codec.http.HttpHeaderNames;
@@ -90,6 +91,30 @@ public void testValidUserReturnsExpectedIdentity() throws UserIdentityExtraction
     Assert.assertEquals(testUserId, identity.getUserIdentity().getUsername());
   }
 
+  @Test
+  public void testValidInternalUserReturnsExpectedInternalIdentity() throws UserIdentityExtractionException {
+    String testUserId = "test-user-id";
+    String testUserIdHeader = "X-User-Id";
+    String testAuthToken = "test-auth-token";
+    CConfiguration config = Mockito.mock(CConfiguration.class);
+    when(config.get(Constants.Security.Authentication.PROXY_USER_ID_HEADER)).thenReturn(testUserIdHeader);
+
+    ProxyUserIdentityExtractor extractor = new ProxyUserIdentityExtractor(config);
+
+    DefaultHttpHeaders headers = new DefaultHttpHeaders();
+    headers.add(HttpHeaderNames.AUTHORIZATION, String.format("CDAP-Internal %s", testAuthToken));
+    headers.add(testUserIdHeader, testUserId);
+    DefaultHttpRequest request = new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET,
+        "http://www.example.com", headers);
+    UserIdentityExtractionResponse response = extractor.extract(request);
+    Assert.assertTrue(response.success());
+    UserIdentityPair identity = response.getIdentityPair();
+
+    Assert.assertEquals(testAuthToken, identity.getUserCredential());
+    Assert.assertEquals(testUserId, identity.getUserIdentity().getUsername());
+    Assert.assertEquals(identity.getUserIdentity().getIdentifierType(), IdentifierType.INTERNAL);
+  }
+
   @Test
   public void testValidUserWithoutCredentialReturnsExpectedIdentity() throws UserIdentityExtractionException {
     String testUserId = "test-user-id";