feat: Add AWS Cognito as OIDC Identity Provider (#325)

[TamarW0]

No description provided.

[vbelouso]

✅ Snyk checks have passed. No issues have been found so far.

Status	Scan Engine	Critical	High	Medium	Low	Total (0)
✅	Code Security	0	0	0	0	0 issues

💻 Catch issues earlier using the plugins for VS Code, JetBrains IDEs, Visual Studio, and Eclipse.

[zvigrinberg]

@TamarW0 LGTM Approved.
Don't forget tests in a follow up PR.

[tmihalac]

Code Review — 8 findings

Critical / High (must fix before merge)

1. Missing wildcard case in get_auth_header match statement — cve_http_output.py:308

The match handles "basic", "keycloak", "bearer", "cognito", and None, but has no case "disabled": and no case _: wildcard. A typo like AUTH_TYPE=Cognito (capital C) or AUTH_TYPE=coginto silently returns None with zero log output, causing all requests to go out unauthenticated.

Fix: Add explicit "disabled" case and a wildcard:

case "disabled" | None:
    return None
case _:
    logger.error(
        "Unrecognized auth_type '%s'. Valid values: basic, bearer, keycloak, cognito, disabled",
        http_config.auth_type,
    )
    return None

---

2. http_auth_header_context(None) silently falls back to SA token — credential_client.py:197

When Cognito token fetch fails, get_auth_header() returns None. Callers pass None into http_auth_header_context(None). Inside fetch_and_decrypt_credential, _http_auth_header_ctx.get() returns None — indistinguishable from "never set" — so it silently falls back to SA token auth. If the credential backend rejects SA tokens, the error says status=401 with no mention that Cognito was the intended auth method.

Fix: Log a warning at call sites when auth was configured but returned None:

auth_header = get_auth_header(http_output_config)
if auth_header is None and http_output_config.auth_type not in ("disabled", None):
    logger.warning(
        "Auth type is '%s' but get_auth_header returned None — "
        "credential backend calls will fall back to SA token",
        http_output_config.auth_type,
    )

---

3. cognito_domain field description shows bare domain, but code needs https:// — cve_http_output.py:103, 287

The field description says "Cognito domain (e.g. myapp.auth.us-east-1.amazoncognito.com)" — no scheme. But _fetch_cognito_token does f"{http_config.cognito_domain}/oauth2/token", so deployers following the docs will get requests.exceptions.MissingSchema. Compare with keycloak_server which documents "https://keycloak.example.com".

Fix: Either prepend https:// in code (safer, since Cognito always uses HTTPS):

domain = http_config.cognito_domain.rstrip("/")
if not domain.startswith("https://"):
    domain = f"https://{domain}"
token_url = f"{domain}/oauth2/token"

Or update the field description to require the scheme.

---

4. get_auth_header(None) crashes with AttributeError — cve_http_output.py:308-309

The type signature accepts None but the first line does match http_config.auth_type:. The 3 new callers all pass builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) which returns None if the config key is missing.

Fix:

def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None:
    if http_config is None:
        return None
    match http_config.auth_type:
        ...

---

High (should fix)

5. No warning when configured auth fails — cve_http_output.py:194-196

When get_auth_header returns None, the Authorization header is silently omitted. The backend returns 401, but the log says "Unable to send output response" with no indication that auth was configured and failed.

Fix: Add an else branch logging a warning when auth_type is not "disabled" but no header was obtained.

---

6. Broad except Exception in _fetch_cognito_token — cve_http_output.py:299-305

Catches ConnectionError, Timeout, SSLError, HTTPError, KeyError, JSONDecodeError identically. A KeyError('access_token') (wrong scope config) looks identical to a network timeout. At minimum, separate requests.RequestException from KeyError/ValueError.

---

Medium (recommended)

7. Token fetched 4× per pipeline run — cve_clone_and_deps.py:105, cve_generate_vdbs.py:224, cve_segmentation.py:228, cve_http_output.py:194

Each pipeline node independently calls get_auth_header() → _fetch_cognito_token(). Cognito M2M tokens have ~1hr TTL; one fetch could serve the entire run. Not a bug, but adds latency and could trigger rate limiting under load.

---

8. Zero test coverage for new code paths

_fetch_cognito_token, get_auth_header with "cognito", http_auth_header_context, and the None fallback path are all untested. Given the edge cases above (case sensitivity, None handling, missing scheme), tests would catch these.

[TamarW0]

/test vulnerability-analysis-on-pr

[tmihalac]

PR Review — AWS Cognito OIDC Provider (APPENG-5865)

Solid work on the Cognito flow itself — get_auth_header is nicely hardened (explicit None-config guard, a disabled | None case, and a case _ catch-all that logs unrecognized auth_type, with a test for the "Cognito" typo). _fetch_cognito_token has clean error handling: specific exception catches, no credential leakage in logs, a 30s timeout, correct Basic-auth client_credentials flow, and https/trailing-slash normalization. Test coverage for the Cognito path is thorough and behavior-focused. A few things worth addressing below.

Important

• The cve_http_output auth header is now reused to authenticate to the credential backend — please confirm this is intended. Previously fetch_and_decrypt_credential always authenticated to {backend_url}/api/v1/credentials/{id} via _resolve_jwt_token (K8s SA token / CLIENT_JWT_TOKEN). Now clone_and_deps, generate_vdbs, and cve_segmentation resolve the header from the cve_http_output config (the config for sending reports) and inject it via http_auth_header_context, overriding the SA-token path. These are almost certainly different services with different token audiences/IdPs — reusing the output token against the credential backend works only if both trust the same issuer/audience. With auth_type: basic the backend would get a Basic ... header and reject it (401/403). auth_type: disabled (default) yields None → falls back to the SA token, so default/local setups are unaffected, but any deployment with real auth configured silently changes behavior. See inline comment.

• builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) can raise and break three core stages. clone_and_deps, generate_vdbs, and cve_segmentation now hard-depend on a cve_http_output function being registered. If it isn't, get_function_config typically raises — breaking stages that previously had no such dependency. get_auth_header already tolerates None, so guard the lookup (try/except or a get-or-None variant) so a missing output config degrades gracefully instead of crashing the pipeline. See inline comment.

Suggestion

• Cognito token is fetched fresh on every stage, with no caching. Each of the three stages independently calls get_auth_header, so a cognito deployment makes 3 separate oauth2/token round-trips per scan. Consider resolving once and threading it through, or caching with expiry. Also note a long-running stage could outlive the token's expiry (default ~1h) since the header is captured at stage entry. Not blocking.

[tmihalac]

This override means the credential backend is now authenticated with whatever get_auth_header(cve_http_output_config) returns, instead of the K8s SA token / CLIENT_JWT_TOKEN from _resolve_jwt_token. The credential backend and the HTTP report-output endpoint are almost certainly different services with different token audiences — this only works if both trust the same issuer/audience. With auth_type: basic the backend would receive a Basic ... header and reject it (raises AuthenticationError on 401/403). Please confirm the credential backend actually accepts the cve_http_output token before merging. (disabled/default is unaffected — it falls back to the SA token here.)

[TamarW0]

Regarding the comment about reusing the cve_http_output token for the credential backend:

This pattern was introduced in commit 06d6368 ("fix: add generic idp authentication for credentials endpoint") for keycloak — the Cognito PR follows the same existing design. The credential backend is configured to accept the same IdP tokens as the HTTP output endpoint.

[tmihalac]

get_function_config typically raises if cve_http_output isn't registered, which would break this stage entirely — it previously had no dependency on the output config. Since get_auth_header already tolerates None, consider guarding this lookup (try/except or a get-or-None variant) so a missing output config degrades gracefully rather than crashing clone/VDB/segmentation. Same pattern applies in cve_generate_vdbs.py and cve_segmentation.py.

[tmihalac]

Automated review of the Cognito auth changes. The core OAuth2 client_credentials logic is correct, secure (no secret leakage), and well unit-tested. Flagging 5 findings inline — 2 Critical (silent auth downgrade + untested integration path) and 3 Important (duplication, over-broad except, redundant token fetch). Details on the relevant lines.

[tmihalac]

🔴 Critical — Cognito failure silently downgrades to unauthenticated, and the lost report is reported as success.

On any Cognito failure (bad secret, network error, non-200, missing access_token), _fetch_cognito_token returns None, so this branch returns None — indistinguishable from auth_type: disabled. Downstream in output_to_http._arun, the Authorization header is attached only if auth_header is not None, so the POST goes out unauthenticated. The backend 401s, request_with_retry raises after retries, and the pre-existing broad except Exception swallows it — _arun then returns message as success while the report was never delivered.

The three states "auth disabled", "auth required but failed", and "auth code bug" all collapse into one None.

Fix: when a non-disabled auth_type fails to produce a token, raise (or return a distinct sentinel) so the sender refuses to send unauthenticated; and don't let a delivery failure return success.

[tmihalac]

🔴 Critical — the entire integration half of this feature is untested.

This new context-var branch (and the http_auth_header_context manager above it) has no test. test_credential_client.py wasn't touched, and there are no tests for the try/except wiring added to cve_clone_and_deps.py / cve_generate_vdbs.py / cve_segmentation.py. The unit tests cover _fetch_cognito_token/get_auth_header well, but not the load-bearing wiring.

Please add (extend the untouched test_credential_client.py):

• uses the context header when set, and _resolve_jwt_token is not called;
• falls back to Bearer <jwt> when the context value is None;
• the ContextVar is reset after the with exits (concurrency-critical — this codebase runs parallel async tasks);
• per pipeline site: config present → header propagated; missing config raises → auth_header=None fallback.

Also verify the real exception type builder.get_function_config raises when the function isn't registered is actually within (KeyError, ValueError, AttributeError) — if not, the step crashes instead of falling back.

[tmihalac]

🟠 Important — this ~10-line auth-resolution block is copy-pasted into three files.

The identical block appears here, in cve_generate_vdbs.py, and in cve_segmentation.py. Extract a single helper (e.g. resolve_http_auth_header(builder) -> str | None in cve_http_output.py) so the fallback policy lives in one place — and can be unit-tested once rather than three times (see the coverage note on the credential client).

[tmihalac]

🟠 Important — this except is too broad, mislabels the cause, and logs too quietly.

• Catching AttributeError swallows genuine programming bugs (attr typo, config-shape refactor, a bug in get_auth_header) and silently downgrades to SA-token auth.
• The except wraps both get_function_config and get_auth_header; the latter already handles its own errors and returns None, so if it ever raises, the log "HTTP output config unavailable" is false (the config was available).
• Switching from an operator-configured Cognito/Keycloak identity to the SA token is a trust-boundary downgrade, but it's logged at INFO — routinely filtered in prod, so the degradation is invisible.

Fix: wrap only builder.get_function_config in a narrow except KeyError; call get_auth_header outside the try (it returns None on all internal failures); and log at WARNING/ERROR when a configured auth_type yields a None header.

[tmihalac]

🟠 Important — redundant token fetch per pipeline node + mid-operation expiry risk.

With auth_type: cognito, each of the three nodes (clone_and_deps, generate_vdbs, segmentation) independently calls this → a fresh requests.post (30s timeout) to /oauth2/token on every scan, and the token is discarded after that node's credential fetch. That's 3+ token round-trips per scan.

Also: the header is resolved once at the start of a minutes-long _arun and held in the ContextVar for the whole operation — a Cognito token (typical expires_in 3600s) can expire mid-operation with no refresh (observable via 401, but fragile).

Consider: caching the token respecting expires_in, or resolving it once per scan and threading it through.

[TamarW0]

Suggestion

• Cognito token is fetched fresh on every stage, with no caching. Each of the three stages independently calls get_auth_header, so a cognito deployment makes 3 separate oauth2/token round-trips per scan. Consider resolving once and threading it through, or caching with expiry. Also note a long-running stage could outlive the token's expiry (default ~1h) since the header is captured at stage entry. Not blocking.

Token caching is out of scope for this PR. Created a follow-up ticket to address both the redundant token fetches and the mid-operation expiry risk TC-6316

[TamarW0]

/test vulnerability-analysis-on-pr

[tmihalac]

The HTTP request is still inside try/except Exception, logs the error, then the else is skipped and _arun() still returns message regardless.

[tmihalac]

/test vulnerability-analysis-on-pr