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 <impala-public-jenkins@cloudera.com>
Tested-by: Michael Smith <michael.smith@cloudera.com>
This commit is contained in:
Michael Smith
2024-06-21 16:45:56 -07:00
parent aea057f095
commit 100693d5ad
9 changed files with 121 additions and 34 deletions

View File

@@ -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");

View File

@@ -19,12 +19,15 @@ package org.apache.impala.customcluster;
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 @@ public class LdapImpalaShellTest {
// 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 @@ public class LdapImpalaShellTest {
/**
* 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 @@ public class LdapImpalaShellTest {
"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;

View File

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

View File

@@ -94,8 +94,8 @@ public class LdapSimpleBindImpalaShellTest extends LdapImpalaShellTest {
*/
@Test
public void testShellLdapAuth() throws Exception {
setUp("");
testShellLdapAuthImpl();
setUp("--test_cookie=impala.ldap=testShellLdapAuth");
testShellLdapAuthImpl("impala.ldap");
}
/**

View File

@@ -27,12 +27,25 @@ import java.io.InputStreamReader;
* 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 @@ class RunShellCommand {
/**
* 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 @@ class RunShellCommand {
// 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);
}
}

View File

@@ -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 six.moves import urllib, http_client
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 @@ class ImpalaHttpClient(TTransportBase):
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 @@ class ImpalaHttpClient(TTransportBase):
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 @@ class ImpalaHttpClient(TTransportBase):
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 @@ class ImpalaHttpClient(TTransportBase):
# 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):

View File

@@ -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:

View File

@@ -414,11 +414,13 @@ class ImpalaClient(object):
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

View File

@@ -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', "