This is an automated email from the ASF dual-hosted git repository. joemcdonnell pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/impala.git
commit 929130b735f83216aa7eed898ab6ee42e1672cfa Author: gaurav1086 <[email protected]> AuthorDate: Fri Mar 7 13:53:26 2025 -0800 IMPALA-13813: OAuth/JWT Avoid key verification on every rpc call This patch optimizes the OAuth/JWT flow by setting cookies in order to avoid token verification in every RPC call. The default cookie expiry time is 1 day. This is only valid for hs2-http protocol. Testing: Modified existing custom cluster tests: test_jwt_auth_valid and test_oauth_auth_valid: - total jwt token verification success count = 1: Reason: Verify jwt/oauth token only the first time and then set the cookie so do not need to re-verify the token for subsequent rpc queries. - total cookie auth success = rpc count - 1: Reason: After first verification, all subsequent authentication will be cookie auth based. - Benchmarking the query SELECT 1; executed 10,000 times with OAuth authentication showed a total time of 2.16s with the cookie enabled vs. 2.38s without the cookie. This indicates a modest performance gain (~9%) when cookie support is enabled. The time command output in both scenarios are: With cookie enabled: - real 2.16 - user 0.99 - sys 0.21 With cookie disabled: - real 2.38 - user 1.12 - sys 0.22 Change-Id: I0e3e5d9cf8bdb99920611b06571515e05e15164e Reviewed-on: http://gerrit.cloudera.org:8080/22600 Reviewed-by: Impala Public Jenkins <[email protected]> Tested-by: Impala Public Jenkins <[email protected]> --- be/src/rpc/authentication.cc | 9 +++++++-- tests/custom_cluster/test_shell_jwt_auth.py | 15 ++++++++++++--- tests/custom_cluster/test_shell_oauth_auth.py | 15 ++++++++++++--- 3 files changed, 31 insertions(+), 8 deletions(-) diff --git a/be/src/rpc/authentication.cc b/be/src/rpc/authentication.cc index 8b022b7fd..7f1739385 100644 --- a/be/src/rpc/authentication.cc +++ b/be/src/rpc/authentication.cc @@ -800,7 +800,10 @@ error_description=\"$0 \"", status.GetDetail())); return false; } connection_context->username = username; - // TODO: cookies are not added, but are not needed right now + + // Create a cookie to return. + connection_context->return_headers.push_back( + Substitute("Set-Cookie: $0", GenerateCookie(username, hash))); return true; } @@ -839,8 +842,10 @@ error_description=\"$0 \"", status.GetDetail())); return false; } connection_context->username = username; - // TODO: cookies are not added, but are not needed right now + // Create a cookie to return. + connection_context->return_headers.push_back( + Substitute("Set-Cookie: $0", GenerateCookie(username, hash))); return true; } diff --git a/tests/custom_cluster/test_shell_jwt_auth.py b/tests/custom_cluster/test_shell_jwt_auth.py index d46d3a5aa..0bfd25479 100644 --- a/tests/custom_cluster/test_shell_jwt_auth.py +++ b/tests/custom_cluster/test_shell_jwt_auth.py @@ -76,10 +76,19 @@ class TestImpalaShellJWTAuth(CustomClusterTestSuite): self.cluster.get_first_impalad().service.wait_for_metric_value( "impala-server.backend-num-queries-executed", 1, timeout=15) - # Ensure the Impala coordinator is correctly reporting the jwt auth metrics - # must be done before the cluster shuts down since it calls to the coordinator + cookie_auth_count = self.cluster.get_first_impalad().service.get_metric_value( + "impala.thrift-server.hiveserver2-http-frontend.total-cookie-auth-success") + + # Impala sets a cookie with the expiry time of 1 day. + # After the first successful token verification, the authentication + # will happen via cookie auth, hence this count will be 1. + self.__assert_success_fail_metric(success_count=1) + + # Total cookie auth success should be 1 less than total rpc_count + # since after the 1st rpc count, the cookie is set and no more jwt token + # verification happens. query_rpc_count = self.__get_rpc_count() - before_rpc_count - self.__assert_success_fail_metric(success_count=query_rpc_count) + assert cookie_auth_count == query_rpc_count - 1, "Incorrect Cookie Auth Count" # Shut down cluster to ensure logs flush to disk. self._stop_impala_cluster() diff --git a/tests/custom_cluster/test_shell_oauth_auth.py b/tests/custom_cluster/test_shell_oauth_auth.py index 653ef5bda..bb559a43c 100644 --- a/tests/custom_cluster/test_shell_oauth_auth.py +++ b/tests/custom_cluster/test_shell_oauth_auth.py @@ -75,10 +75,19 @@ class TestImpalaShellOAuthAuth(CustomClusterTestSuite): self.cluster.get_first_impalad().service.wait_for_metric_value( "impala-server.backend-num-queries-executed", 1, timeout=15) - # Ensure the Impala coordinator is correctly reporting the oauth auth metrics - # must be done before the cluster shuts down since it calls to the coordinator + cookie_auth_count = self.cluster.get_first_impalad().service.get_metric_value( + "impala.thrift-server.hiveserver2-http-frontend.total-cookie-auth-success") + + # Impala sets a cookie with the expiry time of 1 day. + # After the first successful token verification, the authentication + # will happen via cookie auth, hence this count will be 1. + self.__assert_success_fail_metric(success_count=1) + + # Total cookie auth success should be 1 less than total rpc_count + # since after the 1st rpc count, the cookie is set and no more jwt token + # verification happens. query_rpc_count = self.__get_rpc_count() - before_rpc_count - self.__assert_success_fail_metric(success_count=query_rpc_count) + assert cookie_auth_count == query_rpc_count - 1, "Incorrect Cookie Auth Count" # Shut down cluster to ensure logs flush to disk. self._stop_impala_cluster()
