---
title: 'HMAC Middleware + Correlation Headers for Airflow Plugin'
slug: 'de-airflow-hmac-correlation-middleware'
created: '2026-03-24'
status: 'implementation-complete'
stepsCompleted: [1, 2, 3, 4]
tech_stack: ['FastAPI', 'Starlette BaseHTTPMiddleware', 'textiq_hmac', 'plain logging', 'SQLAlchemy (sync)', 'boto3/SeaweedFS']
files_to_modify: ['airflow/plugins/api/middleware.py (NEW)', 'airflow/plugins/api/upload.py', 'airflow/plugins/api/auth.py', 'tests/airflow/test_plugins/test_middleware.py (NEW)']
code_patterns: ['Separate RequestId + CorrelationId middleware', 'BaseHTTPMiddleware with request.state (ref: src/api/middleware.py:33)', 'Internal service variant (SPEC-CORR §3.1 adapted for plain logging)', 'Middleware stack order (SPEC-CORR §3.2)', 'DAG context propagation (SPEC-CORR §4.3)', 'JSONResponse error returns from middleware (ref: api/responses.py)']
test_patterns: ['pytest with TestClient(upload_app)', 'unittest.mock.patch for textiq_hmac', 'Test vectors from SPEC-CORR §7', 'No existing plugin tests — new test file needed']
---

# Tech-Spec: HMAC Middleware + Correlation Headers for Airflow Plugin

**Created:** 2026-03-24

## Overview

### Problem Statement

The Airflow upload plugin (`airflow/plugins/api/`) has HMAC auth called inline in the endpoint handler and no correlation headers support. This means:
- Cross-service request tracing is impossible (no X-Correlation-ID)
- X-Request-ID is not generated or propagated at the plugin level
- The inline auth pattern won't scale as new API endpoints are added
- DAG runs cannot be correlated back to the originating business transaction

### Solution

Add three FastAPI middleware classes to `upload_app`:
1. **RequestIdMiddleware** (outermost) — accepts client-provided X-Request-ID or generates new UUIDv4 if absent (per SPEC-CORR §2.2), echoes in response, stores on `request.state`
2. **CorrelationIdMiddleware** — forwards X-Correlation-ID from upstream (warns + auto-generates if missing), echoes in response, stores on `request.state`
3. **HMACVerificationMiddleware** — wraps existing `verify_hmac_headers()` logic, always enabled

Register in correct order: RequestId (outermost) → CorrelationId → HMAC → routes. Remove inline auth from upload endpoint. Propagate correlation_id + request_id to Airflow DAG run conf. Use plain `logging` (not structlog) throughout.

### Scope

**In Scope:**
- `RequestIdMiddleware`: accept client-provided X-Request-ID or generate new UUIDv4 if absent, echo in response, store on `request.state`
- `CorrelationIdMiddleware`: forward X-Correlation-ID from upstream, warn + auto-generate if missing, echo in response, store on `request.state`
- `HMACVerificationMiddleware`: always verify HMAC headers, no toggle
- Middleware registration order: RequestId (outermost) → CorrelationId → HMAC → routes
- Propagate `correlation_id` + `request_id` to DAG run conf
- Remove inline `verify_hmac_headers()` call from upload endpoint
- Plain `logging` with correlation_id/request_id in log messages
- Store correlation_id + request_id on `request.state` for handler access
- Tests for both middleware + DAG conf propagation

**Out of Scope:**
- API key fallback (deprecated)
- Structlog / contextvars (Airflow container uses plain logging)
- Audit log schema changes
- Other services (APA, DG, CB, IAM)
- Downstream HTTP client propagation helper (no outbound HTTP calls from plugin)
- Socket.IO propagation

## Context for Development

### Codebase Patterns

- Airflow plugin registers a plain `FastAPI()` sub-app via `AirflowPlugin.fastapi_apps` at `/api/v1/documents`
- Upload endpoint is sync (not async) — runs in threadpool for blocking boto3/SQLAlchemy
- Current auth: inline `verify_hmac_headers(request)` in upload handler returns `ctx` dict with `tenant_id`, `node_id`, `user_id`, `role`
- `textiq_hmac` package provides `load_key_ring()`, `verify_context_headers()`, `VerificationError`
- HMAC key ring loaded via `HMAC_KEY_RING` env var (handled by `textiq_hmac.load_key_ring()`)
- DAG trigger conf currently: `document_id, bucket, object_key, tenant_id, node_id`
- Plain `logging.getLogger(__name__)` used throughout plugin (not structlog)
- SPEC-CORR §3.1 internal variant is the reference pattern (adapted from structlog to plain logging)
- SPEC-CORR §4.3 defines DAG context propagation: include `correlation_id` + `request_id` in conf
- Correlation headers NEVER cause request rejection (graceful degradation per SPEC-CORR §5.3)

### Files to Reference

| File | Purpose |
| ---- | ------- |
| `airflow/plugins/api/upload.py` | Upload endpoint — remove inline auth, read from `request.state` |
| `airflow/plugins/api/auth.py` | Existing HMAC verification — reuse in middleware |
| `airflow/plugins/api/dag.py` | DAG trigger — receives conf dict (reference only, no changes needed) |
| `airflow/plugins/__init__.py` | Plugin registration — imports `upload_app`, has try/except for Airflow imports |
| `airflow/plugins/api/responses.py` | `error_response()` / `success_response()` — reuse in HMAC middleware for 401 |
| `textiq-blueprint/docs/planning-artifacts/2-planning/specs/request-correlation-headers-spec.md` | SPEC-CORR — authoritative spec |
| `textiq-blueprint/docs/planning-artifacts/2-planning/specs/hmac-context-headers-spec.md` | SPEC-HMAC — HMAC signing spec |
| `src/api/middleware.py` | Main app's RequestIdMiddleware — reference pattern (structlog variant) |

### Technical Decisions

- **Plain logging over structlog**: Airflow plugin runs in Airflow webserver process with its own Python environment. Structlog is not available and adding it would couple plugin to non-Airflow dependencies.
- **API key fallback dropped**: API key auth is deprecated. HMAC is the sole auth mechanism, always enabled — no toggle.
- **request.state for context passing**: Middleware stores `correlation_id`, `request_id`, and HMAC `ctx` on `request.state`. Handlers read from there — no global/threadlocal state.
- **Separate RequestId + CorrelationId middleware**: Intentional deviation from SPEC-CORR's single `CorrelationHeadersMiddleware` for testability and single-responsibility. Order dependency: RequestId must be registered before CorrelationId (both before HMAC).
- **Middleware on upload_app**: Middleware added directly to the `upload_app` FastAPI instance (not the parent Airflow app) so it only applies to our plugin routes.
- **Correlation headers never block**: Per SPEC-CORR, missing/invalid correlation headers are handled gracefully — auto-generate and warn, never reject.
- **New middleware.py file**: Three middleware classes in a new `airflow/plugins/api/middleware.py` file — keeps concerns separated from `auth.py` (verification logic) and `upload.py` (endpoint logic).
- **Sync endpoint, async middleware**: Upload handler is sync (threadpool). Starlette `BaseHTTPMiddleware.dispatch()` is async — this is compatible since Starlette wraps sync handlers automatically.
- **HMAC middleware returns JSONResponse on failure**: Uses `error_response()` from `api/responses.py` for consistent error format, short-circuits before reaching the handler.

## Implementation Plan

### Tasks

- [x] Task 1: Create `RequestIdMiddleware` in new `middleware.py`
  - File: `airflow/plugins/api/middleware.py` (NEW)
  - Action: Create `RequestIdMiddleware(BaseHTTPMiddleware)` that:
    - Accepts client-provided `X-Request-ID` if present, otherwise generates new UUIDv4: `request.headers.get("X-Request-ID") or str(uuid.uuid4())`
    - Stores on `request.state.request_id`
    - Echoes `X-Request-ID` in response headers
  - Notes: Matches SPEC-CORR §2.2 generation rules and `src/api/middleware.py:33-64` reference pattern. Uses plain `logging` instead of structlog.

- [x] Task 2: Create `CorrelationIdMiddleware` in `middleware.py`
  - File: `airflow/plugins/api/middleware.py`
  - Action: Create `CorrelationIdMiddleware(BaseHTTPMiddleware)` that:
    - Reads `X-Correlation-ID` from request headers
    - If present: use it (forwarded from upstream)
    - If missing: auto-generate UUIDv4 + log WARNING with path
    - Stores on `request.state.correlation_id`
    - Echoes `X-Correlation-ID` in response headers
  - Notes: Per SPEC-CORR §5.3 — never reject requests for missing/invalid correlation headers. Graceful degradation only.

- [x] Task 3: Create `HMACVerificationMiddleware` in `middleware.py`
  - File: `airflow/plugins/api/middleware.py`
  - Action: Create `HMACVerificationMiddleware(BaseHTTPMiddleware)` that:
    - Always calls `verify_hmac_headers(request)` from `api/auth.py`
    - On success: store result dict on `request.state.hmac_ctx`, set `request.state.auth_method = "hmac"`
    - On `VerificationError`: return `error_response("UNAUTHORIZED", str(e), status_code=401)` — short-circuit, don't call `call_next`
    - On any other `Exception` (e.g., `load_key_ring()` fails due to bad JSON, missing env var): log the error, return `error_response("INTERNAL_ERROR", "Authentication service unavailable", status_code=500)` — do not let unhandled exceptions propagate from middleware
    - Log auth method per request
  - Notes: Import `verify_hmac_headers` and `VerificationError` from `api.auth`. Import `error_response` from `api.responses`. Consider caching `load_key_ring()` result at module level rather than per-request (see `auth.py` line 40).

- [x] Task 4: Register middleware on `upload_app`
  - File: `airflow/plugins/api/upload.py`
  - Action: After `upload_app = FastAPI(...)`, register middleware in reverse order (Starlette executes last-added first):
    ```python
    upload_app.add_middleware(HMACVerificationMiddleware)
    upload_app.add_middleware(CorrelationIdMiddleware)
    upload_app.add_middleware(RequestIdMiddleware)
    ```
  - Notes: This gives execution order: RequestId → CorrelationId → HMAC → routes. Import from `api.middleware`.

- [x] Task 5: Remove inline auth from upload endpoint
  - File: `airflow/plugins/api/upload.py`
  - Action:
    - Remove the `try: ctx = verify_hmac_headers(request)` block (lines 62–65)
    - Replace with: `ctx = request.state.hmac_ctx` (set by middleware — always populated since HMAC is always enabled)
    - Remove the outer `except VerificationError` handler (lines 175–176) — middleware handles this now
    - Remove import of `verify_hmac_headers` and `VerificationError` from `api.auth`
  - Notes: `ctx` is guaranteed to be a dict (never `None`) because HMAC middleware always verifies and short-circuits on failure before reaching the handler.

- [x] Task 6: Add correlation_id + request_id to DAG conf
  - File: `airflow/plugins/api/upload.py`
  - Action: Update the `conf` dict passed to `trigger_dag()` to include:
    ```python
    conf = {
        "document_id": str(document_id),
        "bucket": bucket,
        "object_key": s3_key,
        "tenant_id": tenant_id,
        "node_id": node_id,
        "correlation_id": getattr(request.state, "correlation_id", None),
        "request_id": getattr(request.state, "request_id", None),
    }
    ```
  - Notes: Uses `getattr` with fallback for safety. Per SPEC-CORR §4.3.

- [x] Task 7: Add correlation_id + request_id to upload log messages
  - File: `airflow/plugins/api/upload.py`
  - Action: Update the success log line (line 169) to include both IDs:
    ```python
    log.info("Upload successful: document_id=%s, key=%s, correlation_id=%s, request_id=%s",
             document_id, s3_key,
             getattr(request.state, "correlation_id", None),
             getattr(request.state, "request_id", None))
    ```
  - Notes: Keep plain logging format. Add to error log paths too where request is available.

- [x] Task 8: Write tests for `RequestIdMiddleware`
  - File: `tests/airflow/test_plugins/test_middleware.py` (NEW)
  - Action: Test cases:
    - `test_request_id_generated`: request without X-Request-ID header → response has X-Request-ID with valid UUIDv4
    - `test_request_id_uses_client_value`: request WITH X-Request-ID header → response echoes the same value
    - `test_request_id_on_state`: after middleware, `request.state.request_id` is set and matches response header
  - Notes: Use minimal FastAPI app with single test route that returns `request.state.request_id`.

- [x] Task 9: Write tests for `CorrelationIdMiddleware`
  - File: `tests/airflow/test_plugins/test_middleware.py`
  - Action: Test cases (adapted from SPEC-CORR §7 test vectors):
    - `test_correlation_id_forwarded`: request with X-Correlation-ID → same value in response (forwarded)
    - `test_correlation_id_missing_generates_fallback`: request without X-Correlation-ID → response has auto-generated UUIDv4
    - `test_correlation_id_missing_logs_warning`: request without X-Correlation-ID → WARNING log emitted with path
    - `test_correlation_id_on_state`: after middleware, `request.state.correlation_id` is set
  - Notes: Use `caplog` fixture to assert WARNING log message.

- [x] Task 10: Write tests for `HMACVerificationMiddleware`
  - File: `tests/airflow/test_plugins/test_middleware.py`
  - Action: Test cases:
    - `test_hmac_valid_passes`: mock `verify_hmac_headers` returns ctx → 200, `request.state.hmac_ctx` populated, `auth_method = "hmac"`
    - `test_hmac_invalid_returns_401`: mock `verify_hmac_headers` raises `VerificationError` → 401 JSON error response, handler NOT called
    - `test_hmac_key_ring_error_returns_500`: mock `verify_hmac_headers` raises generic `Exception` (e.g., key ring misconfigured) → 500 JSON error response with `"INTERNAL_ERROR"`
  - Notes: Mock `api.auth.verify_hmac_headers`.

- [x] Task 11: Write integration test for full middleware stack
  - File: `tests/airflow/test_plugins/test_middleware.py`
  - Action: Test cases:
    - `test_middleware_stack_order`: request through full stack → response has both X-Request-ID and X-Correlation-ID, HMAC ctx on state
    - `test_middleware_hmac_failure_still_has_headers`: HMAC fails → 401 response still includes X-Request-ID and X-Correlation-ID (correlation middleware runs before HMAC)
  - Notes: Use `upload_app` with all three middleware registered. Mock HMAC internals.

### Acceptance Criteria

- [x] AC1: Given a request without X-Request-ID header, when processed by RequestIdMiddleware, then a new UUIDv4 is generated, echoed in the response X-Request-ID header, and stored on `request.state.request_id`.

- [x] AC2: Given a request with X-Request-ID header, when processed by RequestIdMiddleware, then the client-provided value is used (not replaced), echoed in the response, and stored on `request.state.request_id`.

- [x] AC3: Given a request with X-Correlation-ID header from upstream, when processed by CorrelationIdMiddleware, then the same value is forwarded in the response X-Correlation-ID header and stored on `request.state.correlation_id`.

- [x] AC4: Given a request WITHOUT X-Correlation-ID header, when processed by CorrelationIdMiddleware, then a fallback UUIDv4 is generated, a WARNING is logged with the request path, and the generated value is echoed in the response.

- [x] AC5: Given valid HMAC headers, when a request is processed by HMACVerificationMiddleware, then `request.state.hmac_ctx` contains `{tenant_id, node_id, user_id, role}` and `request.state.auth_method = "hmac"`.

- [x] AC6: Given invalid/missing HMAC headers, when a request is processed by HMACVerificationMiddleware, then it returns HTTP 401 with `{"error": {"code": "UNAUTHORIZED", "message": "..."}}` and the endpoint handler is NOT called.

- [x] AC7: Given HMAC verification fails (401), when the response is returned, then it still contains both X-Request-ID and X-Correlation-ID headers (correlation middleware runs before HMAC).

- [x] AC8: Given a successful upload, when the DAG is triggered, then `dag_run.conf` contains `correlation_id` and `request_id` fields alongside existing `document_id`, `bucket`, `object_key`, `tenant_id`, `node_id`.

- [x] AC9: Given a successful upload, when the response is logged, then the log message includes `correlation_id` and `request_id`.

## Additional Context

### Dependencies

- `textiq_hmac>=0.1.0` (already installed — provides `load_key_ring`, `verify_context_headers`, `VerificationError`)
- No new external dependencies required
- Depends on: TEXTIQ-181 sub-task 1 (textiq-hmac dependency) — already complete

### Testing Strategy

- **Framework:** pytest + `fastapi.testclient.TestClient`
- **New file:** `tests/airflow/test_plugins/test_middleware.py`
- **Approach:** Unit test each middleware in isolation using minimal FastAPI test apps, then integration test the full stack
- **Mocking:** `unittest.mock.patch` for `api.auth.verify_hmac_headers`
- **Test vectors:** Adapted from SPEC-CORR §7 (vectors 1–3 apply to internal service variant)
- **Coverage targets:**
  - RequestIdMiddleware: accepts client value or generates new, echoes in response, stores on request.state
  - CorrelationIdMiddleware: forwards from upstream, warns if missing, auto-generates fallback, echoes in response
  - HMACVerificationMiddleware: passes when HMAC valid, 401 when invalid
  - Integration: middleware order correct, request.state populated for handler, DAG conf includes both IDs, 401 still has correlation headers

### Notes

- **Jira:** TEXTIQ-181 — [HMAC] Story 1.3 — Dual-Auth + Correlation Headers on DE
- **Specs:** SPEC-CORR (request-correlation-headers-spec.md), SPEC-HMAC (hmac-context-headers-spec.md)
- **Scope boundary:** This covers the Airflow plugin side only. The main `src/` app (`RequestIdMiddleware` in `src/api/middleware.py`) needs separate extension for X-Correlation-ID.
- **Risk — BaseHTTPMiddleware with sync endpoints:** Starlette's `BaseHTTPMiddleware` wraps sync handlers in `run_in_threadpool`. This is the same pattern the upload endpoint already uses. No additional risk.
- **Future consideration:** When additional endpoints are added to the Airflow plugin, they automatically get all three middleware — no per-endpoint auth wiring needed.
