From 99cd7572f113cbbb8f5ef89af60af3734fe8cfa7 Mon Sep 17 00:00:00 2001
From: Sunil G <sunilg@apache.org>
Date: Tue, 1 Oct 2019 20:06:21 +0530
Subject: [PATCH] YARN-9801. SchedConfCli does not work wiwith https mode.
 Contributed by Prabhu Joseph

---
 .../hadoop/yarn/client/cli/SchedConfCLI.java  | 62 +++++++++++++++++--
 1 file changed, 58 insertions(+), 4 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/SchedConfCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/SchedConfCLI.java
index daf4addd72..146db9ec49 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/SchedConfCLI.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/SchedConfCLI.java
@@ -23,6 +23,8 @@
 import com.sun.jersey.api.client.ClientResponse;
 import com.sun.jersey.api.client.WebResource;
 import com.sun.jersey.api.client.WebResource.Builder;
+import com.sun.jersey.client.urlconnection.HttpURLConnectionFactory;
+import com.sun.jersey.client.urlconnection.URLConnectionClientHandler;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.GnuParser;
 import org.apache.commons.cli.MissingArgumentException;
@@ -31,6 +33,8 @@
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.conf.Configured;
+import org.apache.hadoop.security.authentication.client.AuthenticatedURL;
+import org.apache.hadoop.security.ssl.SSLFactory;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.util.Tool;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
@@ -41,6 +45,9 @@
 
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response.Status;
+import java.io.IOException;
+import java.net.HttpURLConnection;
+import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -156,7 +163,12 @@ public int run(String[] args) throws Exception {
   @VisibleForTesting
   int formatSchedulerConf(String webAppAddress, WebResource resource)
       throws Exception {
-    Client webServiceClient = Client.create();
+    Configuration conf = getConf();
+    SSLFactory clientSslFactory = null;
+    if (YarnConfiguration.useHttps(conf)) {
+      clientSslFactory = new SSLFactory(SSLFactory.Mode.CLIENT, conf);
+    }
+    Client webServiceClient = createWebServiceClient(clientSslFactory);
     ClientResponse response = null;
     resource = (resource != null) ? resource :
         webServiceClient.resource(webAppAddress);
@@ -194,14 +206,24 @@ int formatSchedulerConf(String webAppAddress, WebResource resource)
       if (response != null) {
         response.close();
       }
-      webServiceClient.destroy();
+      if (webServiceClient != null) {
+        webServiceClient.destroy();
+      }
+      if (clientSslFactory != null) {
+        clientSslFactory.destroy();
+      }
     }
   }
 
   @VisibleForTesting
   int updateSchedulerConfOnRMNode(String webAppAddress,
       SchedConfUpdateInfo updateInfo) throws Exception {
-    Client webServiceClient = Client.create();
+    Configuration conf = getConf();
+    SSLFactory clientSslFactory = null;
+    if (YarnConfiguration.useHttps(conf)) {
+      clientSslFactory = new SSLFactory(SSLFactory.Mode.CLIENT, conf);
+    }
+    Client webServiceClient = createWebServiceClient(clientSslFactory);
     ClientResponse response = null;
     WebResource resource = webServiceClient.resource(webAppAddress);
 
@@ -236,10 +258,42 @@ int updateSchedulerConfOnRMNode(String webAppAddress,
       if (response != null) {
         response.close();
       }
-      webServiceClient.destroy();
+      if (webServiceClient != null) {
+        webServiceClient.destroy();
+      }
+      if (clientSslFactory != null) {
+        clientSslFactory.destroy();
+      }
     }
   }
 
+  private Client createWebServiceClient(SSLFactory clientSslFactory) {
+    Client webServiceClient = new Client(new URLConnectionClientHandler(
+        new HttpURLConnectionFactory() {
+        @Override
+        public HttpURLConnection getHttpURLConnection(URL url)
+            throws IOException {
+          AuthenticatedURL.Token token = new AuthenticatedURL.Token();
+          AuthenticatedURL aUrl;
+          HttpURLConnection conn = null;
+          try {
+            if (clientSslFactory != null) {
+              clientSslFactory.init();
+              aUrl = new AuthenticatedURL(null, clientSslFactory);
+            } else {
+              aUrl = new AuthenticatedURL();
+            }
+            conn = aUrl.openConnection(url, token);
+          } catch (Exception e) {
+            throw new IOException(e);
+          }
+          return conn;
+        }
+      }));
+    webServiceClient.setChunkedEncodingSize(null);
+    return webServiceClient;
+  }
+
 
   @VisibleForTesting
   void addQueues(String args, SchedConfUpdateInfo updateInfo) {