[tez] branch master updated: TEZ-4102: Let session credentials be merged before merging am launch credentials (László Bodor reviewed by Ashutosh Chauhan)

[abstractdog]

This is an automated email from the ASF dual-hosted git repository.

abstractdog pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/tez.git


The following commit(s) were added to refs/heads/master by this push:
     new 8b5d40f  TEZ-4102: Let session credentials be merged before merging am 
launch credentials (László Bodor reviewed by Ashutosh Chauhan)
8b5d40f is described below

commit 8b5d40fe831a33c05464c84d40a0f622cfd64d35
Author: László Bodor <bodorlaszlo0...@gmail.com>
AuthorDate: Thu Jan 23 11:21:21 2020 +0100

    TEZ-4102: Let session credentials be merged before merging am launch 
credentials (László Bodor reviewed by Ashutosh Chauhan)
    
    Signed-off-by: Laszlo Bodor <bodorlaszlo0...@gmail.com>
---
 .../java/org/apache/tez/client/TezClientUtils.java | 34 +++++++++++++---------
 .../org/apache/tez/client/TestTezClientUtils.java  | 26 +++++++++++++++++
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java 
b/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java
index 2b21024..2aeabe4 100644
--- a/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java
+++ b/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java
@@ -470,19 +470,8 @@ public class TezClientUtils {
     // Setup required Credentials for the AM launch. DAG specific credentials
     // are handled separately.
     ByteBuffer securityTokens = null;
-    // Setup security tokens
-    Credentials amLaunchCredentials = new Credentials();
-    if (amConfig.getCredentials() != null) {
-      amLaunchCredentials.addAll(amConfig.getCredentials());
-    }
-
-    // Add Staging dir creds to the list of session credentials.
-    TokenCache.obtainTokensForFileSystems(sessionCreds, new 
Path[]{binaryConfPath}, conf);
-
-    populateTokenCache(conf, sessionCreds);
-
-    // Add session specific credentials to the AM credentials.
-    amLaunchCredentials.mergeAll(sessionCreds);
+    Credentials amLaunchCredentials =
+        prepareAmLaunchCredentials(amConfig, sessionCreds, conf, 
binaryConfPath);
 
     DataOutputBuffer dob = new DataOutputBuffer();
     amLaunchCredentials.writeTokenStorageToStream(dob);
@@ -703,6 +692,25 @@ public class TezClientUtils {
 
   }
 
+  static Credentials prepareAmLaunchCredentials(AMConfiguration amConfig, 
Credentials sessionCreds,
+      TezConfiguration conf, Path binaryConfPath) throws IOException {
+    // Setup security tokens
+    Credentials amLaunchCredentials = new Credentials();
+
+    // Add Staging dir creds to the list of session credentials.
+    TokenCache.obtainTokensForFileSystems(sessionCreds, new Path[] 
{binaryConfPath }, conf);
+
+    populateTokenCache(conf, sessionCreds);
+
+    // Add session specific credentials to the AM credentials.
+    amLaunchCredentials.mergeAll(sessionCreds);
+
+    if (amConfig.getCredentials() != null) {
+      amLaunchCredentials.mergeAll(amConfig.getCredentials());
+    }
+    return amLaunchCredentials;
+  }
+
   //get secret keys and tokens and store them into TokenCache
   private static void populateTokenCache(TezConfiguration conf, Credentials 
credentials)
           throws IOException{
diff --git 
a/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java 
b/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java
index 581d722..edcec49 100644
--- a/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java
+++ b/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java
@@ -46,8 +46,10 @@ import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.permission.FsPermission;
 import org.apache.hadoop.hdfs.DistributedFileSystem;
 import org.apache.hadoop.io.DataInputByteBuffer;
+import org.apache.hadoop.io.Text;
 import org.apache.hadoop.security.Credentials;
 import org.apache.hadoop.security.token.Token;
+import org.apache.hadoop.security.token.TokenIdentifier;
 import org.apache.hadoop.yarn.api.ApplicationConstants;
 import org.apache.hadoop.yarn.api.ApplicationConstants.Environment;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
@@ -890,5 +892,29 @@ public class TestTezClientUtils {
 
   }
 
+  @Test
+  public void testSessionCredentialsMergedBeforeAmConfigCredentials() throws 
Exception {
+    TezConfiguration conf = new TezConfiguration();
+    Text tokenType = new Text("TEST_TOKEN_TYPE");
+    Text tokenKind = new Text("TEST_TOKEN_KIND");
+    Text tokenService = new Text("TEST_TOKEN_SERVICE");
+
+    Credentials amConfigCredentials = new Credentials();
+    amConfigCredentials.addToken(tokenType,
+        new Token<>("id1".getBytes(), null, tokenKind, tokenService));
+
+    Credentials sessionCredentials = new Credentials();
+    Token<TokenIdentifier> sessionToken =
+        new Token<>("id2".getBytes(), null, tokenKind, tokenService);
+    sessionCredentials.addToken(tokenType, sessionToken);
 
+    AMConfiguration amConfig = new AMConfiguration(conf, null, 
amConfigCredentials);
+
+    Credentials amLaunchCredentials =
+        TezClientUtils.prepareAmLaunchCredentials(amConfig, 
sessionCredentials, conf, null);
+
+    // if there is another token in am conf creds of the same token type,
+    // session token should be applied while creating ContainerLaunchContext
+    Assert.assertEquals(sessionToken, amLaunchCredentials.getToken(tokenType));
+  }
 }