-
Notifications
You must be signed in to change notification settings - Fork 145
feat(kernel): JWT private-key M2M auth on use_kernel=True #921
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -156,11 +156,16 @@ def kernel_auth_kwargs( | |
|
|
||
| (``azure-oauth`` is rejected as unsupported before these guards — | ||
| PECOBLR-4120.) | ||
| 1. **OAuth M2M** — ``oauth_client_id`` + ``oauth_client_secret`` | ||
| 1. **OAuth M2M (JWT private key)** — ``oauth_jwt_key_file`` present → | ||
| forward the private-key + ``oauth_client_id`` + ``oauth_jwt_kid`` | ||
| to the kernel's ``oauth-m2m-jwt`` (RFC 7523 client assertion). The | ||
| kernel signs the assertion and owns the token lifecycle. Checked | ||
| first because a private-key file is unambiguous JWT M2M intent. | ||
| 2. **OAuth M2M** — ``oauth_client_id`` + ``oauth_client_secret`` | ||
| both present → forward raw creds to the kernel's ``oauth-m2m``. | ||
| 2. **PAT** — the built provider is (or wraps) an | ||
| 3. **PAT** — the built provider is (or wraps) an | ||
| ``AccessTokenAuthProvider`` → extract the bearer token. | ||
| 3. **OAuth U2M** — ``auth_type`` is ``databricks-oauth`` → forward the | ||
| 4. **OAuth U2M** — ``auth_type`` is ``databricks-oauth`` → forward the | ||
| connector's coupled ``databricks-sql-python`` bundle (``client_id`` | ||
| + ``redirect_ports`` list, defaulting scopes to ``PYSQL_OAUTH_SCOPES`` | ||
| when the caller supplies none) to the kernel's ``oauth-u2m``, so a | ||
|
|
@@ -169,9 +174,9 @@ def kernel_auth_kwargs( | |
| ``databricks-sql-connector`` default (PECOBLR-4039/4040). Unlike the | ||
| Thrift path, a caller-supplied ``oauth_scopes`` is honored here. | ||
| ``azure-oauth`` is rejected as unsupported (PECOBLR-4120). | ||
| 4. **Custom credentials_provider** → ``NotSupportedError`` (opaque | ||
| 5. **Custom credentials_provider** → ``NotSupportedError`` (opaque | ||
| token source; no raw creds for the kernel to own). | ||
| 5. Anything else → ``NotSupportedError``. | ||
| 6. Anything else → ``NotSupportedError``. | ||
|
|
||
| M2M is checked before PAT so that a workload passing both an | ||
| access token *and* M2M creds resolves to the (refreshing) M2M path | ||
|
|
@@ -186,7 +191,12 @@ def kernel_auth_kwargs( | |
| client_secret = opts.get("oauth_client_secret") | ||
| federation_client_id = opts.get("identity_federation_client_id") | ||
| auth_type = opts.get("auth_type") | ||
| jwt_key_file = opts.get("oauth_jwt_key_file") | ||
| has_m2m = bool(client_id and client_secret) | ||
| # A private-key file is unambiguous JWT client-assertion M2M intent | ||
| # (RFC 7523): the kernel signs a short-lived assertion with the key | ||
| # rather than sending a client secret. | ||
| has_jwt_m2m = bool(jwt_key_file) | ||
|
|
||
| # azure-oauth (Azure AD U2M) is not yet supported on the kernel path. | ||
| # Reject it up front — before any M2M/U2M routing — so ANY azure-oauth | ||
|
|
@@ -223,10 +233,78 @@ def kernel_auth_kwargs( | |
| "(machine-to-machine). Drop oauth_client_secret for U2M, or drop " | ||
| "auth_type for M2M." | ||
| ) | ||
| if has_jwt_m2m and client_secret: | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🟡 Medium — The new JWT M2M path is missing an ambiguity guard against U2M that its shared-secret sibling has. The shared-secret M2M path rejects The JWT branch adds guards for Consider adding a parallel guard, e.g.: if has_jwt_m2m and auth_type == "databricks-oauth":
raise NotSupportedError(
"Ambiguous auth on use_kernel=True: auth_type='databricks-oauth' "
"selects the U2M browser flow, but oauth_jwt_key_file (JWT "
"private-key M2M) was also provided. Drop oauth_jwt_key_file for "
"U2M, or drop auth_type for JWT M2M."
)A unit test alongside |
||
| raise NotSupportedError( | ||
| "Ambiguous auth on use_kernel=True: both oauth_jwt_key_file " | ||
| "(JWT private-key M2M) and oauth_client_secret (shared-secret " | ||
| "M2M) were provided. Pass exactly one — a private key for " | ||
| "JWT client-assertion M2M, or a client secret for shared-secret M2M." | ||
| ) | ||
| if has_jwt_m2m and opts.get("credentials_provider") is not None: | ||
| raise NotSupportedError( | ||
| "Ambiguous auth on use_kernel=True: both a custom " | ||
| "credentials_provider and oauth_jwt_key_file were provided. " | ||
| "Pass exactly one — oauth_client_id + oauth_jwt_key_file for " | ||
| "kernel-managed JWT private-key M2M, or use the Thrift backend " | ||
| "(default) for credentials_provider." | ||
| ) | ||
| if has_jwt_m2m and auth_type == "databricks-oauth": | ||
| raise NotSupportedError( | ||
| f"Ambiguous auth on use_kernel=True: auth_type={auth_type!r} selects " | ||
| "the U2M browser flow, but oauth_jwt_key_file was also provided " | ||
| "(JWT private-key M2M). Drop oauth_jwt_key_file for U2M, or drop " | ||
| "auth_type for JWT M2M." | ||
| ) | ||
|
|
||
| # 1. OAuth M2M — raw client-credentials pair forwarded to the kernel. | ||
| if has_m2m: | ||
| # 1. OAuth M2M (JWT private-key client assertion) — the kernel signs a | ||
| # short-lived assertion with the private key and runs the | ||
| # client-credentials grant. Checked before shared-secret M2M and PAT | ||
| # because a private-key file is unambiguous JWT M2M intent. Requires | ||
| # oauth_client_id (the service principal / OAuth client) and | ||
| # oauth_jwt_kid (the key id the IdP uses to select the registered | ||
| # public key). Optional oauth_jwt_passphrase / oauth_jwt_algorithm / | ||
| # oauth_scopes / token_url are forwarded when present; the kernel | ||
| # fills defaults (RS256 algorithm, all-apis scope, OIDC discovery) | ||
| # for any omitted. | ||
| if has_jwt_m2m: | ||
| if not client_id: | ||
| raise ProgrammingError( | ||
| "use_kernel=True JWT private-key M2M (oauth_jwt_key_file) " | ||
| "requires oauth_client_id (the service principal / OAuth " | ||
| "client id used as the assertion issuer and subject)." | ||
| ) | ||
| jwt_kid = opts.get("oauth_jwt_kid") | ||
| if not jwt_kid: | ||
| raise ProgrammingError( | ||
| "use_kernel=True JWT private-key M2M (oauth_jwt_key_file) " | ||
| "requires oauth_jwt_kid (the key id written into the JWT " | ||
| "header so the IdP can select the registered public key)." | ||
| ) | ||
| kwargs: Dict[str, Any] = { | ||
| "auth_type": "oauth-m2m-jwt", | ||
| "client_id": client_id, | ||
| "jwt_key_file": jwt_key_file, | ||
| "jwt_kid": jwt_kid, | ||
| } | ||
| jwt_passphrase = opts.get("oauth_jwt_passphrase") | ||
| if jwt_passphrase: | ||
| kwargs["jwt_passphrase"] = jwt_passphrase | ||
| jwt_algorithm = opts.get("oauth_jwt_algorithm") | ||
| if jwt_algorithm: | ||
| kwargs["jwt_algorithm"] = jwt_algorithm | ||
| token_url = opts.get("token_url") | ||
| if token_url: | ||
| kwargs["token_url"] = token_url | ||
| scopes = _normalize_scopes(opts.get("oauth_scopes")) | ||
| if scopes is not None: | ||
| kwargs["oauth_scopes"] = scopes | ||
| if federation_client_id: | ||
| kwargs["identity_federation_client_id"] = federation_client_id | ||
| return kwargs | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🔵 Low — Inline step-number comments weren't renumbered when the JWT branch was inserted. The docstring was correctly updated to 1=JWT, 2=M2M, 3=PAT, 4=U2M, 5=creds_provider, 6=else, but the inline |
||
|
|
||
| # 2. OAuth M2M — raw client-credentials pair forwarded to the kernel. | ||
| if has_m2m: | ||
| kwargs = { | ||
| "auth_type": "oauth-m2m", | ||
| "client_id": client_id, | ||
| "client_secret": client_secret, | ||
|
|
@@ -238,7 +316,7 @@ def kernel_auth_kwargs( | |
| kwargs["identity_federation_client_id"] = federation_client_id | ||
| return kwargs | ||
|
|
||
| # 2. PAT (including TokenFederationProvider-wrapped PAT). | ||
| # 3. PAT (including TokenFederationProvider-wrapped PAT). | ||
| if _is_pat(auth_provider): | ||
| token = _extract_bearer_token(auth_provider) | ||
| if not token: | ||
|
|
@@ -251,7 +329,7 @@ def kernel_auth_kwargs( | |
| kwargs["identity_federation_client_id"] = federation_client_id | ||
| return kwargs | ||
|
|
||
| # 3. OAuth U2M — browser authorization-code flow; the kernel runs it. | ||
| # 4. OAuth U2M — browser authorization-code flow; the kernel runs it. | ||
| # Only databricks-oauth reaches here (azure-oauth rejected up front). | ||
| # Forward the connector's own databricks-sql-python bundle instead of | ||
| # the kernel's databricks-sql-connector default, for parity with the | ||
|
|
@@ -283,7 +361,7 @@ def kernel_auth_kwargs( | |
| kwargs["identity_federation_client_id"] = federation_client_id | ||
| return kwargs | ||
|
|
||
| # 4. Custom credentials_provider — the connector's primary M2M path | ||
| # 5. Custom credentials_provider — the connector's primary M2M path | ||
| # on Thrift/SEA, but unusable on the kernel: it's an opaque token | ||
| # source with no extractable client_id/secret, so the kernel | ||
| # can't own the token lifecycle. Point the caller at the raw | ||
|
|
@@ -297,7 +375,7 @@ def kernel_auth_kwargs( | |
| "credentials_provider." | ||
| ) | ||
|
|
||
| # 5. Everything else (including no usable credentials at all — | ||
| # 6. Everything else (including no usable credentials at all — | ||
| # ``auth_provider`` is None on the kernel path when no access | ||
| # token was supplied and no OAuth kwargs resolved above). | ||
| provider_desc = ( | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🟡 Medium — Missing ambiguity guard for JWT M2M + U2M
auth_type.The bridge explicitly rejects shared-secret M2M colliding with a U2M request (
client_secret and auth_type == "databricks-oauth"→NotSupportedError), on the stated principle that conflicting auth signals must "fail loudly at session-open rather than silently resolving to one flow." The new JWT branch adds guards againstoauth_jwt_key_file+oauth_client_secretandoauth_jwt_key_file+credentials_provider, but there is no guard foroauth_jwt_key_file+auth_type="databricks-oauth".Concretely, a caller who passes
auth_type="databricks-oauth"(clear browser-U2M intent) while anoauth_jwt_key_fileis also present (e.g. leftover ambient config) silently gets routed tooauth-m2m-jwt— the browser flow they asked for never runs, and they authenticate as the service principal instead. This is exactly the failure mode theclient_secret+U2M guard was written to prevent, so the JWT path should mirror it. Consider adding:before the JWT branch, and a corresponding unit test.
(Anchored to the nearest changed line — see the description for the exact location.)