IMPALA-12093: impala-shell to preserve all cookies

Updates impala-shell to preserve all cookies by default, defined as
setting 'http_cookie_names=*'. Prior behavior of restricting cookies to
a user-specified list is preserved when 'http_cookie_names' is given any
value besides '*'. Setting 'http_cookie_names=' prevents any cookies
from being preserved.

Adds verbose output that prints all cookies that are preserved by the
HTTP client.

Existing cookie tests with LDAP still work. Adds a test where Impala
returns an extra cookie, and test verifies that verbose mode prints all
expected cookies.

Change-Id: Ic81f790288460b086ab218e6701e8115a996dfa7
Reviewed-on: http://gerrit.cloudera.org:8080/19827
Reviewed-by: Impala Public Jenkins <[email protected]>
Tested-by: Michael Smith <[email protected]>

Author: MikaelSmith

be/src/rpc/authentication.cc

@@ -218,6 +218,9 @@ DEFINE_bool(enable_group_filter_check_for_authenticated_kerberos_user, false,
     "is the same Active Directory server). "
     "The default value is false, which provides backwards-compatible behavior.");
 
+DEFINE_string_hidden(test_cookie, "",
+    "Adds Set-Cookie: <test_cookie> to BasicAuth for testing.");
+
 namespace impala {
 
 // Sasl callbacks.  Why are these here?  Well, Sasl isn't that bright, and
@@ -674,6 +677,10 @@ bool BasicAuth(ThriftServer::ConnectionContext* connection_context,
     // Create a cookie to return.
     connection_context->return_headers.push_back(
         Substitute("Set-Cookie: $0", GenerateCookie(username, hash)));
+    if (!FLAGS_test_cookie.empty()) {
+      connection_context->return_headers.push_back(
+          Substitute("Set-Cookie: $0", FLAGS_test_cookie));
+    }
     return true;
   }
   connection_context->return_headers.push_back("WWW-Authenticate: Basic");

fe/src/test/java/org/apache/impala/customcluster/LdapImpalaShellTest.java

@@ -19,12 +19,15 @@
 
 import static org.apache.impala.testutil.LdapUtil.*;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import com.google.common.collect.Range;
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import org.apache.directory.server.annotations.CreateLdapServer;
 import org.apache.directory.server.annotations.CreateTransport;
@@ -76,7 +79,8 @@ protected boolean pythonSupportsSSLContext() throws Exception {
     // python -c "import ssl; print hasattr(ssl, 'create_default_context')"
     String[] cmd = {
         "python", "-c", "import ssl; print(hasattr(ssl, 'create_default_context'))"};
-    return Boolean.parseBoolean(RunShellCommand.Run(cmd, true, "", "").replace("\n", ""));
+    return Boolean.parseBoolean(
+        RunShellCommand.Run(cmd, true, "", "").stdout.replace("\n", ""));
   }
 
   /**
@@ -129,33 +133,34 @@ protected String[] buildCommand(
   /**
    * Tests ldap authentication using impala-shell.
    */
-  protected void testShellLdapAuthImpl() throws Exception {
+  protected void testShellLdapAuthImpl(String extra_cookie) throws Exception {
     String query = "select logged_in_user()";
     // Templated shell commands to test a simple 'show tables' command.
     // 1. Valid username, password and default http_cookie_names. Should succeed with
     // cookies are being used.
     String[] validCommand = {"impala-shell.sh", "", "--ldap", "--auth_creds_ok_in_clear",
+        "--verbose",
         String.format("--user=%s", TEST_USER_1),
         String.format("--ldap_password_cmd=printf %s", TEST_PASSWORD_1),
         String.format("--query=%s", query)};
     // 2. Valid username, password and matching http_cookie_names. Should succeed with
     // cookies are being used.
     String[] validCommandMatchingCookieNames = {"impala-shell.sh", "", "--ldap",
-        "--auth_creds_ok_in_clear", "--http_cookie_names=impala.auth",
+        "--auth_creds_ok_in_clear", "--verbose", "--http_cookie_names=impala.auth",
         String.format("--user=%s", TEST_USER_1),
         String.format("--ldap_password_cmd=printf %s", TEST_PASSWORD_1),
         String.format("--query=%s", query)};
     // 3. Valid username and password, but not matching http_cookie_names. Should succeed
     // with cookies are not being used.
     String[] validCommandMismatchingCookieNames = {"impala-shell.sh", "", "--ldap",
-        "--auth_creds_ok_in_clear", "--http_cookie_names=impala.conn",
+        "--auth_creds_ok_in_clear", "--verbose", "--http_cookie_names=impala.conn",
         String.format("--user=%s", TEST_USER_1),
         String.format("--ldap_password_cmd=printf %s", TEST_PASSWORD_1),
         String.format("--query=%s", query)};
     // 4. Valid username and password, but empty http_cookie_names. Should succeed with
     // cookies are not being used.
     String[] validCommandEmptyCookieNames = {"impala-shell.sh", "", "--ldap",
-        "--auth_creds_ok_in_clear", "--http_cookie_names=",
+        "--auth_creds_ok_in_clear", "--verbose", "--http_cookie_names=",
         String.format("--user=%s", TEST_USER_1),
         String.format("--ldap_password_cmd=printf %s", TEST_PASSWORD_1),
         String.format("--query=%s", query)};
@@ -168,35 +173,53 @@ protected void testShellLdapAuthImpl() throws Exception {
         "impala-shell.sh", "", String.format("--query=%s", query)};
     String protocolTemplate = "--protocol=%s";
 
+    // Sorted list of cookies for validCommand, where all cookies are preserved.
+    List<String> preservedCookiesList = new ArrayList<>();
+    preservedCookiesList.add("impala.auth");
+    if (extra_cookie != null) {
+      preservedCookiesList.add(extra_cookie);
+    }
+    Collections.sort(preservedCookiesList);
+    String preservedCookies = String.join(", ", preservedCookiesList);
+
     for (String p : getProtocolsToTest()) {
       String protocol = String.format(protocolTemplate, p);
       validCommand[1] = protocol;
-      RunShellCommand.Run(validCommand, /*shouldSucceed*/ true, TEST_USER_1,
+      RunShellCommand.Output result = RunShellCommand.Run(validCommand,
+          /*shouldSucceed*/ true, TEST_USER_1,
           "Starting Impala Shell with LDAP-based authentication");
       if (p.equals("hs2-http")) {
         // Check that cookies are being used.
         verifyMetrics(Range.atLeast(1L), zero, Range.atLeast(1L), zero);
+        assertTrue(result.stderr,
+            result.stderr.contains("Preserving cookies: " + preservedCookies));
       }
       validCommandMatchingCookieNames[1] = protocol;
-      RunShellCommand.Run(validCommandMatchingCookieNames, /*shouldSucceed*/ true,
-          TEST_USER_1, "Starting Impala Shell with LDAP-based authentication");
+      result = RunShellCommand.Run(validCommandMatchingCookieNames,
+          /*shouldSucceed*/ true, TEST_USER_1,
+          "Starting Impala Shell with LDAP-based authentication");
       if (p.equals("hs2-http")) {
         // Check that cookies are being used.
         verifyMetrics(Range.atLeast(2L), zero, Range.atLeast(2L), zero);
+        assertTrue(result.stderr,
+            result.stderr.contains("Preserving cookies: impala.auth"));
       }
       validCommandMismatchingCookieNames[1] = protocol;
-      RunShellCommand.Run(validCommandMismatchingCookieNames, /*shouldSucceed*/ true,
-          TEST_USER_1, "Starting Impala Shell with LDAP-based authentication");
+      result = RunShellCommand.Run(validCommandMismatchingCookieNames,
+          /*shouldSucceed*/ true, TEST_USER_1,
+          "Starting Impala Shell with LDAP-based authentication");
       if (p.equals("hs2-http")) {
         // Check that cookies are NOT being used.
         verifyMetrics(Range.atLeast(2L), zero, Range.atLeast(2L), zero);
+        assertFalse(result.stderr, result.stderr.contains("Preserving cookies:"));
       }
       validCommandEmptyCookieNames[1] = protocol;
-      RunShellCommand.Run(validCommandEmptyCookieNames, /*shouldSucceed*/ true,
+      result = RunShellCommand.Run(validCommandEmptyCookieNames, /*shouldSucceed*/ true,
           TEST_USER_1, "Starting Impala Shell with LDAP-based authentication");
       if (p.equals("hs2-http")) {
         // Check that cookies are NOT being used.
         verifyMetrics(Range.atLeast(2L), zero, Range.atLeast(2L), zero);
+        assertFalse(result.stderr, result.stderr.contains("Preserving cookies:"));
       }
 
       invalidCommand[1] = protocol;

fe/src/test/java/org/apache/impala/customcluster/LdapSearchBindImpalaShellTest.java

@@ -97,7 +97,7 @@ private String getKerberosArgs() throws IOException {
   public void testShellLdapAuth() throws Exception {
     setUp("--ldap_user_search_basedn=dc=myorg,dc=com "
         + "--ldap_user_filter=(&(objectClass=person)(cn={0}))");
-    testShellLdapAuthImpl();
+    testShellLdapAuthImpl(null);
   }
 
   /**

fe/src/test/java/org/apache/impala/customcluster/LdapSimpleBindImpalaShellTest.java

@@ -94,8 +94,8 @@ private String getKerberosArgs() throws IOException {
    */
   @Test
   public void testShellLdapAuth() throws Exception {
-    setUp("");
-    testShellLdapAuthImpl();
+    setUp("--test_cookie=impala.ldap=testShellLdapAuth");
+    testShellLdapAuthImpl("impala.ldap");
   }
 
   /**

fe/src/test/java/org/apache/impala/customcluster/RunShellCommand.java

@@ -27,12 +27,25 @@
  * Helper class to run a shell command.
  */
 class RunShellCommand {
+  /**
+   * Tuple wrapping stdout, stderr from Run.
+   */
+  public static class Output {
+    public Output(String out, String err) {
+      stdout = out;
+      stderr = err;
+    }
+
+    public String stdout;
+    public String stderr;
+  }
+
   /**
    * Run a shell command 'cmd'. If 'shouldSucceed' is true, the command is expected to
-   * succeed, otherwise it is expected to fail. Returns the output (stdout) of the
+   * succeed, otherwise it is expected to fail. Returns the output (stdout, stderr) of the
    * command.
    */
-  public static String Run(String[] cmd, boolean shouldSucceed, String expectedOut,
+  public static Output Run(String[] cmd, boolean shouldSucceed, String expectedOut,
       String expectedErr) throws Exception {
     // run the command with the env variables inherited from the current process
     return Run(cmd, null, shouldSucceed, expectedOut, expectedErr);
@@ -41,10 +54,10 @@ public static String Run(String[] cmd, boolean shouldSucceed, String expectedOut
   /**
    * Run a shell command 'cmd' with custom 'env' variables.
    * If 'shouldSucceed' is true, the command is expected to
-   * succeed, otherwise it is expected to fail. Returns the output (stdout) of the
+   * succeed, otherwise it is expected to fail. Returns the output (stdout, stderr) of the
    * command.
    */
-  public static String Run(String[] cmd, String[] env, boolean shouldSucceed,
+  public static Output Run(String[] cmd, String[] env, boolean shouldSucceed,
                            String expectedOut, String expectedErr) throws Exception {
     Runtime rt = Runtime.getRuntime();
     Process process = rt.exec(cmd, env);
@@ -71,6 +84,6 @@ public static String Run(String[] cmd, String[] env, boolean shouldSucceed,
     // If the query succeeds, assert that the output is correct.
     String stdout = stdoutBuf.toString();
     if (shouldSucceed) assertTrue(stdout, stdout.contains(expectedOut));
-    return stdout;
+    return new Output(stdout, stderr);
   }
 }

shell/ImpalaHttpClient.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+from __future__ import print_function, unicode_literals
 
 from io import BytesIO
 import os
@@ -31,7 +32,7 @@
 
 from thrift.transport.TTransport import TTransportBase
 from shell_exceptions import HttpError, AuthenticationException
-from cookie_util import get_all_matching_cookies, get_cookie_expiry
+from cookie_util import get_all_matching_cookies, get_all_cookies, get_cookie_expiry
 
 import six
 
@@ -52,7 +53,8 @@ class ImpalaHttpClient(TTransportBase):
   MIN_REQUEST_SIZE_FOR_EXPECT = 1024
 
   def __init__(self, uri_or_host, port=None, path=None, cafile=None, cert_file=None,
-      key_file=None, ssl_context=None, http_cookie_names=None, socket_timeout_s=None):
+      key_file=None, ssl_context=None, http_cookie_names=None, socket_timeout_s=None,
+      verbose=False):
     """ImpalaHttpClient supports two different types of construction:
 
     ImpalaHttpClient(host, port, path) - deprecated
@@ -66,7 +68,8 @@ def __init__(self, uri_or_host, port=None, path=None, cafile=None, cert_file=Non
     http_cookie_names is used to specify a comma-separated list of possible cookie names
     used for cookie-based authentication or session management. If a cookie with one of
     these names is returned in an http response by the server or an intermediate proxy
-    then it will be included in each subsequent request for the same connection.
+    then it will be included in each subsequent request for the same connection. If it
+    is set as wildcards, all cookies in an http response will be preserved.
     """
     if port is not None:
       warnings.warn(
@@ -111,16 +114,22 @@ def __init__(self, uri_or_host, port=None, path=None, cafile=None, cert_file=Non
     else:
       self.realhost = self.realport = self.proxy_auth = None
     if (not http_cookie_names) or (str(http_cookie_names).strip() == ""):
+      self.__preserve_all_cookies = False
       self.__http_cookie_dict = None
       self.__auth_cookie_names = None
+    elif str(http_cookie_names).strip() == "*":
+      self.__preserve_all_cookies = True
+      self.__http_cookie_dict = dict()
+      self.__auth_cookie_names = set()
     else:
+      self.__preserve_all_cookies = False
       # Build a dictionary that maps cookie name to namedtuple.
       cookie_names = http_cookie_names.split(',')
       self.__http_cookie_dict = \
           {cn: Cookie(cookie=None, expiry_time=None) for cn in cookie_names}
       # Store the auth cookie names in __auth_cookie_names.
       # Assume auth cookie names end with ".auth".
-      self.__auth_cookie_names = [cn for cn in cookie_names if cn.endswith(".auth")]
+      self.__auth_cookie_names = {cn for cn in cookie_names if cn.endswith(".auth")}
     # Set __are_matching_cookies_found as True if matching cookies are found in response.
     self.__are_matching_cookies_found = False
     self.__wbuf = BytesIO()
@@ -134,6 +143,7 @@ def __init__(self, uri_or_host, port=None, path=None, cafile=None, cert_file=Non
     self.__basic_auth = None
     self.__kerb_service = None
     self.__add_custom_headers_funcs = []
+    self.__verbose = verbose
 
   @staticmethod
   def basic_proxy_auth_header(proxy):
@@ -309,13 +319,25 @@ def getHttpCookieHeaderForRequest(self):
   # Extract cookies from response and save those cookies for which the cookie names
   # are in the cookie name list specified in the __init__().
   def extractHttpCookiesFromResponse(self):
-    if self.__http_cookie_dict:
+    if self.__preserve_all_cookies:
+      matching_cookies = get_all_cookies(self.path, self.headers)
+    elif self.__http_cookie_dict:
       matching_cookies = get_all_matching_cookies(
           self.__http_cookie_dict.keys(), self.path, self.headers)
-      if matching_cookies:
-        self.__are_matching_cookies_found = True
-        for c in matching_cookies:
-          self.__http_cookie_dict[c.key] = Cookie(c, get_cookie_expiry(c))
+    else:
+      matching_cookies = None
+
+    if matching_cookies:
+      if self.__verbose:
+        names = sorted([c.key for c in matching_cookies])
+        print("Preserving cookies: " + ', '.join(names), file=sys.stderr)
+        sys.stderr.flush()
+
+      self.__are_matching_cookies_found = True
+      for c in matching_cookies:
+        self.__http_cookie_dict[c.key] = Cookie(c, get_cookie_expiry(c))
+        if c.key.endswith(".auth"):
+          self.__auth_cookie_names.add(c.key)
 
   # Return True if there are any saved cookies which are sent in previous request.
   def areHttpCookiesSaved(self):

shell/cookie_util.py

@@ -50,8 +50,7 @@ def get_cookie_expiry(c):
   return None
 
 
-def get_all_matching_cookies(cookie_names, path, resp_headers):
-  matching_cookies = None
+def get_cookies(resp_headers):
   if 'Set-Cookie' not in resp_headers:
     return None
 
@@ -63,9 +62,28 @@ def get_all_matching_cookies(cookie_names, path, resp_headers):
       cookie_headers = resp_headers.get_all('Set-Cookie')
       for header in cookie_headers:
         cookies.load(header)
+    return cookies
   except Exception:
     return None
 
+
+def get_all_cookies(path, resp_headers):
+  cookies = get_cookies(resp_headers)
+  if not cookies:
+    return None
+
+  matching_cookies = []
+  for c in cookies.values():
+    if c and cookie_matches_path(c, path):
+      matching_cookies.append(c)
+  return matching_cookies
+
+
+def get_all_matching_cookies(cookie_names, path, resp_headers):
+  cookies = get_cookies(resp_headers)
+  if not cookies:
+    return None
+
   matching_cookies = []
   for cn in cookie_names:
     if cn in cookies:

shell/impala_client.py

@@ -414,11 +414,13 @@ def _get_http_transport(self, connect_timeout_ms):
       url = "https://{0}/{1}".format(host_and_port, self.http_path)
       transport = ImpalaHttpClient(url, ssl_context=ssl_ctx,
                                    http_cookie_names=self.http_cookie_names,
-                                   socket_timeout_s=self.http_socket_timeout_s)
+                                   socket_timeout_s=self.http_socket_timeout_s,
+                                   verbose=self.verbose)
     else:
       url = "http://{0}/{1}".format(host_and_port, self.http_path)
       transport = ImpalaHttpClient(url, http_cookie_names=self.http_cookie_names,
-                                   socket_timeout_s=self.http_socket_timeout_s)
+                                   socket_timeout_s=self.http_socket_timeout_s,
+                                   verbose=self.verbose)
 
     if self.use_ldap:
       # Set the BASIC authorization

shell/option_parser.py

@@ -333,12 +333,14 @@ def get_option_parser(defaults):
                     "spooling is disabled only a single row batch can be fetched at a "
                     "time regardless of the specified fetch_size.")
   parser.add_option("--http_cookie_names", dest="http_cookie_names",
-                    default="impala.auth,impala.session.id",
+                    default="*",
                     help="A comma-separated list of HTTP cookie names that are supported "
                     "by the impala-shell. If a cookie with one of these names is "
                     "returned in an http response by the server or an intermediate proxy "
                     "then it will be included in each subsequent request for the same "
-                    "connection.")
+                    "connection. If set to wildcard (*), all cookies in an http response "
+                    "will be preserved. The name of an authentication cookie must end "
+                    "with '.auth', for example 'impala.auth'.")
   parser.add_option("--no_http_tracing", dest="no_http_tracing",
                     action="store_true",
                     help="Tracing http headers 'X-Request-Id', 'X-Impala-Session-Id', "