diff --git a/kagenti-operator/go.mod b/kagenti-operator/go.mod index e99c769d..1a21b84d 100644 --- a/kagenti-operator/go.mod +++ b/kagenti-operator/go.mod @@ -133,14 +133,14 @@ require ( go.uber.org/zap v1.28.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/crypto v0.52.0 // indirect + golang.org/x/crypto v0.53.0 // indirect golang.org/x/exp v0.0.0-20251219203646-944ab1f22d93 // indirect golang.org/x/mod v0.36.0 // indirect golang.org/x/net v0.55.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect - golang.org/x/sys v0.45.0 // indirect - golang.org/x/term v0.43.0 // indirect - golang.org/x/text v0.37.0 // indirect + golang.org/x/sys v0.46.0 // indirect + golang.org/x/term v0.44.0 // indirect + golang.org/x/text v0.38.0 // indirect golang.org/x/time v0.15.0 // indirect golang.org/x/tools v0.45.0 // indirect gomodules.xyz/jsonpatch/v2 v2.5.0 // indirect diff --git a/kagenti-operator/go.sum b/kagenti-operator/go.sum index 45d1ec59..d2d00330 100644 --- a/kagenti-operator/go.sum +++ b/kagenti-operator/go.sum @@ -437,8 +437,8 @@ go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= -golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988= -golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc= +golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= +golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= golang.org/x/exp v0.0.0-20251219203646-944ab1f22d93 h1:fQsdNF2N+/YewlRZiricy4P1iimyPKZ/xwniHj8Q2a0= golang.org/x/exp v0.0.0-20251219203646-944ab1f22d93/go.mod h1:EPRbTFwzwjXj9NpYyyrvenVh9Y+GFeEvMNh7Xuz7xgU= golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= @@ -449,12 +449,12 @@ golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= -golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4= -golang.org/x/term v0.43.0/go.mod h1:lrhlHNdQJHO+1qVYiHfFKVuVioJIheAc3fBSMFYEIsk= -golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= -golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= +golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= +golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= +golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= diff --git a/kagenti-operator/hack/kind-reload-all.sh b/kagenti-operator/hack/kind-reload-all.sh index 8301ec33..5b416a29 100755 --- a/kagenti-operator/hack/kind-reload-all.sh +++ b/kagenti-operator/hack/kind-reload-all.sh @@ -5,6 +5,7 @@ # Services: # 1. kagenti-controller-manager (operator) # 2. bundle-service (OPA bundle server) +# 3. token-broker (OAuth token management) # # Usage: # ./hack/kind-reload-all.sh [kind-cluster-name] [namespace] @@ -25,6 +26,7 @@ IMAGE_TAG="$(git -C "${ROOT_DIR}" rev-parse --short HEAD)" OPERATOR_IMG="localhost/kagenti-operator:${IMAGE_TAG}" BUNDLE_IMG="localhost/bundle-service:${IMAGE_TAG}" +TOKEN_BROKER_IMG="localhost/token-broker:${IMAGE_TAG}" echo "============================================" echo " Building and loading to Kind: ${CLUSTER}" @@ -52,16 +54,31 @@ echo "" echo "==> Building kagenti-operator image" ${CONTAINER_TOOL} build -t "${OPERATOR_IMG}" -f "${ROOT_DIR}/Dockerfile" "${ROOT_DIR}" -echo "" -echo "==> Building bundle-service image" -${CONTAINER_TOOL} build -t "${BUNDLE_IMG}" -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" "${ROOT_DIR}" +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + echo "" + echo "==> Building bundle-service image" + ${CONTAINER_TOOL} build -t "${BUNDLE_IMG}" -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" "${ROOT_DIR}" +fi + +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + echo "" + echo "==> Building token-broker image" + ${CONTAINER_TOOL} build -t "${TOKEN_BROKER_IMG}" -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" "${ROOT_DIR}" +fi # --- Load into Kind --- echo "" echo "==> Loading images into Kind cluster '${CLUSTER}'" kind load docker-image "${OPERATOR_IMG}" --name "${CLUSTER}" -kind load docker-image "${BUNDLE_IMG}" --name "${CLUSTER}" + +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + kind load docker-image "${BUNDLE_IMG}" --name "${CLUSTER}" +fi + +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + kind load docker-image "${TOKEN_BROKER_IMG}" --name "${CLUSTER}" +fi # --- Deploy kagenti-controller-manager --- @@ -76,8 +93,9 @@ else echo " (the controller-manager requires webhook certs, RBAC, and CRDs from kustomize)" fi -# --- Deploy bundle-service --- +# --- Deploy bundle-service (if source exists) --- +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then echo "" echo "==> Deploying bundle-service" @@ -210,14 +228,71 @@ spec: targetPort: http protocol: TCP EOF +fi + +# --- Deploy token-broker (if source exists) --- + +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then +echo "" +echo "==> Deploying token-broker" + +# Load OAuth credentials from .env +if [ -f "${ROOT_DIR}/.env" ]; then + export $(grep -v '^#' "${ROOT_DIR}/.env" | grep -v '^\s*$' | xargs) +fi + +if [ -z "${GITHUB_OAUTH_CLIENT_ID:-}" ] || [ -z "${GITHUB_OAUTH_CLIENT_SECRET:-}" ]; then + echo " ERROR: OAuth credentials not set." + echo " Create ${ROOT_DIR}/.env with:" + echo " GITHUB_OAUTH_CLIENT_ID=" + echo " GITHUB_OAUTH_CLIENT_SECRET=" + echo " See .env.example for reference." + exit 1 +fi + +# Create/update OAuth secret +kubectl delete secret github-oauth-credentials -n "${NAMESPACE}" 2>/dev/null || true +kubectl create secret generic github-oauth-credentials \ + --from-literal=client-id="${GITHUB_OAUTH_CLIENT_ID}" \ + --from-literal=client-secret="${GITHUB_OAUTH_CLIENT_SECRET}" \ + --namespace="${NAMESPACE}" +echo " OAuth secret created" + +# Apply the deployment manifests (single source of truth). +# These pin namespace kagenti-system and include the OAuth-callback HTTPRoute +# (attached to the shared "http" gateway). The secret is created above from +# .env, so the secret example manifest is intentionally skipped. +DEPLOY_DIR="${ROOT_DIR}/cmd/token-broker/deploy" +kubectl apply \ + -f "${DEPLOY_DIR}/00-serviceaccount.yaml" \ + -f "${DEPLOY_DIR}/02-deployment.yaml" \ + -f "${DEPLOY_DIR}/03-service.yaml" \ + -f "${DEPLOY_DIR}/04-httproute.yaml" + +# Inject the freshly built, git-tagged image (the manifest carries a placeholder +# tag with imagePullPolicy: IfNotPresent; kind-loaded images need Never). +kubectl set image deployment/token-broker \ + token-broker="${TOKEN_BROKER_IMG}" \ + -n "${NAMESPACE}" +kubectl patch deployment/token-broker -n "${NAMESPACE}" --type=json \ + -p='[{"op":"replace","path":"/spec/template/spec/containers/0/imagePullPolicy","value":"Never"}]' +echo " token-broker manifests applied (image ${TOKEN_BROKER_IMG})" +fi # --- Delete pods to pick up the new images --- echo "" echo "==> Deleting pods to pick up new images" -kubectl delete pods -n "${NAMESPACE}" -l app=bundle-service --wait=false kubectl delete pods -n "${NAMESPACE}" -l control-plane=controller-manager --wait=false 2>/dev/null || true +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + kubectl delete pods -n "${NAMESPACE}" -l app=bundle-service --wait=false +fi + +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + kubectl delete pods -n "${NAMESPACE}" -l app=token-broker --wait=false +fi + # --- Wait for rollouts --- echo "" @@ -225,8 +300,15 @@ echo "==> Waiting for kagenti-controller-manager rollout" kubectl rollout status deployment/kagenti-controller-manager -n "${NAMESPACE}" --timeout=120s 2>/dev/null || \ echo " WARNING: kagenti-controller-manager rollout did not complete" -echo "==> Waiting for bundle-service rollout" -kubectl rollout status deployment/bundle-service -n "${NAMESPACE}" --timeout=60s +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + echo "==> Waiting for bundle-service rollout" + kubectl rollout status deployment/bundle-service -n "${NAMESPACE}" --timeout=60s +fi + +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + echo "==> Waiting for token-broker rollout" + kubectl rollout status deployment/token-broker -n "${NAMESPACE}" --timeout=60s +fi # --- Summary --- @@ -236,12 +318,35 @@ echo " Done!" echo "" echo " Images loaded:" echo " ${OPERATOR_IMG}" -echo " ${BUNDLE_IMG}" +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + echo " ${BUNDLE_IMG}" +fi +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + echo " ${TOKEN_BROKER_IMG}" +fi echo "" echo " Namespace: ${NAMESPACE}" echo " - kagenti-controller-manager" -echo " - bundle-service" +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + echo " - bundle-service" +fi +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + echo " - token-broker" +fi echo "" -echo " bundle-service URL: http://bundle-service.${NAMESPACE}.svc.cluster.local:8080" -echo " To port-forward: kubectl port-forward -n ${NAMESPACE} svc/bundle-service 8080:8080" +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + echo " bundle-service URL: http://bundle-service.${NAMESPACE}.svc.cluster.local:8080" +fi +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + echo " token-broker URL: http://token-broker.${NAMESPACE}.svc.cluster.local:8190" + echo " OAuth callback: http://token-broker.localtest.me:8080/oauth/callback (via 'http' gateway)" +fi +echo "" +echo " To port-forward:" +if [ -f "${ROOT_DIR}/cmd/bundle-service/Dockerfile" ]; then + echo " kubectl port-forward -n ${NAMESPACE} svc/bundle-service 8080:8080" +fi +if [ -f "${ROOT_DIR}/cmd/token-broker/Dockerfile" ]; then + echo " kubectl port-forward -n ${NAMESPACE} svc/token-broker 8190:8190" +fi echo "============================================" diff --git a/kagenti-operator/token-broker b/kagenti-operator/token-broker new file mode 100755 index 00000000..d8263978 Binary files /dev/null and b/kagenti-operator/token-broker differ diff --git a/token-broker/Dockerfile b/token-broker/Dockerfile new file mode 100644 index 00000000..fd6ea413 --- /dev/null +++ b/token-broker/Dockerfile @@ -0,0 +1,21 @@ +FROM --platform=$BUILDPLATFORM docker.io/golang:1.26 AS builder +ARG TARGETOS +ARG TARGETARCH + +WORKDIR /workspace +COPY go.mod go.mod +COPY go.sum go.sum +RUN go mod download + +COPY cmd/ cmd/ +COPY internal/ internal/ +COPY pkg/ pkg/ + +RUN CGO_ENABLED=0 GOOS=${TARGETOS:-linux} GOARCH=${TARGETARCH} go build -o token-broker ./cmd/ + +FROM gcr.io/distroless/static:nonroot +WORKDIR / +COPY --from=builder /workspace/token-broker . +USER 65532:65532 + +ENTRYPOINT ["/token-broker"] diff --git a/token-broker/README.md b/token-broker/README.md new file mode 100644 index 00000000..fc3e2a6a --- /dev/null +++ b/token-broker/README.md @@ -0,0 +1,418 @@ +# Token Broker Service + +The Token Broker is a standalone Kagenti service that enables **HITL (Human-in-the-Loop) Authorization** — allowing agents to obtain additional user permissions at runtime, beyond those provided when the task was submitted. + +## Overview + +Kagenti supports three route types for outbound agent requests: + +- **Passthrough** — forward the token as-is +- **Token Exchange** — use a subset of the permissions already in the agent's token +- **Token Broker** — obtain additional permissions with the user's help (this service) + +The Token Broker handles the third case. When an agent needs to access an external or internal resource that requires permissions not already present in its token, the Token Broker brings the user into the loop via an OAuth 2.0 PKCE flow — obtaining just-in-time, user-scoped credentials. Neither the agent nor Kagenti ever holds the user's credentials; the Token Broker only caches the resulting access token and controls which agents (those in the same user session) may use it. + +The Token Broker: +- **Acts as the OAuth client** with OAuth credentials (client_id, client_secret) +- **Generates PKCE** (code_verifier, code_challenge) for each OAuth flow +- **Exchanges tokens directly** with OAuth provider (not through the resource server) +- Manages OAuth sessions with timeout and lifecycle management +- Caches tokens per (session, resource_url) with JWT expiry parsing +- Discovers OAuth endpoints and scopes from resource servers' `.well-known` metadata +- Coordinates OAuth flows with the Backend via long-polling events + +Resources that can be brokered include MCP servers, LLM APIs, direct REST APIs, and other agents — anything that accepts an OAuth bearer token. + +## Architecture + +``` +┌─────────────┐ ┌──────────────┐ ┌──────────────────┐ +│ Backend │◄────────┤ Token Broker ├────────►│ Resource Server │ +│ │ Events │ │ Discovery│ (MCP/LLM/API/..) │ +└─────────────┘ └──────┬───────┘ └──────────────────┘ + │ + │ Token Cache + │ Session Store + ▼ + ┌──────────────┐ + │ Storage │ + │ (In-Memory) │ + └──────────────┘ +``` + +## Building + +```bash +go build -o token-broker ./cmd/ +``` + +## Configuration + +### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `TOKEN_BROKER_PORT` | `8190` | HTTP server port | +| `OAUTH_CLIENT_ID` | **(required)** | OAuth client ID for the OAuth App | +| `OAUTH_CLIENT_SECRET` | **(required)** | OAuth client secret for the OAuth App | +| `OAUTH_CALLBACK_URL` | `http://localhost:8190/oauth/callback` | URL the OAuth provider redirects to after user consent | +| `TOKEN_BROKER_SESSION_TIMEOUT` | `60s` | Idle session timeout | +| `TOKEN_BROKER_MAX_SESSIONS_PER_USER` | `5` | Max concurrent sessions per user | +| `TOKEN_BROKER_TOKEN_WAIT_TIMEOUT` | `300s` | Max time to wait for OAuth completion | +| `ALLOWED_REDIRECT_HOSTS` | *(none — all permitted)* | Comma-separated permitted hostnames for `backend_session_redirect_url`; set in production | +| `JWT_JWKS_URL` | *(none — validation disabled)* | JWKS endpoint for JWT signature validation; set in production | +| `JWT_ISSUER` | *(none)* | Expected JWT issuer; required when `JWT_JWKS_URL` is set | +| `JWT_AUDIENCE` | *(none)* | Comma-separated expected audiences; optional | +| `OAUTH_AUTHORIZATION_ENDPOINT` | *(discovered)* | Override OAuth authorization endpoint — skips `.well-known` discovery when set together with `OAUTH_TOKEN_ENDPOINT` | +| `OAUTH_TOKEN_ENDPOINT` | *(discovered)* | Override OAuth token endpoint — skips `.well-known` discovery when set together with `OAUTH_AUTHORIZATION_ENDPOINT` | +| `OAUTH_SCOPES_SUPPORTED` | *(discovered)* | Comma-separated global scope list, used only when discovery is skipped | +| `RESOURCE_CONFIG` | *(none)* | Per-resource OAuth config (scopes + endpoint overrides) — see below | +| `RESOURCE_SCOPES` | *(none)* | Per-resource scope override, legacy format — see below | + +### Security-critical variables + +**`JWT_JWKS_URL`** — when unset, JWT signatures are **not validated** and any caller can forge a token. Always set this in production. The Token Broker logs a WARNING at startup when it is unset. + +**`ALLOWED_REDIRECT_HOSTS`** — when unset, `backend_session_redirect_url` accepts any host. Always set this to the hostname(s) of your backend UI in production. The Token Broker logs a WARNING at startup when it is unset. + +#### Per-resource OAuth configuration (`RESOURCE_CONFIG`) + +`RESOURCE_CONFIG` is a JSON map of resource URL to per-resource OAuth overrides. It takes precedence over `RESOURCE_SCOPES` and over `.well-known` discovery for matching resources. + +Each entry supports: +- `scopes` — list of OAuth scopes to request (replaces discovered scopes) +- `authorization_endpoint` — override the authorization endpoint +- `token_endpoint` — override the token endpoint + +```bash +export RESOURCE_CONFIG='{ + "http://mcp-server.example.com": {"scopes": ["read:user", "user:email"]}, + "http://llm-api.example.com": { + "scopes": ["model:read"], + "authorization_endpoint": "https://idp.example.com/oauth/authorize", + "token_endpoint": "https://idp.example.com/oauth/token" + } +}' +``` + +The resource URL must match the value in the `X-Server-Url` header exactly (no trailing-slash normalization). + +#### Per-resource scope override, legacy (`RESOURCE_SCOPES`) + +Deprecated in favor of `RESOURCE_CONFIG`. Overrides scopes only, no endpoint overrides. + +Format: comma-separated `= ` entries. + +```bash +export RESOURCE_SCOPES="\ + https://mcp.example.com=read:user user:email,\ + https://mcp2.example.com=read:user user:email" +``` + +### Example Configuration + +```bash +export TOKEN_BROKER_PORT=8190 +export OAUTH_CLIENT_ID=Ov23liXXXXXXXXXXXXXX +export OAUTH_CLIENT_SECRET=your_oauth_app_secret +export OAUTH_CALLBACK_URL=https://token-broker.example.com/oauth/callback +export ALLOWED_REDIRECT_HOSTS=app.example.com +export JWT_JWKS_URL=https://keycloak.example.com/realms/kagenti/protocol/openid-connect/certs +export JWT_ISSUER=https://keycloak.example.com/realms/kagenti +export TOKEN_BROKER_SESSION_TIMEOUT=60s +export TOKEN_BROKER_MAX_SESSIONS_PER_USER=5 +export TOKEN_BROKER_TOKEN_WAIT_TIMEOUT=300s + +# Optional: per-resource OAuth config (scopes + endpoint overrides) +export RESOURCE_CONFIG='{"http://mcp.example.com": {"scopes": ["read:user", "user:email"]}}' +``` + +## Running + +### Local Development + +```bash +# Set required configuration +export OAUTH_CLIENT_ID=Ov23liXXXXXXXXXXXXXX +export OAUTH_CLIENT_SECRET=your_secret + +# Run the service +./token-broker +``` + +The service will start on port 8190 by default and log to stdout in JSON format. + +### Docker + +```bash +docker build -t token-broker . +docker run -p 8190:8190 \ + -e OAUTH_CLIENT_ID=Ov23liXXXXXXXXXXXXXX \ + -e OAUTH_CLIENT_SECRET=your_secret \ + token-broker +``` + +## API Endpoints + +### POST /sessions + +Create a new OAuth session. + +**Headers:** +- `Authorization`: Bearer JWT token (required) + - JWT claims: `sub` (user_id), `session_uid` (session_key) or `jti` (session_key, fallback) + +**Body:** +```json +{ "backend_session_redirect_url": "https://app.example.com/oauth-complete" } +``` + +**Response:** `201 Created` (empty body) + +**Errors:** +- `400` - Missing/invalid request or redirect URL host not in `ALLOWED_REDIRECT_HOSTS` +- `401` - Invalid JWT +- `429` - Max sessions per user exceeded + +--- + +### POST /sessions/token + +Request a token for a resource (MCP server, LLM API, etc.). Blocks until token is available or timeout. + +**Headers:** +- `Authorization`: Bearer JWT token (required) +- `X-Server-Url`: resource base URL (required) + +**Response:** `200 OK` +```json +{ "token": "access-token" } +``` + +**Errors:** +- `401` - Invalid JWT or session +- `408` - Timeout waiting for OAuth completion +- `503` - OAuth flow failed + +--- + +### POST /sessions/broker-events + +Long-poll for OAuth events from the Token Broker. + +**Headers:** +- `Authorization`: Bearer JWT token (required) + +**Response:** `200 OK` + +Authorization URL ready (Backend should open OAuth consent window): +```json +{ + "type": "oauth_url_ready", + "auth_url": "https://github.com/login/oauth/authorize?client_id=...&code_challenge=...&state=..." +} +``` + +Or error event: +```json +{ + "type": "error", + "message": "Session expired", + "code": "session_expired" +} +``` + +--- + +### POST /sessions/end + +Terminate a session and release all resources. + +**Headers:** +- `Authorization`: Bearer JWT token (required) + +**Response:** `200 OK` + +--- + +### GET /oauth/callback + +OAuth provider callback endpoint. Receives the authorization code from the OAuth provider after the user consents. + +**Query Parameters:** +- `code`: Authorization code (required) +- `state`: OAuth state parameter (required, format: `sessionKey.nonce`) + +**Response:** `302 Found` — redirects user to `backend_session_redirect_url` + +**Note:** Called by the OAuth provider, not by the Backend or AuthBridge. + +--- + +## Testing + +### Manual Testing with curl + +1. **Generate session key and JWT token:** +```bash +SESSION_KEY=$(uuidgen | tr '[:upper:]' '[:lower:]') +USER_ID="test-user" +JWT_HEADER='{"alg":"none","typ":"JWT"}' +JWT_PAYLOAD="{\"sub\":\"$USER_ID\",\"session_uid\":\"$SESSION_KEY\"}" +JWT_TOKEN=$(echo -n "$JWT_HEADER" | base64 | tr -d '=' | tr '/+' '_-').$(echo -n "$JWT_PAYLOAD" | base64 | tr -d '=' | tr '/+' '_-'). +``` + +2. **Create a session:** +```bash +curl -X POST http://localhost:8190/sessions \ + -H "Authorization: Bearer $JWT_TOKEN" \ + -H "Content-Type: application/json" \ + -d '{"backend_session_redirect_url": "http://localhost:3000/oauth-complete"}' +``` + +3. **Start event long-poll (in another terminal):** +```bash +curl -X POST "http://localhost:8190/sessions/broker-events" \ + -H "Authorization: Bearer $JWT_TOKEN" +``` + +4. **Request a token (in another terminal):** +```bash +curl -X POST "http://localhost:8190/sessions/token" \ + -H "Authorization: Bearer $JWT_TOKEN" \ + -H "X-Server-Url: http://localhost:8082" +``` + +5. **End session:** +```bash +curl -X POST "http://localhost:8190/sessions/end" \ + -H "Authorization: Bearer $JWT_TOKEN" +``` + +## Logging + +The Token Broker logs to stdout in JSON format for Kubernetes compatibility. + +**Log Levels:** +- `INFO` - Normal operations (session creation, token acquisition, OAuth flows) +- `DEBUG` - Detailed flow information (full auth URLs, cache hits, semaphore acquisition) +- `WARN` - Recoverable issues and misconfiguration (missing JWT_JWKS_URL, missing ALLOWED_REDIRECT_HOSTS) +- `ERROR` - Errors requiring attention (OAuth failures, internal errors) + +## Session Lifecycle + +``` +1. Backend creates session → Session created (active) +2. Backend polls /broker-events → Timer reset +3. Backend disconnects → Timer starts (60s) +4. Backend reconnects → Timer reset +5. Timer expires OR Backend calls /end → Session terminated +``` + +## Token Caching + +Tokens are cached per `(session_key, resource_url)`: +- JWT tokens: Expiry parsed from `exp` claim +- Non-JWT tokens: Treated as long-lived (1 year) +- Near-expiry tokens (< 5 min): Treated as expired +- Only agents in the same user session may use a cached token + +## Concurrency Model + +- **Per-session semaphore**: Only one OAuth flow per session at a time +- **Double-checked locking**: Check cache before and after semaphore acquisition +- **Session isolation**: Different sessions run OAuth flows independently + +## Error Handling + +Two error response formats are used depending on the caller: + +**AuthBridge API** (`POST /sessions/token`) — flat format: +```json +{ "code": "error_code", "message": "Human-readable message" } +``` + +**Backend API** (`POST /sessions`, `/sessions/broker-events`, `/sessions/end`) — nested format: +```json +{ "error": { "code": "error_code", "message": "Human-readable message" } } +``` + +**Error Codes:** +- `invalid_request` - Missing or invalid parameters +- `unauthorized` - Invalid session or user mismatch +- `session_not_found` - Session does not exist +- `session_expired` - Session timed out +- `too_many_sessions` - User exceeded max sessions +- `timeout` - OAuth flow did not complete in time +- `oauth_failed` - OAuth discovery or exchange failed +- `internal_error` - Internal server error + +## Security Considerations + +- Session keys are UUID v4 (cryptographically random) +- User ID validation on every request; session ownership strictly enforced +- **JWT signature validation** via JWKS when `JWT_JWKS_URL` is configured (recommended for production) +- **Redirect host validation** via `ALLOWED_REDIRECT_HOSTS` (recommended for production) +- **Token Broker acts as OAuth client** with client credentials +- **PKCE generated by Token Broker** (code_verifier never leaves Token Broker) +- **Token exchange happens directly** between Token Broker and OAuth provider +- `code_verifier` never exposed to Backend or resource server +- `client_secret` held only by Token Broker (not the resource server) +- Tokens stored in memory only (no persistence) +- Only agents belonging to the same user session may use a cached token + +## Troubleshooting + +### Session expires too quickly +Increase `TOKEN_BROKER_SESSION_TIMEOUT`: +```bash +export TOKEN_BROKER_SESSION_TIMEOUT=120s +``` + +### OAuth flow times out +Increase `TOKEN_BROKER_TOKEN_WAIT_TIMEOUT`: +```bash +export TOKEN_BROKER_TOKEN_WAIT_TIMEOUT=600s +``` + +### Too many sessions error +Increase `TOKEN_BROKER_MAX_SESSIONS_PER_USER`: +```bash +export TOKEN_BROKER_MAX_SESSIONS_PER_USER=10 +``` + +### Token not cached +Check logs for JWT parsing errors. Non-JWT tokens are treated as long-lived. + +### Redirect URL rejected (400) +The `backend_session_redirect_url` host is not in `ALLOWED_REDIRECT_HOSTS`. Add the backend hostname to the env var. + +## Development + +### Project Structure +``` +cmd/ + main.go # Service bootstrap and configuration + +internal/ + api/ # HTTP handlers + cache/ # Token cache with JWT expiry parsing + core/ # Token acquisition orchestration and interfaces + oauth/ # OAuth client (PKCE, URL building, token exchange) and discovery + session/ # Session lifecycle, semaphore, event coordination + auth/ # JWT claim extraction + +pkg/ + oauth/ # PKCE challenge generation (GeneratePKCEChallenge, GenerateState) +``` + +### Adding New Features + +1. Update interfaces in `internal/core/interfaces.go` +2. Implement in the appropriate package +3. Add unit tests +4. Update this README + +## License + +See LICENSE file in repository root. diff --git a/token-broker/cmd/config_test.go b/token-broker/cmd/config_test.go new file mode 100644 index 00000000..64da324e --- /dev/null +++ b/token-broker/cmd/config_test.go @@ -0,0 +1,204 @@ +/* +Copyright 2025. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "os" + "testing" + "time" +) + +func TestLoadConfig_Defaults(t *testing.T) { + clearEnv() + + cfg := LoadConfig() + + if cfg.ListenPort != 8190 { + t.Errorf("Expected default ListenPort 8190, got %d", cfg.ListenPort) + } + + if cfg.SessionTimeout != 60*time.Second { + t.Errorf("Expected default SessionTimeout 60s, got %v", cfg.SessionTimeout) + } + + if cfg.MaxSessionsPerUser != 5 { + t.Errorf("Expected default MaxSessionsPerUser 5, got %d", cfg.MaxSessionsPerUser) + } + + if cfg.TokenWaitTimeout != 300*time.Second { + t.Errorf("Expected default TokenWaitTimeout 300s, got %v", cfg.TokenWaitTimeout) + } + + if cfg.AuthorizationEndpoint != "" { + t.Errorf("Expected empty AuthorizationEndpoint, got %s", cfg.AuthorizationEndpoint) + } + + if cfg.TokenEndpoint != "" { + t.Errorf("Expected empty TokenEndpoint, got %s", cfg.TokenEndpoint) + } + + if len(cfg.ScopesSupported) != 0 { + t.Errorf("Expected empty ScopesSupported, got %v", cfg.ScopesSupported) + } +} + +func TestLoadConfig_OAuthMetadata(t *testing.T) { + clearEnv() + + _ = os.Setenv("OAUTH_CLIENT_ID", "test-client-id") + _ = os.Setenv("OAUTH_CLIENT_SECRET", "test-client-secret") + _ = os.Setenv("OAUTH_CALLBACK_URL", "http://localhost:8190/oauth/callback") + _ = os.Setenv("OAUTH_AUTHORIZATION_ENDPOINT", "https://github.com/login/oauth/authorize") + _ = os.Setenv("OAUTH_TOKEN_ENDPOINT", "https://github.com/login/oauth/access_token") + _ = os.Setenv("OAUTH_SCOPES_SUPPORTED", "repo,user,read:org") + defer clearEnv() + + cfg := LoadConfig() + + if cfg.ClientID != "test-client-id" { + t.Errorf("Expected ClientID 'test-client-id', got %s", cfg.ClientID) + } + + if cfg.ClientSecret != "test-client-secret" { + t.Errorf("Expected ClientSecret 'test-client-secret', got %s", cfg.ClientSecret) + } + + if cfg.AuthorizationEndpoint != "https://github.com/login/oauth/authorize" { + t.Errorf("Expected AuthorizationEndpoint, got %s", cfg.AuthorizationEndpoint) + } + + if cfg.TokenEndpoint != "https://github.com/login/oauth/access_token" { + t.Errorf("Expected TokenEndpoint, got %s", cfg.TokenEndpoint) + } + + expectedScopes := []string{"repo", "user", "read:org"} + if len(cfg.ScopesSupported) != len(expectedScopes) { + t.Fatalf("Expected %d scopes, got %d", len(expectedScopes), len(cfg.ScopesSupported)) + } + + for i, scope := range expectedScopes { + if cfg.ScopesSupported[i] != scope { + t.Errorf("Expected scope[%d] '%s', got '%s'", i, scope, cfg.ScopesSupported[i]) + } + } +} + +func TestLoadConfig_ScopesWithSpaces(t *testing.T) { + clearEnv() + + _ = os.Setenv("OAUTH_CLIENT_ID", "id") + _ = os.Setenv("OAUTH_CLIENT_SECRET", "secret") + _ = os.Setenv("OAUTH_SCOPES_SUPPORTED", " repo , user , read:org ") + defer clearEnv() + + cfg := LoadConfig() + + expectedScopes := []string{"repo", "user", "read:org"} + if len(cfg.ScopesSupported) != len(expectedScopes) { + t.Fatalf("Expected %d scopes, got %d", len(expectedScopes), len(cfg.ScopesSupported)) + } + + for i, scope := range expectedScopes { + if cfg.ScopesSupported[i] != scope { + t.Errorf("Expected scope[%d] '%s', got '%s'", i, scope, cfg.ScopesSupported[i]) + } + } +} + +func TestValidate(t *testing.T) { + tests := []struct { + name string + config *Config + expectError bool + }{ + { + name: "valid config", + config: &Config{ + ListenPort: 8190, + ClientID: "test-client-id", + ClientSecret: "test-client-secret", + CallbackURL: "http://localhost:8190/oauth/callback", + SessionTimeout: 60 * time.Second, + MaxSessionsPerUser: 5, + TokenWaitTimeout: 300 * time.Second, + }, + expectError: false, + }, + { + name: "missing client ID", + config: &Config{ + ListenPort: 8190, + ClientSecret: "test-client-secret", + CallbackURL: "http://localhost:8190/oauth/callback", + SessionTimeout: 60 * time.Second, + MaxSessionsPerUser: 5, + TokenWaitTimeout: 300 * time.Second, + }, + expectError: true, + }, + { + name: "missing client secret", + config: &Config{ + ListenPort: 8190, + ClientID: "test-client-id", + CallbackURL: "http://localhost:8190/oauth/callback", + SessionTimeout: 60 * time.Second, + MaxSessionsPerUser: 5, + TokenWaitTimeout: 300 * time.Second, + }, + expectError: true, + }, + { + name: "invalid port", + config: &Config{ + ListenPort: 0, + ClientID: "test-client-id", + ClientSecret: "test-client-secret", + CallbackURL: "http://localhost:8190/oauth/callback", + SessionTimeout: 60 * time.Second, + MaxSessionsPerUser: 5, + TokenWaitTimeout: 300 * time.Second, + }, + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.config.Validate() + if tt.expectError && err == nil { + t.Error("Expected error but got none") + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error but got: %v", err) + } + }) + } +} + +func clearEnv() { + _ = os.Unsetenv("TOKEN_BROKER_PORT") + _ = os.Unsetenv("OAUTH_CLIENT_ID") + _ = os.Unsetenv("OAUTH_CLIENT_SECRET") + _ = os.Unsetenv("OAUTH_CALLBACK_URL") + _ = os.Unsetenv("TOKEN_BROKER_SESSION_TIMEOUT") + _ = os.Unsetenv("TOKEN_BROKER_MAX_SESSIONS_PER_USER") + _ = os.Unsetenv("TOKEN_BROKER_TOKEN_WAIT_TIMEOUT") + _ = os.Unsetenv("OAUTH_AUTHORIZATION_ENDPOINT") + _ = os.Unsetenv("OAUTH_TOKEN_ENDPOINT") + _ = os.Unsetenv("OAUTH_SCOPES_SUPPORTED") +} diff --git a/token-broker/cmd/main.go b/token-broker/cmd/main.go new file mode 100644 index 00000000..b5bc1a3d --- /dev/null +++ b/token-broker/cmd/main.go @@ -0,0 +1,372 @@ +/* +Copyright 2025. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "net/http" + "os" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/go-chi/chi/v5" + "github.com/go-chi/chi/v5/middleware" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/jwtvalidation/validation" + "github.com/kagenti/token-broker/internal/api" + "github.com/kagenti/token-broker/internal/cache" + "github.com/kagenti/token-broker/internal/core" + "github.com/kagenti/token-broker/internal/oauth" + "github.com/kagenti/token-broker/internal/session" +) + +// Config holds the Token Broker configuration. +type Config struct { + ListenPort int + ClientID string + ClientSecret string + CallbackURL string + SessionTimeout time.Duration + MaxSessionsPerUser int + TokenWaitTimeout time.Duration + + // OAuth metadata (optional - if not set, discovery will be used) + AuthorizationEndpoint string + TokenEndpoint string + ScopesSupported []string + + // ResourceScopes maps resource server URL to OAuth scopes (legacy, scopes-only). + // Format of RESOURCE_SCOPES env var: "url1=scope1 scope2,url2=scope3" + // Deprecated: use ResourceConfigs (RESOURCE_CONFIG) for richer per-server config. + ResourceScopes map[string][]string + + // ResourceConfigs maps resource server URL to full per-server OAuth config overrides. + // Format of RESOURCE_CONFIG env var: JSON map of server URL to ResourceConfig. + ResourceConfigs map[string]oauth.ResourceConfig + + // AllowedRedirectHosts is the list of permitted hostnames for backend_session_redirect_url. + // Format of ALLOWED_REDIRECT_HOSTS env var: comma-separated hostnames. + AllowedRedirectHosts []string + + // JWT validation (required for production). + JWKSUrl string + JWTIssuer string + JWTAudiences []string +} + +// DefaultConfig returns the default configuration. +func DefaultConfig() *Config { + return &Config{ + ListenPort: 8190, + ClientID: "", + ClientSecret: "", + CallbackURL: "http://localhost:8190/oauth/callback", + SessionTimeout: 60 * time.Second, + MaxSessionsPerUser: 5, + TokenWaitTimeout: 300 * time.Second, + } +} + +// LoadConfig loads configuration from environment variables. +func LoadConfig() *Config { + cfg := DefaultConfig() + loadCoreConfig(cfg) + loadOAuthDiscoveryConfig(cfg) + loadResourceConfig(cfg) + loadSecurityConfig(cfg) + return cfg +} + +func loadCoreConfig(cfg *Config) { + if port := os.Getenv("TOKEN_BROKER_PORT"); port != "" { + if _, err := fmt.Sscanf(port, "%d", &cfg.ListenPort); err != nil { + slog.Default().Warn("invalid TOKEN_BROKER_PORT, using default", + "value", port, "default", cfg.ListenPort, "error", err) + } + } + if v := os.Getenv("OAUTH_CLIENT_ID"); v != "" { + cfg.ClientID = v + } + if v := os.Getenv("OAUTH_CLIENT_SECRET"); v != "" { + cfg.ClientSecret = v + } + if v := os.Getenv("OAUTH_CALLBACK_URL"); v != "" { + cfg.CallbackURL = v + } + if v := os.Getenv("TOKEN_BROKER_SESSION_TIMEOUT"); v != "" { + if d, err := time.ParseDuration(v); err == nil { + cfg.SessionTimeout = d + } + } + if v := os.Getenv("TOKEN_BROKER_MAX_SESSIONS_PER_USER"); v != "" { + if _, err := fmt.Sscanf(v, "%d", &cfg.MaxSessionsPerUser); err != nil { + slog.Default().Warn("invalid TOKEN_BROKER_MAX_SESSIONS_PER_USER, using default", + "value", v, "default", cfg.MaxSessionsPerUser, "error", err) + } + } + if v := os.Getenv("TOKEN_BROKER_TOKEN_WAIT_TIMEOUT"); v != "" { + if d, err := time.ParseDuration(v); err == nil { + cfg.TokenWaitTimeout = d + } + } +} + +func loadOAuthDiscoveryConfig(cfg *Config) { + if v := os.Getenv("OAUTH_AUTHORIZATION_ENDPOINT"); v != "" { + cfg.AuthorizationEndpoint = v + } + if v := os.Getenv("OAUTH_TOKEN_ENDPOINT"); v != "" { + cfg.TokenEndpoint = v + } + if scopes := os.Getenv("OAUTH_SCOPES_SUPPORTED"); scopes != "" { + for _, s := range strings.Split(scopes, ",") { + if trimmed := strings.TrimSpace(s); trimmed != "" { + cfg.ScopesSupported = append(cfg.ScopesSupported, trimmed) + } + } + } +} + +func loadResourceConfig(cfg *Config) { + if serverScopes := os.Getenv("RESOURCE_SCOPES"); serverScopes != "" { + cfg.ResourceScopes = make(map[string][]string) + for _, entry := range strings.Split(serverScopes, ",") { + parts := strings.SplitN(strings.TrimSpace(entry), "=", 2) + if len(parts) != 2 { + continue + } + serverURL := strings.TrimSpace(parts[0]) + if serverURL == "" { + continue + } + cfg.ResourceScopes[serverURL] = append(cfg.ResourceScopes[serverURL], strings.Fields(parts[1])...) + } + } + if serverConfig := os.Getenv("RESOURCE_CONFIG"); serverConfig != "" { + if err := json.Unmarshal([]byte(serverConfig), &cfg.ResourceConfigs); err != nil { + slog.Default().Warn("invalid RESOURCE_CONFIG, ignoring", "error", err) + } + } +} + +func loadSecurityConfig(cfg *Config) { + if allowedHosts := os.Getenv("ALLOWED_REDIRECT_HOSTS"); allowedHosts != "" { + for _, h := range strings.Split(allowedHosts, ",") { + if trimmed := strings.TrimSpace(h); trimmed != "" { + cfg.AllowedRedirectHosts = append(cfg.AllowedRedirectHosts, trimmed) + } + } + } + if v := os.Getenv("JWT_JWKS_URL"); v != "" { + cfg.JWKSUrl = v + } + if v := os.Getenv("JWT_ISSUER"); v != "" { + cfg.JWTIssuer = v + } + if audiences := os.Getenv("JWT_AUDIENCE"); audiences != "" { + for _, a := range strings.Split(audiences, ",") { + if trimmed := strings.TrimSpace(a); trimmed != "" { + cfg.JWTAudiences = append(cfg.JWTAudiences, trimmed) + } + } + } +} + +// Validate checks if the configuration is valid. +func (c *Config) Validate() error { + if c.ListenPort <= 0 || c.ListenPort > 65535 { + return fmt.Errorf("invalid listen port: %d", c.ListenPort) + } + + if c.ClientID == "" { + return fmt.Errorf("OAuth client ID is required (set OAUTH_CLIENT_ID)") + } + + if c.ClientSecret == "" { + return fmt.Errorf("OAuth client secret is required (set OAUTH_CLIENT_SECRET)") + } + + if c.CallbackURL == "" { + return fmt.Errorf("OAuth callback URL is required (set OAUTH_CALLBACK_URL)") + } + + if c.SessionTimeout <= 0 { + return fmt.Errorf("session timeout must be positive") + } + + if c.MaxSessionsPerUser <= 0 { + return fmt.Errorf("max sessions per user must be positive") + } + + if c.TokenWaitTimeout <= 0 { + return fmt.Errorf("token wait timeout must be positive") + } + + return nil +} + +func main() { + logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{ + Level: slog.LevelInfo, + })) + slog.SetDefault(logger) + + logger.Info("Starting Token Broker service") + + cfg := LoadConfig() + if err := cfg.Validate(); err != nil { + logger.Error("Invalid configuration", "error", err) + os.Exit(1) + } + + logger.Info("Configuration loaded", + "listen_port", cfg.ListenPort, + "callback_url", cfg.CallbackURL, + "session_timeout", cfg.SessionTimeout, + "max_sessions_per_user", cfg.MaxSessionsPerUser, + "token_wait_timeout", cfg.TokenWaitTimeout, + "has_configured_auth_endpoint", cfg.AuthorizationEndpoint != "", + "has_configured_token_endpoint", cfg.TokenEndpoint != "", + "configured_scopes_count", len(cfg.ScopesSupported), + "allowed_redirect_hosts", cfg.AllowedRedirectHosts, + "jwt_validation_enabled", cfg.JWKSUrl != "", + "jwt_issuer", cfg.JWTIssuer) + + if len(cfg.AllowedRedirectHosts) == 0 { + logger.Warn("ALLOWED_REDIRECT_HOSTS is not set — all redirect hosts are permitted; set this in production") + } + if cfg.JWKSUrl == "" { + logger.Warn("JWT_JWKS_URL is not set — JWT signatures are NOT validated", + "recommendation", "set JWT_JWKS_URL, JWT_ISSUER, and JWT_AUDIENCE in production") + } + + // Initialize components + clock := &core.RealClock{} + tokenCache := cache.NewTokenCache(clock) + sessionManager := session.NewSessionManager(cfg.SessionTimeout, cfg.MaxSessionsPerUser, clock, logger) + + oauthConfig := &oauth.Config{ + ClientID: cfg.ClientID, + ClientSecret: cfg.ClientSecret, + } + oauthClient := oauth.NewClient(oauthConfig, logger) + + discovererConfig := &oauth.DiscovererConfig{ + AuthorizationEndpoint: cfg.AuthorizationEndpoint, + TokenEndpoint: cfg.TokenEndpoint, + ScopesSupported: cfg.ScopesSupported, + ResourceScopes: cfg.ResourceScopes, + ResourceConfigs: cfg.ResourceConfigs, + } + discoverer := oauth.NewDiscovererWithConfig(discovererConfig, logger) + + broker := core.NewTokenBroker( + sessionManager, + tokenCache, + discoverer, + oauthClient, + cfg.CallbackURL, + cfg.TokenWaitTimeout, + logger, + ) + + // Create HTTP router + r := chi.NewRouter() + r.Use(middleware.RequestID) + // middleware.RealIP intentionally omitted: it trusts X-Forwarded-For / + // X-Real-IP headers (staticcheck SA1019), which can be spoofed. For a + // security-sensitive service we keep r.RemoteAddr as the real peer address. + r.Use(middleware.Recoverer) + r.Use(middleware.Timeout(cfg.TokenWaitTimeout + 10*time.Second)) + + // Health endpoints for Kubernetes probes — registered without Logger middleware + // to avoid flooding logs with probe traffic. + r.Get("/healthz", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + r.Get("/readyz", func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + // All other routes get request logging. + r.Group(func(r chi.Router) { + r.Use(middleware.Logger) + // Register API handlers + apiHandler := api.NewHandler(broker, sessionManager, logger, cfg.AllowedRedirectHosts) + if cfg.JWKSUrl != "" { + verifier := validation.NewLazyJWKSVerifier(cfg.JWKSUrl, cfg.JWTIssuer) + apiHandler.WithJWTVerifier(verifier, cfg.JWTAudiences) + } + apiHandler.RegisterRoutes(r) + }) + + // Create HTTP server + addr := fmt.Sprintf(":%d", cfg.ListenPort) + srv := &http.Server{ + Addr: addr, + Handler: r, + ReadTimeout: 30 * time.Second, + WriteTimeout: cfg.TokenWaitTimeout + 30*time.Second, + IdleTimeout: 120 * time.Second, + } + + go func() { + logger.Info("Token Broker listening", "address", addr) + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + logger.Error("Server error", "error", err) + os.Exit(1) + } + }() + + // Wait for interrupt signal + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + go func() { + ticker := time.NewTicker(5 * time.Minute) + defer ticker.Stop() + for { + select { + case <-ticker.C: + removed := tokenCache.CleanupExpiredTokens() + if removed > 0 { + logger.Info("Periodic token cache cleanup", "removed", removed) + } + case <-ctx.Done(): + return + } + } + }() + + <-ctx.Done() + + logger.Info("Shutting down Token Broker") + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if err := srv.Shutdown(shutdownCtx); err != nil { + logger.Error("Server shutdown error", "error", err) + } + + sessionManager.Shutdown() + logger.Info("Token Broker stopped") +} diff --git a/token-broker/deploy/00-serviceaccount.yaml b/token-broker/deploy/00-serviceaccount.yaml new file mode 100644 index 00000000..35533c73 --- /dev/null +++ b/token-broker/deploy/00-serviceaccount.yaml @@ -0,0 +1,12 @@ +# ServiceAccount for the Token Broker. +# +# The Token Broker is a stateless, in-memory service. It does NOT call the +# Kubernetes API, so it needs no Role/RoleBinding — just an identity to run as. +--- +apiVersion: v1 +kind: ServiceAccount +metadata: + name: token-broker + namespace: kagenti-system + labels: + app: token-broker diff --git a/token-broker/deploy/01-secret.example.yaml b/token-broker/deploy/01-secret.example.yaml new file mode 100644 index 00000000..ec224f6b --- /dev/null +++ b/token-broker/deploy/01-secret.example.yaml @@ -0,0 +1,22 @@ +# OAuth client credentials for the Token Broker. +# +# DO NOT commit real values. This is a TEMPLATE — copy it, fill in the GitHub +# (or other provider) OAuth app credentials, and apply it separately, or wire +# these keys to an external secret store (e.g. ExternalSecrets / Vault). +# +# Create directly without committing a file: +# kubectl create secret generic github-oauth-credentials -n kagenti-system \ +# --from-literal=client-id= \ +# --from-literal=client-secret= +--- +apiVersion: v1 +kind: Secret +metadata: + name: github-oauth-credentials + namespace: kagenti-system + labels: + app: token-broker +type: Opaque +stringData: + client-id: "REPLACE_WITH_OAUTH_CLIENT_ID" + client-secret: "REPLACE_WITH_OAUTH_CLIENT_SECRET" diff --git a/token-broker/deploy/02-deployment.yaml b/token-broker/deploy/02-deployment.yaml new file mode 100644 index 00000000..51323fc6 --- /dev/null +++ b/token-broker/deploy/02-deployment.yaml @@ -0,0 +1,115 @@ +# Token Broker Deployment. +# +# Stateless service that brokers OAuth flows and caches tokens in memory. +# Placeholders to set per environment: +# - image: the tag your build script produces +# - OAUTH_CALLBACK_URL: the externally-reachable URL routed by 04-httproute.yaml; +# its host MUST match the HTTPRoute hostname. +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: token-broker + namespace: kagenti-system + labels: + app: token-broker +spec: + replicas: 1 # In-memory sessions/token cache: do NOT scale >1 without shared state. + selector: + matchLabels: + app: token-broker + template: + metadata: + labels: + app: token-broker + annotations: + # Broker terminates its own OAuth flows; no mesh sidecar required. + sidecar.istio.io/inject: "false" + spec: + serviceAccountName: token-broker + securityContext: + runAsNonRoot: true + containers: + - name: token-broker + image: localhost/token-broker:latest # REPLACE with the tag your build script produces. + imagePullPolicy: IfNotPresent + ports: + - name: http + containerPort: 8190 + protocol: TCP + env: + - name: TOKEN_BROKER_PORT + value: "8190" + - name: OAUTH_CLIENT_ID + valueFrom: + secretKeyRef: + name: github-oauth-credentials + key: client-id + - name: OAUTH_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: github-oauth-credentials + key: client-secret + # Externally-reachable callback URL; host must match the HTTPRoute hostname. + - name: OAUTH_CALLBACK_URL + value: "http://token-broker.localtest.me:8080/oauth/callback" + # Permitted hostnames for backend_session_redirect_url. + # Comma-separated. Set to the hostname(s) of your backend UI. + - name: ALLOWED_REDIRECT_HOSTS + value: "app-demo.localtest.me" + # JWT validation (required for production — leave unset for dev/test only). + # JWT_JWKS_URL: JWKS endpoint of your identity provider. + # JWT_ISSUER: Expected issuer claim in incoming JWTs. + # JWT_AUDIENCE: Comma-separated expected audience(s). + # - name: JWT_JWKS_URL + # value: "https://keycloak.example.com/realms/kagenti/protocol/openid-connect/certs" + # - name: JWT_ISSUER + # value: "https://keycloak.example.com/realms/kagenti" + # - name: JWT_AUDIENCE + # value: "token-broker" + # Optional tunables (defaults in cmd/token-broker/main.go): + # - name: TOKEN_BROKER_SESSION_TIMEOUT + # value: "60s" + # - name: TOKEN_BROKER_MAX_SESSIONS_PER_USER + # value: "5" + # - name: TOKEN_BROKER_TOKEN_WAIT_TIMEOUT + # value: "300s" + # Optional: skip OAuth discovery by configuring endpoints directly. + # - name: OAUTH_AUTHORIZATION_ENDPOINT + # value: "https://github.com/login/oauth/authorize" + # - name: OAUTH_TOKEN_ENDPOINT + # value: "https://github.com/login/oauth/access_token" + # - name: OAUTH_SCOPES_SUPPORTED + # value: "repo,read:org,read:user" + # Per-server OAuth config (scopes + endpoint overrides). JSON format. + # Each key is the resource server URL (must match X-Server-Url exactly). + # Fields: scopes ([]string), authorization_endpoint (string), token_endpoint (string). + - name: RESOURCE_CONFIG + value: >- + { + "http://mcp-server-service.kagenti-demo.svc.cluster.local:8184": {"scopes": ["read:user", "user:email"]} + } + livenessProbe: + httpGet: + path: /healthz + port: http + initialDelaySeconds: 5 + periodSeconds: 10 + readinessProbe: + httpGet: + path: /readyz + port: http + initialDelaySeconds: 2 + periodSeconds: 5 + resources: + requests: + cpu: 50m + memory: 64Mi + limits: + memory: 256Mi + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: + - ALL diff --git a/token-broker/deploy/03-service.yaml b/token-broker/deploy/03-service.yaml new file mode 100644 index 00000000..848839bc --- /dev/null +++ b/token-broker/deploy/03-service.yaml @@ -0,0 +1,21 @@ +# Token Broker Service (in-cluster). +# +# Consumed by AuthBridge and the Backend at: +# http://token-broker.kagenti-system.svc.cluster.local:8190 +--- +apiVersion: v1 +kind: Service +metadata: + name: token-broker + namespace: kagenti-system + labels: + app: token-broker +spec: + type: ClusterIP + selector: + app: token-broker + ports: + - name: http + port: 8190 + targetPort: http + protocol: TCP diff --git a/token-broker/deploy/04-httproute.yaml b/token-broker/deploy/04-httproute.yaml new file mode 100644 index 00000000..50a9118a --- /dev/null +++ b/token-broker/deploy/04-httproute.yaml @@ -0,0 +1,38 @@ +# HTTPRoute exposing the Token Broker OAuth callback through the shared gateway. +# +# The OAuth provider (e.g. GitHub) redirects the user's browser to +# OAUTH_CALLBACK_URL after authorization. That URL must resolve to GET +# /oauth/callback on the broker — otherwise the redirect 404s and the broker +# waits for a callback that never arrives. +# +# This route lives in kagenti-system (same namespace as the Service AND the +# gateway), so: +# - the parentRef is local (gateway "http" in kagenti-system), and +# - the backendRef is same-namespace, so NO ReferenceGrant is needed. +# +# The kagenti-system namespace carries the label shared-gateway-access: "true", +# which the "http" gateway's listener selects — same as the working app-demo route. +# +# The hostname MUST match the host in the Deployment's OAUTH_CALLBACK_URL. +--- +apiVersion: gateway.networking.k8s.io/v1 +kind: HTTPRoute +metadata: + name: token-broker-oauth-callback + namespace: kagenti-system + labels: + app: token-broker +spec: + parentRefs: + - name: http + namespace: kagenti-system + hostnames: + - "token-broker.localtest.me" + rules: + - matches: + - path: + type: PathPrefix + value: /oauth/callback + backendRefs: + - name: token-broker + port: 8190 diff --git a/token-broker/deploy/README.md b/token-broker/deploy/README.md new file mode 100644 index 00000000..2944f2e2 --- /dev/null +++ b/token-broker/deploy/README.md @@ -0,0 +1,56 @@ +# Token Broker — Kubernetes deployment + +Manifests that accompany the `token-broker` image produced by the build script. +All objects live in the `kagenti-system` namespace. + +## Manifests (apply in order) + +| File | Object | Notes | +|------|--------|-------| +| `00-serviceaccount.yaml` | ServiceAccount | No RBAC — the broker does not call the Kubernetes API. | +| `01-secret.example.yaml` | Secret (template) | OAuth client credentials. **Do not commit real values** — see below. | +| `02-deployment.yaml` | Deployment | Set `image:` and `OAUTH_CALLBACK_URL`. `replicas: 1` (in-memory state). | +| `03-service.yaml` | Service (ClusterIP) | `token-broker.kagenti-system.svc.cluster.local:8190`. | +| `04-httproute.yaml` | HTTPRoute | Exposes `GET /oauth/callback` via the shared `http` gateway. | + +## Credentials + +Do not apply `01-secret.example.yaml` as-is. Create the secret out of band: + +```sh +kubectl create secret generic github-oauth-credentials -n kagenti-system \ + --from-literal=client-id= \ + --from-literal=client-secret= +``` + +## Apply + +```sh +kubectl apply -f 00-serviceaccount.yaml +# create the secret (see above) instead of applying the example +kubectl apply -f 02-deployment.yaml -f 03-service.yaml -f 04-httproute.yaml +``` + +## The callback route must match + +`OAUTH_CALLBACK_URL` in the Deployment and the `hostnames:` in the HTTPRoute +must use the **same host**. The OAuth provider redirects the browser to that +URL; the gateway routes the host + `/oauth/callback` to the broker. A mismatch +(or a missing/unattached route) makes the post-consent redirect 404, and the +broker then blocks waiting for a callback that never arrives. + +The route attaches to the gateway `http` in `kagenti-system`. That listener +admits routes from namespaces labelled `shared-gateway-access: "true"`, which +`kagenti-system` already carries. + +## Migrating from the old ad-hoc objects + +Earlier deployments placed the HTTPRoute in `kagenti-demo` with a parentRef to a +non-existent `kagenti-gateway`/`istio-system`, which never attached (404 on the +callback). The token broker does not deploy into `kagenti-demo`. Remove the +stale objects: + +```sh +kubectl delete httproute token-broker-oauth-callback -n kagenti-demo --ignore-not-found +kubectl delete referencegrant allow-demo-httproute-to-token-broker -n kagenti-system --ignore-not-found +``` diff --git a/token-broker/docs/API_SPECIFICATION.md b/token-broker/docs/API_SPECIFICATION.md new file mode 100644 index 00000000..afce08a1 --- /dev/null +++ b/token-broker/docs/API_SPECIFICATION.md @@ -0,0 +1,727 @@ +# API Specification: Token Broker Service Interfaces + +The Token Broker enables HITL (Human-in-the-Loop) Authorization: when an agent needs access to a resource (MCP server, LLM API, direct REST API, another agent) that requires permissions beyond those already in its token, the Token Broker brings the user into the loop via an OAuth 2.0 consent flow. + +This document provides detailed API specifications for the interfaces between: +1. **Token Broker Service** ↔ **AuthBridge Sidecar** (inside AI Agent Pod) +2. **Token Broker Service** ↔ **Backend** + +## Table of Contents + +- [Overview](#overview) +- [Token Broker Service API](#token-broker-service-api) +- [Interface 1: Token Broker ↔ AuthBridge Sidecar](#interface-1-token-broker--authbridge-sidecar) +- [Interface 2: Token Broker ↔ Backend](#interface-2-token-broker--backend) +- [Error Responses](#error-responses) +- [Headers Reference](#headers-reference) + +--- + +## Overview + +The Token Broker Service is a centralized OAuth session and token management service that coordinates OAuth flows between the Backend (user-facing) and the AuthBridge Sidecar (token injection). + +### Architecture Context + +``` +┌─────────────┐ ┌──────────────┐ ┌──────────────────┐ +│ Backend │◄────────┤ Token Broker ├────────►│ Resource Server │ +│ │ Events │ Service │Discovery│ (MCP/LLM/API/..) │ +└──────┬──────┘ └──────┬───────┘ └──────────────────┘ + │ │ + │ X-OAuth-Session-Key │ Token Cache + │ │ Session Store + ▼ ▼ +┌──────────────────────────────────────────────┐ +│ AI Agent Pod │ +│ ┌─────────┐ ┌───────────┐ ┌───────────┐ │ +│ │ Envoy │ │AuthBridge │ │ AI Agent │ │ +│ │ Sidecar │ │ Sidecar │ │ Container │ │ +│ │ │ │ │ │ │ │ +│ │Inbound: │ │ ext_proc │ │Listens on │ │ +│ │ 15124 │◄─┤ gRPC │ │127.0.0.1: │ │ +│ │ │ │ :9090 │ │ 8186 │ │ +│ │Outbound:│ │ │ │ │ │ +│ │ 15123 │◄─┤ Calls │ │ │ │ +│ │ │ │ Token │ │ │ │ +│ │ │ │ Broker │ │ │ │ +│ └─────────┘ └───────────┘ └───────────┘ │ +└──────────────────────────────────────────────┘ +``` + +--- + +## Token Broker Service API + +**Base URL**: `http://token-broker-service:8190` + +**Service Port**: 8190 (ClusterIP) + +**Implementation**: [`internal/api/handlers.go`](../../internal/api/handlers.go) + +--- + +## Interface 1: Token Broker ↔ AuthBridge Sidecar + +The AuthBridge Sidecar calls the Token Broker to obtain OAuth tokens for resource server requests. + +### 1.1 Request Token + +**Endpoint**: `POST /sessions/token` + +**Purpose**: Request an OAuth token for a specific resource server using JWT-based authentication. This call blocks until the token is available or the request times out. + +**Called By**: AuthBridge Sidecar (from kagenti-extensions/authbridge) + +**Authentication**: JWT Bearer token containing user ID and session key + +**Request Headers**: +``` +Authorization: Bearer +X-Server-Url: +``` + +**JWT Claims** (required): +- `sub` (string): User ID +- `session_uid` (string): Session key (preferred), or `jti` (string) as fallback +- `iss` (string): Issuer (OIDC service URL) +- `aud` (string): Audience +- `exp` (number): Expiration timestamp +- `iat` (number): Issued at timestamp + +**Security Note**: JWT signature validation is enabled when `JWT_JWKS_URL` is configured (recommended for production). When unset, the Token Broker logs a WARNING and trusts claims without signature verification — acceptable only in dev/test environments with network policies restricting access. + +**Request Body**: None + +**Success Response** (200 OK): +```json +{ + "token": "ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" +} +``` + +**Behavior**: +1. Extracts JWT from `Authorization: Bearer ` header +2. Validates JWT (signature, expiry, issuer, audience when `JWT_JWKS_URL` is configured; claims-only otherwise) +3. Extracts `sub` (user ID) and `jti` (session key) from JWT claims +4. Validates session exists and belongs to the specified user +5. Checks token cache for `(user_id, resource_url)` pair +6. If token exists and not expired, returns immediately +7. If token missing: + - Acquires per-session semaphore (ensures only one OAuth flow per session) + - Rechecks cache (another request may have obtained token) + - Initiates OAuth discovery and flow: + - Calls `GET /.well-known/oauth-protected-resource` on resource server to discover endpoints and scopes + - **Generates PKCE at Token Broker** (code_verifier, code_challenge) + - **Builds authorization URL at Token Broker** with PKCE parameters and Token Broker's callback URL + - Publishes `oauth_url_ready` event to Backend with complete auth URL + - Waits for OAuth provider to call `GET /oauth/callback` with authorization code + - **Exchanges code directly with OAuth provider** (not through resource server) + - Blocks waiting for OAuth completion (up to 300s by default) +8. Returns token when available + +**Session Validation Security**: The Token Broker validates session ownership BEFORE checking the token cache. This prevents session hijacking where User B could use User A's stolen session key to obtain User A's cached tokens. The validation order is: +1. Parse JWT and extract claims +2. Validate session exists and belongs to the user from JWT +3. Check token cache +4. If cache miss, initiate OAuth flow + +**Error Responses** (flat format): +- `400 Bad Request`: Missing required headers + ```json + { + "code": "invalid_request", + "message": "Missing X-Server-Url header" + } + ``` +- `401 Unauthorized`: Invalid JWT or session + ```json + { + "code": "invalid_token", + "message": "Invalid JWT: missing required claim: sub (user ID)" + } + ``` + ```json + { + "code": "unauthorized", + "message": "Invalid session or user mismatch" + } + ``` +- `408 Request Timeout`: OAuth flow did not complete in time + ```json + { + "code": "timeout", + "message": "OAuth flow did not complete in time" + } + ``` +- `503 Service Unavailable`: Failed to obtain token + ```json + { + "code": "oauth_failed", + "message": "Failed to obtain token:
" + } + ``` + +**Timeout**: 300 seconds (5 minutes) - allows time for user to complete OAuth flow + +**Implementation**: [`HandleGetToken()`](../../internal/api/handlers.go) + +--- + + +## Interface 2: Token Broker ↔ Backend + +The Backend creates sessions, polls for OAuth events, and submits OAuth completion callbacks. + +### 2.1 Create Session + +**Endpoint**: `POST /sessions` + +**Purpose**: Create a new OAuth session for a user with backend redirect URL + +**Called By**: Backend (when user submits a task) + +**Authentication**: JWT Bearer token containing user ID and session key + +**Request Headers**: +``` +Authorization: Bearer +Content-Type: application/json +``` + +**JWT Claims** (required): +- `sub` (string): User ID +- `session_uid` (string): Session key (preferred), or `jti` (string) as fallback +- `iss` (string): Issuer +- `aud` (string): Audience +- `exp` (number): Expiration timestamp +- `iat` (number): Issued at timestamp + +**Security Note**: JWT signature validation is enabled when `JWT_JWKS_URL` is configured (recommended for production). When unset, the Token Broker logs a WARNING and trusts claims without signature verification — acceptable only in dev/test environments with network policies restricting access. + +**Request Body**: +```json +{ + "backend_session_redirect_url": "https://backend.example.com/oauth/complete?session=abc123" +} +``` + +**Success Response** (201 Created): +Empty response body (session key is already known from JWT `jti` claim) + +**Behavior**: +1. Extracts JWT from `Authorization: Bearer ` header +2. Validates JWT (signature, expiry, issuer, audience when `JWT_JWKS_URL` is configured; claims-only otherwise) +3. Extracts `sub` (user ID) and `jti` (session key) from JWT claims +4. Validates user hasn't exceeded max sessions (default: 5) +5. Creates new session with provided session key +6. Stores `backend_session_redirect_url` for post-OAuth user redirect +7. Starts session timeout timer (default: 60s) + +**Error Responses**: +- `400 Bad Request`: Missing or invalid request body + ```json + { + "error": { + "code": "invalid_request", + "message": "Missing backend_session_redirect_url in request body" + } + } + ``` +- `401 Unauthorized`: Invalid JWT + ```json + { + "error": { + "code": "invalid_token", + "message": "Invalid JWT: missing required claim: sub (user ID)" + } + } + ``` +- `429 Too Many Requests`: Max sessions exceeded + ```json + { + "error": { + "code": "too_many_sessions", + "message": "max sessions per user exceeded" + } + } + ``` + +**Implementation**: [`HandleCreateSession()`](../../internal/api/handlers.go) + +--- + +### 2.2 OAuth Callback (NEW) + +**Endpoint**: `GET /oauth/callback` + +**Purpose**: Receive OAuth authorization code from OAuth provider and redirect user to Backend + +**Called By**: OAuth Provider (e.g., GitHub) after user authorizes + +**Query Parameters**: +- `code` (string, required): OAuth authorization code +- `state` (string, required): State parameter in format `sessionKey.nonce` + +**Success Response** (302 Found): +``` +Location: +``` + +**Behavior**: +1. Extracts `code` and `state` from query parameters +2. Parses state to extract session key (splits on last `.`) +3. Looks up session by session key; validates state matches the stored OAuth transaction +4. Signals the waiting `AcquireToken` goroutine with the authorization code +5. That goroutine exchanges the code directly with the OAuth provider and caches the token +6. Redirects user to Backend's `backend_session_redirect_url` (HTTP 302) + +**State Parameter Format**: `sessionKey.nonce` +- Session key may contain `.` characters +- Nonce is base64url encoded (no `.` characters) +- Parsing: Split on LAST `.` to extract session key + +**Error Responses**: +- `400 Bad Request`: Missing or invalid parameters + ```json + { + "error": { + "code": "invalid_request", + "message": "Missing code or state parameter" + } + } + ``` +- `401 Unauthorized`: Invalid state or session not found + ```json + { + "error": { + "code": "invalid_state", + "message": "Invalid or expired OAuth state" + } + } + ``` +- `503 Service Unavailable`: Token exchange failed + ```json + { + "error": { + "code": "oauth_failed", + "message": "Failed to exchange authorization code" + } + } + ``` + +**Implementation**: [`HandleOAuthCallback()`](../../internal/api/handlers.go) + +--- + +### 2.3 Poll for Events (Long-Polling) + +**Endpoint**: `POST /sessions/broker-events` + +**Purpose**: Long-poll for OAuth events (e.g., `oauth_required`, `oauth_url_ready`) + +**Called By**: Backend + +**Authentication**: JWT Bearer token containing user ID and session key + +**Request Headers**: +``` +Authorization: Bearer +``` + +**JWT Claims** (required): +- `sub` (string): User ID +- `session_uid` (string): Session key (preferred), or `jti` (string) as fallback +- `iss` (string): Issuer +- `aud` (string): Audience +- `exp` (number): Expiration timestamp +- `iat` (number): Issued at timestamp + +**Security Note**: JWT signature validation is enabled when `JWT_JWKS_URL` is configured (recommended for production). When unset, the Token Broker logs a WARNING and trusts claims without signature verification — acceptable only in dev/test environments with network policies restricting access. + +**Request Body**: None + +**Behavior**: Blocks until an event is available or the session closes. Resets the session idle timer while a Backend request is in flight. + +**Success Response** (200 OK) - Authorization URL Ready Event: +```json +{ + "type": "oauth_url_ready", + "auth_url": "https://github.com/login/oauth/authorize?client_id=...&code_challenge=...&state=..." +} +``` + +**Backend Action**: Open OAuth popup with auth_url + +**Success Response** (200 OK) - Error Event: +```json +{ + "type": "error", + "message": "Session expired", + "code": "session_expired" +} +``` + +**Event Processing**: +1. Extracts JWT from `Authorization: Bearer ` header +2. Validates JWT (signature, expiry, issuer, audience when `JWT_JWKS_URL` is configured; claims-only otherwise) +3. Extracts `sub` (user ID) and `jti` (session key) from JWT claims +4. Validates session and user +5. Resets session timeout timer (Backend is connected) +6. Waits for event on session's event channel or until session closes +7. Returns event when available or an error event after session close +8. Starts session timeout timer when request completes + +**Error Responses**: +- `401 Unauthorized`: Invalid JWT or session + ```json + { + "error": { + "code": "invalid_token", + "message": "Invalid JWT: missing required claim: sub (user ID)" + } + } + ``` + ```json + { + "error": { + "code": "unauthorized", + "message": "Invalid session or user mismatch" + } + } + ``` + +**Implementation**: [`HandleEvents()`](../../internal/api/handlers.go) + +--- + +### 2.4 End Session + +**Endpoint**: `POST /sessions/end` + +**Purpose**: Explicitly end a session and release resources + +**Called By**: Backend (when task completes or user disconnects) + +**Authentication**: JWT Bearer token containing user ID and session key + +**Request Headers**: +``` +Authorization: Bearer +``` + +**JWT Claims** (required): +- `sub` (string): User ID +- `session_uid` (string): Session key (preferred), or `jti` (string) as fallback +- `iss` (string): Issuer +- `aud` (string): Audience +- `exp` (number): Expiration timestamp +- `iat` (number): Issued at timestamp + +**Security Note**: JWT signature validation is enabled when `JWT_JWKS_URL` is configured (recommended for production). When unset, the Token Broker logs a WARNING and trusts claims without signature verification — acceptable only in dev/test environments with network policies restricting access. + +**Request Body**: None + +**Success Response** (200 OK): +Empty response body + +**Behavior**: +1. Extracts JWT from `Authorization: Bearer ` header +2. Validates JWT (signature, expiry, issuer, audience when `JWT_JWKS_URL` is configured; claims-only otherwise) +3. Extracts `sub` (user ID) and `session_uid`/`jti` (session key) from JWT claims +4. Validates session and user +5. Marks session as ended +6. Wakes all blocked token waiters with session-ended error +7. Releases session resources +8. Stops session timeout timer + +**Error Responses**: +- `401 Unauthorized`: Invalid JWT or session + ```json + { + "error": { + "code": "invalid_token", + "message": "Invalid JWT: missing required claim: sub (user ID)" + } + } + ``` + ```json + { + "error": { + "code": "unauthorized", + "message": "Invalid session or user mismatch" + } + } + ``` +- `500 Internal Server Error`: Failed to end session + +**Implementation**: [`HandleEndSession()`](../../internal/api/handlers.go) + +--- + +## Error Responses + +All Token Broker error responses follow a consistent JSON format: + +```json +{ + "error": { + "code": "", + "message": "" + } +} +``` + +### Common Error Codes + +| Code | HTTP Status | Description | +|------|-------------|-------------| +| `invalid_request` | 400 | Missing required headers or parameters | +| `unauthorized` | 401 | Invalid session or user mismatch | +| `session_expired` | 401 | Session has expired | +| `session_ended` | 401 | Session was explicitly ended | +| `too_many_sessions` | 429 | User has exceeded max concurrent sessions | +| `timeout` | 408 | OAuth flow did not complete in time | +| `oauth_failed` | 503 | Failed to obtain token from OAuth provider | +| `internal_error` | 500 | Internal server error | + +--- + +## Headers Reference + +### Required Headers + +| Header | Used By | Description | +|--------|---------|-------------| +| `Authorization` | All endpoints | JWT Bearer token containing `sub` (user ID) and `session_uid`/`jti` (session key) | +| `X-Server-Url` | `POST /sessions/token` | Target resource server URL for token acquisition | + +### JWT Authentication + +All endpoints use JWT Bearer authentication. The Token Broker extracts `sub` (user ID) and `session_uid` or `jti` (session key) from the JWT claims. When `JWT_JWKS_URL` is configured, full signature/expiry/issuer/audience validation is performed via JWKS. When unset (dev/test mode), claims are trusted without verification and a WARNING is logged at startup. + +**JWT Generation**: JWTs must contain the following claims: +- `sub`: User ID +- `session_uid`: Session key (preferred); `jti` is accepted as fallback +- `iss`: Issuer (must match `JWT_ISSUER` when validation is enabled) +- `aud`: Audience (must match `JWT_AUDIENCE` when validation is enabled) +- `iat`: Issued at timestamp +- `exp`: Expiration timestamp + +**JWT Validation**: When `JWT_JWKS_URL` is configured, the Token Broker validates signatures, expiry, issuer, and audience via JWKS. When unset (dev/test mode), claims are trusted without verification and a WARNING is logged at startup. + +--- + +## Configuration + +The Token Broker Service is configured via environment variables or configuration file: + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `TOKEN_BROKER_PORT` | 8190 | HTTP server port | +| `OAUTH_CLIENT_ID` | **(required)** | OAuth client ID | +| `OAUTH_CLIENT_SECRET` | **(required)** | OAuth client secret | +| `OAUTH_CALLBACK_URL` | **(required)** | Token Broker's public OAuth callback URL | +| `ALLOWED_REDIRECT_HOSTS` | *(none — all permitted, WARNING logged)* | Comma-separated permitted hostnames for `backend_session_redirect_url` | +| `JWT_JWKS_URL` | *(none — validation disabled, WARNING logged)* | JWKS endpoint for JWT signature validation | +| `JWT_ISSUER` | *(none)* | Expected JWT issuer | +| `JWT_AUDIENCE` | *(none)* | Comma-separated expected audiences | +| `TOKEN_BROKER_SESSION_TIMEOUT` | 60s | Idle session timeout after Backend disconnect | +| `TOKEN_BROKER_MAX_SESSIONS_PER_USER` | 5 | Maximum concurrent sessions per user | +| `TOKEN_BROKER_TOKEN_WAIT_TIMEOUT` | 300s | Maximum time to wait for OAuth completion | + +**Note**: The `callback_url` must be registered with the OAuth provider (e.g., in GitHub App settings) and must be publicly accessible. + +--- + +## Token Caching + +Tokens are cached per `(session_key, resource_url)` pair: + +- **Cache Key**: `session_key` + `resource_url` +- **Cache Value**: Access token + expiry time +- **Expiry Logic**: + - Tokens are treated as expired if less than 5 minutes remain + - Expired tokens are automatically removed from cache +- **Session isolation**: Only agents belonging to the same user session may use a cached token + +**Implementation**: [`internal/cache/token_cache.go`](../../internal/cache/token_cache.go) + +--- + +## Session Management + +Sessions are managed with the following lifecycle: + +1. **Creation**: Backend calls `POST /sessions` +2. **Active**: Backend polls `/events`, AuthBridge requests tokens +3. **Timeout**: Session expires after 60s of Backend inactivity +4. **Explicit End**: Backend calls `POST /sessions/end` + +**Session Semaphore**: Each session has a semaphore (capacity: 1) to ensure only one OAuth flow runs at a time per session. + +**Implementation**: [`internal/session/manager.go`](../../internal/session/manager.go) + +--- + +## OAuth Discovery and Token Exchange Flow (CURRENT ARCHITECTURE) + +When a token is needed and not cached, the Token Broker coordinates an OAuth flow where it acts as the OAuth client and receives callbacks directly. + +### Current Flow (Token Broker Receives OAuth Callbacks) + +1. **Discovery**: Token Broker calls `GET /.well-known/oauth-protected-resource` on resource server + - Discovers OAuth authorization server base URL + - Derives authorization_endpoint (base + `/authorize`) + - Derives token_endpoint (base + `/access_token`) + - Extracts scopes_supported +2. **PKCE Generation**: **Token Broker generates PKCE** using `pkg/oauth.GeneratePKCEChallenge()` + - Creates `code_verifier` (random 43-character string) + - Creates `code_challenge` (SHA256 hash of verifier, base64url encoded) + - Generates `state` in format `sessionKey.nonce` using `pkg/oauth.GenerateState()` +3. **Storage**: Token Broker stores PKCE, state, and endpoints in session's OAuth transaction +4. **URL Building**: **Token Broker builds authorization URL** with: + - Token Broker's `client_id` + - Token Broker's `callback_url` (redirect_uri) - configured OAuth callback endpoint + - `code_challenge` and `code_challenge_method=S256` + - `state` (format: `sessionKey.nonce`) + - `scope` (discovered from resource server) +5. **Event**: Token Broker publishes `oauth_url_ready` event with complete authorization URL +6. **User Flow**: Backend opens OAuth popup for user +7. **Authorization**: User logs in and authorizes the application +8. **Callback**: OAuth provider redirects to Token Broker's `/oauth/callback` with `code` and `state` +9. **Session Lookup**: Token Broker extracts session key from state parameter +10. **Token Exchange**: **Token Broker exchanges directly with OAuth provider** + - Sends `client_id`, `client_secret`, `code`, `code_verifier`, `redirect_uri` to OAuth provider's token_endpoint + - OAuth provider validates PKCE and returns access token + - Token Broker receives access token directly (no resource server intermediary) +11. **Cache**: Token Broker caches received token for `(user_id, resource_url)` +12. **User Redirect**: Token Broker redirects user to Backend's `backend_session_redirect_url` (HTTP 302) +13. **Unblock**: Waiting token requests receive the token + +### Why Token Broker Receives OAuth Callbacks + +**Architecture Benefits**: + +- **Token Broker as OAuth Client**: Registered OAuth client with GitHub, has `client_id` and `client_secret` +- **Direct Callback Receipt**: OAuth provider redirects directly to Token Broker (standard OAuth 2.0 flow) +- **No Code Forwarding**: Authorization code never passes through Backend (more secure) +- **resource server Role**: Only provides metadata (endpoints and scopes via `.well-known`) +- **Centralized OAuth**: Single service handles all OAuth operations +- **Better Security**: `code_verifier` and authorization code never leave Token Broker +- **Simpler Backend**: No OAuth callback handling logic needed +- **Standard OAuth Flow**: Follows OAuth 2.0 best practices + +### PKCE Security + +**Implementation**: +- **Token Broker generates PKCE** for each OAuth flow +- `code_verifier` stays in Token Broker (never transmitted to user or Backend) +- `code_challenge` sent to OAuth provider in authorization URL +- `code_verifier` sent to OAuth provider during token exchange +- resource server never sees PKCE parameters +- Authorization code received directly by Token Broker (not forwarded through Backend) + +**Implementation**: [`internal/oauth/`](../../internal/oauth/) + +--- + +## Sequence Diagrams + +### Complete OAuth Flow + +```mermaid +sequenceDiagram + participant User + participant Backend + participant TokenBroker as Token Broker + participant AuthBridge + participant ResourceServer as Resource Server + participant OAuth as GitHub OAuth + + Backend->>+TokenBroker: POST /sessions
(JWT Bearer auth)
{backend_session_redirect_url} + TokenBroker-->>-Backend: 201 Created + + Backend->>TokenBroker: POST /sessions/broker-events
(JWT Bearer auth, long-poll) + + AuthBridge->>+TokenBroker: POST /sessions/token
(JWT Bearer auth) + TokenBroker->>TokenBroker: Parse JWT (no signature validation)
Validate session ownership
Check cache: miss
Acquire semaphore + TokenBroker->>+ResourceServer: GET /.well-known/oauth-protected-resource + ResourceServer-->>-TokenBroker: {authorization_servers, scopes_supported} + Note over TokenBroker: Token Broker generates PKCE:
code_verifier, code_challenge
state = sessionKey.nonce + TokenBroker->>TokenBroker: Store PKCE, state, and endpoints
Build auth URL with Token Broker callback + + TokenBroker-->>Backend: Event: {type: oauth_url_ready, auth_url} + Backend->>User: Open OAuth popup + User->>OAuth: Authorize application + OAuth-->>TokenBroker: GET /oauth/callback?code=...&state=... + + Note over TokenBroker: Extract session key from state
Validate state matches session + TokenBroker->>+OAuth: POST /token
{client_id, client_secret,
code, code_verifier, redirect_uri} + OAuth-->>-TokenBroker: {access_token, expires_in} + TokenBroker->>TokenBroker: Cache token for
(user_id, resource_url) + TokenBroker-->>User: HTTP 302 Redirect to
backend_session_redirect_url + + TokenBroker-->>-AuthBridge: {token} + User->>Backend: Redirected to session URL +``` + +### Cached Token (Fast Path) + +```mermaid +sequenceDiagram + participant Backend + participant TokenBroker as Token Broker + participant AuthBridge + + Backend->>+TokenBroker: POST /sessions
(JWT Bearer auth)
{backend_session_redirect_url} + TokenBroker-->>-Backend: 201 Created + + AuthBridge->>+TokenBroker: POST /sessions/token
(JWT Bearer auth) + TokenBroker->>TokenBroker: Parse JWT
Validate session
Check cache: HIT! + TokenBroker-->>-AuthBridge: {token} +``` + +--- + +## Related Documentation + +- [Architecture Overview](ARCHITECTURE.md) +- [Deployment Guide](DEPLOYMENT.md) +- [Testing Guide](TESTING.md) +- [Token Broker Architecture](../../docs/token_broker_architecture.md) +- [Envoy Sidecar Design](../../docs/envoy_sidecar_claude.md) + +--- + +## Implementation Notes + +### AuthBridge Configuration + +The AuthBridge Sidecar is configured to use the Token Broker via routes: + +```yaml +routes: + rules: + - host: "mcp-server-service" + action: "broker" + token_broker_url: "http://token-broker-service:8190" +``` + +**Configuration**: See deployment YAML files in this directory + +### Envoy Configuration + +Envoy is configured with ext_proc to call AuthBridge: + +- **Outbound Listener** (15123): Agent → MCP Server (token injection) +- **Inbound Listener** (15124): Backend → Agent (JWT validation, bypassed in demo) +- **ext_proc Cluster**: AuthBridge at 127.0.0.1:9090 + +**Configuration**: See deployment YAML files in this directory + +--- + +**Last Updated**: 2026-05-23 +**Version**: 1.0 +**Status**: Current Implementation \ No newline at end of file diff --git a/token-broker/docs/ARCHITECTURE.md b/token-broker/docs/ARCHITECTURE.md new file mode 100644 index 00000000..42ba4416 --- /dev/null +++ b/token-broker/docs/ARCHITECTURE.md @@ -0,0 +1,152 @@ +# Token Broker Architecture + +## Purpose + +The Token Broker enables **HITL (Human-in-the-Loop) Authorization**: when an agent needs access to a resource that requires OAuth permissions not already present in its token, the Token Broker brings the user into the loop via an OAuth 2.0 PKCE flow and caches the resulting access token for the session. + +Kagenti supports three route types for outbound agent requests: +- **Passthrough** — forward the token as-is +- **Token Exchange** — use a subset of permissions already in the agent's token +- **Token Broker** — obtain additional permissions with the user's help (this service) + +--- + +## System Context + +``` +┌─────────────┐ ┌──────────────┐ ┌──────────────────┐ +│ Backend │◄────────┤ Token Broker ├────────►│ Resource Server │ +│ │ Events │ │ Discovery│ (MCP/LLM/API/..) │ +└─────────────┘ └──────┬───────┘ └──────────────────┘ + │ Token Cache + │ Session Store + ▼ + ┌──────────────┐ + │ In-Memory │ + └──────────────┘ +``` + +The Token Broker is a standalone HTTP service. It is not a sidecar; it is a shared service within the cluster. + +- **AuthBridge sidecar** calls `POST /sessions/token` to obtain a token for a resource server. This call blocks until the OAuth flow completes. +- **Backend** creates sessions, long-polls for events (`POST /sessions/broker-events`), and ends sessions (`POST /sessions/end`). +- **OAuth provider** redirects the user's browser to `GET /oauth/callback` after consent. +- **Resource servers** are only contacted for `.well-known/oauth-protected-resource` metadata discovery. + +--- + +## Repository Structure + +``` +cmd/ + main.go # Service bootstrap, configuration, HTTP server, graceful shutdown + +internal/ + api/ # HTTP handlers and route registration + auth/ # JWT claim extraction (without signature validation) + cache/ # Token cache: per-(session_key, resource_url), JWT expiry parsing + core/ # Interfaces, session/transaction types, TokenBroker orchestrator + oauth/ # OAuth client (PKCE gen, URL building, token exchange) + endpoint discovery + session/ # SessionManager lifecycle, semaphore, event/token waiter channels + +pkg/ + oauth/ # PKCE: GeneratePKCEChallenge(), GenerateState() +``` + +--- + +## Package Responsibilities + +### `cmd/main.go` +Service entry point. Reads environment variables, constructs components, wires routes, runs HTTP server, handles graceful shutdown. + +### `internal/core` +- **`interfaces.go`**: `SessionStore`, `TokenCache`, `OAuthDiscoverer`, `Clock`, `Semaphore` interfaces; `Session`, `OAuthTransaction`, `Event`, `TokenResult` types. +- **`broker.go`**: `TokenBroker` — coordinates the full token acquisition flow: cache check → semaphore → OAuth discovery → PKCE generation → auth URL → event → wait for callback → token exchange → cache → unblock waiters. + +### `internal/api` +HTTP handlers. Registers routes on a chi router. Two error formats: flat for the AuthBridge API (`/sessions/token`), nested for the Backend API (all other endpoints). + +### `internal/oauth` +- **`discovery.go`**: `Discoverer` — calls `GET /.well-known/oauth-protected-resource` on the resource server; supports global and per-resource config overrides that skip discovery. +- **`client.go`**: `Client` — `BuildAuthorizationURL()` constructs the OAuth redirect URL with PKCE parameters; `ExchangeToken()` calls the OAuth provider's token endpoint directly. + +### `internal/session` +- **`manager.go`**: `SessionManager` — creates, validates, retrieves, and expires sessions; owns idle timeout timers; notifies token and event waiters on session end. +- **`semaphore.go`**: `SimpleSemaphore` — channel-based semaphore used to enforce one active OAuth flow per session. + +### `internal/cache` +- **`token_cache.go`**: `TokenCache` — stores access tokens keyed by `(session_key, resource_url)`; checks expiry and near-expiry (< 5 min) on retrieval. +- **`jwt_parser.go`**: Extracts the `exp` claim from a JWT to derive cache TTL. Non-JWT tokens are treated as long-lived (1 year). + +### `internal/auth` +JWT claim extraction (`sub`, `session_uid`/`jti`) without signature validation — used as a fallback when `JWT_JWKS_URL` is not configured. When it is configured, signature validation is handled by the `authlib` JWKS verifier. + +### `pkg/oauth` +`GeneratePKCEChallenge()` — generates a random code verifier and its SHA256/base64url challenge (RFC 7636 S256 method). +`GenerateState()` — generates a cryptographically random state nonce. + +--- + +## Token Acquisition Flow + +When `POST /sessions/token` is received and no cached token exists: + +1. Validate session ownership (user in JWT must own the session). +2. Check token cache — return immediately if found. +3. Acquire per-session semaphore (only one OAuth flow per session at a time). +4. Recheck token cache (another request may have completed while waiting). +5. Discover OAuth endpoints from resource server's `.well-known/oauth-protected-resource` (or use configured overrides). +6. Generate PKCE (`code_verifier`, `code_challenge`) and state (`sessionKey.nonce`) at the Token Broker. +7. Build authorization URL with Token Broker's `callback_url`, PKCE, and discovered scopes. +8. Publish `oauth_url_ready` event — Backend receives this via long-poll and opens the OAuth consent window. +9. Wait for the OAuth provider to call `GET /oauth/callback` with `code` and `state`. +10. Exchange `code` + `code_verifier` directly with the OAuth provider's token endpoint. +11. Cache the resulting access token. +12. Unblock any other waiters for the same `(session, resource_url)`. +13. Return token to the AuthBridge caller. + +### Key security properties +- `code_verifier` never leaves the Token Broker. +- `client_secret` is held only by the Token Broker. +- Authorization code goes directly from OAuth provider to Token Broker (not through Backend). +- Session ownership is validated before the cache is checked — prevents session hijacking. + +--- + +## Concurrency Model + +- **Per-session semaphore** (capacity 1): at most one OAuth flow runs per session at a time. +- **Double-checked locking**: cache is checked before *and* after semaphore acquisition. +- **`session.Done` channel**: closed when a session ends; unblocks any goroutine waiting on OAuth completion so it returns promptly with a "session ended" error. +- **`session.EventWaiters` channel** (buffered, cap 1): delivers one event per long-poll response. + +--- + +## Session Lifecycle + +``` +1. POST /sessions → session created, timer not yet started +2. POST /sessions/broker-events → timer reset while Backend is connected +3. request completes → idle timer starts (default 60s) +4. next poll arrives → timer reset +5. timer fires → session expired, all waiters unblocked + OR + POST /sessions/end → session terminated explicitly +``` + +--- + +## Error Response Formats + +Two formats are used to match the consumers: + +**AuthBridge API** (`POST /sessions/token`) — flat: +```json +{ "code": "error_code", "message": "description" } +``` + +**Backend API** (all other endpoints) — nested: +```json +{ "error": { "code": "error_code", "message": "description" } } +``` diff --git a/token-broker/go.mod b/token-broker/go.mod new file mode 100644 index 00000000..0ea5ff9e --- /dev/null +++ b/token-broker/go.mod @@ -0,0 +1,27 @@ +module github.com/kagenti/token-broker + +go 1.26.0 + +require ( + github.com/go-chi/chi/v5 v5.3.0 + github.com/google/uuid v1.6.0 + github.com/kagenti/kagenti-extensions/authbridge/authlib v0.0.0-20260619001334-ce3417655ee8 + github.com/stretchr/testify v1.10.0 +) + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0 // indirect + github.com/goccy/go-json v0.10.3 // indirect + github.com/lestrrat-go/blackmagic v1.0.3 // indirect + github.com/lestrrat-go/httpcc v1.0.1 // indirect + github.com/lestrrat-go/httprc v1.0.6 // indirect + github.com/lestrrat-go/iter v1.0.2 // indirect + github.com/lestrrat-go/jwx/v2 v2.1.6 // indirect + github.com/lestrrat-go/option v1.0.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/segmentio/asm v1.2.0 // indirect + golang.org/x/crypto v0.53.0 // indirect + golang.org/x/sys v0.46.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/token-broker/go.sum b/token-broker/go.sum new file mode 100644 index 00000000..7bcbc58f --- /dev/null +++ b/token-broker/go.sum @@ -0,0 +1,43 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0 h1:NMZiJj8QnKe1LgsbDayM4UoHwbvwDRwnI3hwNaAHRnc= +github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0/go.mod h1:ZXNYxsqcloTdSy/rNShjYzMhyjf0LaoftYK0p+A3h40= +github.com/go-chi/chi/v5 v5.3.0 h1:halUjDxhshgXHMrao5bB8eNBXo/rnzwr8m5m36glehM= +github.com/go-chi/chi/v5 v5.3.0/go.mod h1:R+tYY2hNuVUUjxoPtqUdgBqevM9s9njzkTLutVsOCto= +github.com/goccy/go-json v0.10.3 h1:KZ5WoDbxAIgm2HNbYckL0se1fHD6rz5j4ywS6ebzDqA= +github.com/goccy/go-json v0.10.3/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/kagenti/kagenti-extensions/authbridge/authlib v0.0.0-20260619001334-ce3417655ee8 h1:O5bcG8b5QGbIOGUVOFwrEBHso5pwmaRWg9UNC8OtxFI= +github.com/kagenti/kagenti-extensions/authbridge/authlib v0.0.0-20260619001334-ce3417655ee8/go.mod h1:C9cgtHUtuGLTO/dl0xbXL2zflPwopI2SE0UqDf5C/lA= +github.com/lestrrat-go/blackmagic v1.0.3 h1:94HXkVLxkZO9vJI/w2u1T0DAoprShFd13xtnSINtDWs= +github.com/lestrrat-go/blackmagic v1.0.3/go.mod h1:6AWFyKNNj0zEXQYfTMPfZrAXUWUfTIZ5ECEUEJaijtw= +github.com/lestrrat-go/httpcc v1.0.1 h1:ydWCStUeJLkpYyjLDHihupbn2tYmZ7m22BGkcvZZrIE= +github.com/lestrrat-go/httpcc v1.0.1/go.mod h1:qiltp3Mt56+55GPVCbTdM9MlqhvzyuL6W/NMDA8vA5E= +github.com/lestrrat-go/httprc v1.0.6 h1:qgmgIRhpvBqexMJjA/PmwSvhNk679oqD1RbovdCGW8k= +github.com/lestrrat-go/httprc v1.0.6/go.mod h1:mwwz3JMTPBjHUkkDv/IGJ39aALInZLrhBp0X7KGUZlo= +github.com/lestrrat-go/iter v1.0.2 h1:gMXo1q4c2pHmC3dn8LzRhJfP1ceCbgSiT9lUydIzltI= +github.com/lestrrat-go/iter v1.0.2/go.mod h1:Momfcq3AnRlRjI5b5O8/G5/BvpzrhoFTZcn06fEOPt4= +github.com/lestrrat-go/jwx/v2 v2.1.6 h1:hxM1gfDILk/l5ylers6BX/Eq1m/pnxe9NBwW6lVfecA= +github.com/lestrrat-go/jwx/v2 v2.1.6/go.mod h1:Y722kU5r/8mV7fYDifjug0r8FK8mZdw0K0GpJw/l8pU= +github.com/lestrrat-go/option v1.0.1 h1:oAzP2fvZGQKWkvHa1/SAcFolBEca1oN+mQ7eooNBEYU= +github.com/lestrrat-go/option v1.0.1/go.mod h1:5ZHFbivi4xwXxhxY9XHDe2FHo6/Z7WWmtT7T5nBBp3I= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/segmentio/asm v1.2.0 h1:9BQrFxC+YOHJlTlHGkTrFWf59nbL3XnCoFLTwDCI7ys= +github.com/segmentio/asm v1.2.0/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= +golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/token-broker/internal/api/handlers.go b/token-broker/internal/api/handlers.go new file mode 100644 index 00000000..eecf2d30 --- /dev/null +++ b/token-broker/internal/api/handlers.go @@ -0,0 +1,530 @@ +package api + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "net/url" + + "github.com/go-chi/chi/v5" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/jwtvalidation/validation" + "github.com/kagenti/token-broker/internal/auth" + "github.com/kagenti/token-broker/internal/core" +) + +// TokenBroker defines the interface for token acquisition operations. +type TokenBroker interface { + AcquireToken(ctx context.Context, sessionKey, userID, resourceURL string) (string, error) + CompleteOAuth(sessionKey, userID, code, state string) error + GetSessionStore() core.SessionStore +} + +// Handler handles HTTP requests for the Token Broker API. +type Handler struct { + broker TokenBroker + sessionStore core.SessionStore + logger *slog.Logger + allowedRedirectHosts []string + jwtVerifier validation.Verifier // nil = fallback to no-validation mode + jwtAudiences []string +} + +// NewHandler creates a new API handler. +func NewHandler(broker TokenBroker, sessionStore core.SessionStore, logger *slog.Logger, allowedRedirectHosts []string) *Handler { + return &Handler{ + broker: broker, + sessionStore: sessionStore, + logger: logger, + allowedRedirectHosts: allowedRedirectHosts, + } +} + +// WithJWTVerifier sets the JWT verifier and expected audiences on the handler. +func (h *Handler) WithJWTVerifier(verifier validation.Verifier, audiences []string) *Handler { + h.jwtVerifier = verifier + h.jwtAudiences = audiences + return h +} + +// RegisterRoutes registers all API routes. +func (h *Handler) RegisterRoutes(r chi.Router) { + r.Post("/sessions", h.HandleCreateSession) + r.Post("/sessions/token", h.HandleGetToken) // AuthBridge API endpoint + r.Post("/sessions/broker-events", h.HandleEvents) // Backend API endpoint (JWT auth) + r.Post("/sessions/end", h.HandleEndSession) // Backend API endpoint (JWT auth) + r.Get("/oauth/callback", h.HandleOAuthCallback) // OAuth provider callback +} + +// parseAndValidateJWT extracts and validates the JWT from the Authorization header. +// When a JWKS verifier is configured it performs full signature/expiry/issuer/audience +// validation. When no verifier is configured it falls back to claim-only parsing (dev mode). +// Returns (userID, sessionKey) on success, or writes an error response and returns ("", ""). +func (h *Handler) parseAndValidateJWT(w http.ResponseWriter, r *http.Request, errWriter func(http.ResponseWriter, int, string, string)) (userID, sessionKey string) { + authHeader := r.Header.Get("Authorization") + tokenStr, err := auth.ExtractBearerToken(authHeader) + if err != nil { + h.logger.Warn("Failed to extract bearer token", "error", err) + errWriter(w, http.StatusUnauthorized, "invalid_token", "Missing or invalid Authorization header") + return "", "" + } + + if h.jwtVerifier != nil { + claims, err := h.jwtVerifier.Verify(r.Context(), tokenStr, h.jwtAudiences) + if err != nil { + h.logger.Warn("JWT validation failed", "error", err) + errWriter(w, http.StatusUnauthorized, "invalid_token", fmt.Sprintf("JWT validation failed: %v", err)) + return "", "" + } + // Extract session key from Extra map (session_uid preferred, jti as fallback). + sk, _ := claims.Extra["session_uid"].(string) + if sk == "" { + sk, _ = claims.Extra["jti"].(string) + } + if claims.Subject == "" || sk == "" { + h.logger.Warn("JWT missing required claims", "has_sub", claims.Subject != "", "has_session_key", sk != "") + errWriter(w, http.StatusUnauthorized, "invalid_token", "JWT missing required sub or session_uid/jti claims") + return "", "" + } + return claims.Subject, sk + } + + // Fallback: no-validation mode (dev/test only — WARNING logged at startup). + legacyClaims, err := auth.ParseJWTWithoutValidation(tokenStr) + if err != nil { + h.logger.Warn("Failed to parse JWT", "error", err) + errWriter(w, http.StatusUnauthorized, "invalid_token", fmt.Sprintf("Invalid JWT: %v", err)) + return "", "" + } + return legacyClaims.Sub, legacyClaims.GetSessionKey() +} + +// ErrorResponse represents a standardized error response (flat format). +type ErrorResponse struct { + Code string `json:"code"` + Message string `json:"message"` +} + +// BackendErrorResponse represents the nested error format used by Backend endpoints. +type BackendErrorResponse struct { + Error ErrorDetail `json:"error"` +} + +// ErrorDetail contains error details (for Backend error format). +type ErrorDetail struct { + Code string `json:"code"` + Message string `json:"message"` +} + +// writeJSON writes a JSON response. +func writeJSON(w http.ResponseWriter, status int, data interface{}) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(data); err != nil { + slog.Default().Error("failed to encode JSON response", "error", err) + } +} + +// writeError writes a standardized error response (flat format for AuthBridge API). +func writeError(w http.ResponseWriter, status int, code, message string) { + writeJSON(w, status, ErrorResponse{ + Code: code, + Message: message, + }) +} + +// writeBackendError writes a standardized error response (nested format for Backend API). +func writeBackendError(w http.ResponseWriter, status int, code, message string) { + writeJSON(w, status, BackendErrorResponse{ + Error: ErrorDetail{ + Code: code, + Message: message, + }, + }) +} + +// CreateSessionRequest represents the request body for POST /sessions +type CreateSessionRequest struct { + BackendSessionRedirectURL string `json:"backend_session_redirect_url"` +} + +// HandleCreateSession handles POST /sessions +func (h *Handler) HandleCreateSession(w http.ResponseWriter, r *http.Request) { + userID, sessionKey := h.parseAndValidateJWT(w, r, writeBackendError) + if userID == "" { + return + } + + // Parse request body for backend redirect URL + var req CreateSessionRequest + if r.Body != nil { + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + h.logger.Warn("Failed to parse request body", "error", err) + writeBackendError(w, http.StatusBadRequest, "invalid_request", "Invalid request body") + return + } + } + + // Validate backend redirect URL against allowed hosts. + if req.BackendSessionRedirectURL != "" && len(h.allowedRedirectHosts) > 0 { + parsed, err := url.Parse(req.BackendSessionRedirectURL) + if err != nil || !h.isAllowedRedirectHost(parsed.Hostname()) { + h.logger.Warn("Rejected backend_session_redirect_url: host not in allowlist", + "host", func() string { + if err != nil { + return "" + } + return parsed.Hostname() + }()) + writeBackendError(w, http.StatusBadRequest, "invalid_request", "backend_session_redirect_url host is not permitted") + return + } + } + + h.logger.Info("Creating session", + "user_id", userID, + "session_key", sessionKey, + "backend_redirect_url", req.BackendSessionRedirectURL) + + // Create session with provided session key and redirect URL + err := h.sessionStore.CreateSession(sessionKey, userID, req.BackendSessionRedirectURL) + if err != nil { + if err.Error() == "max sessions per user exceeded" { + writeBackendError(w, http.StatusTooManyRequests, "too_many_sessions", err.Error()) + return + } + h.logger.Error("Failed to create session", "error", err, "user_id", userID, "session_key", sessionKey) + writeBackendError(w, http.StatusInternalServerError, "internal_error", "Failed to create session") + return + } + + h.logger.Info("Session created", "session_key", sessionKey, "user_id", userID) + + // Return success (client already knows the session key from JWT jti) + w.WriteHeader(http.StatusCreated) +} + +// HandleGetToken handles POST /sessions/token (new API with JWT authentication). +func (h *Handler) HandleGetToken(w http.ResponseWriter, r *http.Request) { + userID, sessionKey := h.parseAndValidateJWT(w, r, writeError) + if userID == "" { + return + } + + // Extract resource URL from header. + // This is set by the authbridge sidecar from operator-configured routes — trusted infrastructure. + resourceURL := r.Header.Get("X-Server-Url") + if resourceURL == "" { + writeError(w, http.StatusBadRequest, "invalid_request", "Missing X-Server-Url header") + return + } + + h.logger.Info("Token request (new API)", + "session_key", sessionKey, + "user_id", userID, + "resource_url", resourceURL) + + // Validate session + if err := h.sessionStore.ValidateSession(sessionKey, userID); err != nil { + h.logger.Warn("Session validation failed", + "session_key", sessionKey, + "user_id", userID, + "error", err) + writeError(w, http.StatusUnauthorized, "unauthorized", "Invalid session or user mismatch") + return + } + + // Acquire token (blocks until available or timeout) + ctx := r.Context() + accessToken, err := h.broker.AcquireToken(ctx, sessionKey, userID, resourceURL) + if err != nil { + h.logger.Error("Token acquisition failed", + "session_key", sessionKey, + "user_id", userID, + "resource_url", resourceURL, + "error", err) + + // Classify the error with errors.Is so that wrapping (AcquireToken wraps + // every performOAuthFlow error) does not collapse a specific, recoverable + // condition into a generic 5xx. A 5xx makes agents/sidecars mark this route + // dead and stop retrying, so it is reserved for genuinely retryable + // server/upstream failures — never session/client state. + switch { + case errors.Is(err, core.ErrOAuthTimeout): + // User did not finish the OAuth flow in time — transient, retryable. + writeError(w, http.StatusRequestTimeout, "timeout", "OAuth flow did not complete in time") + case errors.Is(err, core.ErrSessionEnded), errors.Is(err, core.ErrSessionExpired): + // Session gone — permanent for this session; client must re-establish. + writeError(w, http.StatusUnauthorized, "session_expired", "Session ended or expired") + case errors.Is(err, core.ErrOAuthUpstream): + // Discovery / token exchange failed upstream — transient, retryable. + writeError(w, http.StatusBadGateway, "oauth_upstream", "Upstream OAuth provider error") + default: + // Unknown: do not signal a permanently-dead route. Retryable server error. + writeError(w, http.StatusInternalServerError, "internal_error", "Failed to obtain token") + } + return + } + + h.logger.Info("Token acquired (new API)", + "session_key", sessionKey, + "user_id", userID, + "resource_url", resourceURL) + + // Return token + writeJSON(w, http.StatusOK, map[string]string{ + "token": accessToken, + }) +} + +// HandleEvents handles POST /sessions/events +func (h *Handler) HandleEvents(w http.ResponseWriter, r *http.Request) { + userID, sessionKey := h.parseAndValidateJWT(w, r, writeBackendError) + if userID == "" { + return + } + + // Long-poll for events + h.handleEventLongPoll(w, r, sessionKey, userID) +} + +// handleEventLongPoll handles long-polling for events. +func (h *Handler) handleEventLongPoll(w http.ResponseWriter, r *http.Request, sessionKey, userID string) { + h.logger.Debug("Event long-poll started", + "session_key", sessionKey, + "user_id", userID) + + // Validate session + if err := h.sessionStore.ValidateSession(sessionKey, userID); err != nil { + h.logger.Warn("Session validation failed for event poll", + "session_key", sessionKey, + "user_id", userID, + "error", err) + writeBackendError(w, http.StatusUnauthorized, "unauthorized", "Invalid session or user mismatch") + return + } + + // Reset session timer (Backend is connected) + h.sessionStore.ResetSessionTimer(sessionKey) + + // Get session + session, err := h.sessionStore.GetSession(sessionKey) + if err != nil { + writeBackendError(w, http.StatusUnauthorized, "session_not_found", "Session not found") + return + } + + // When we return, start the session timer + defer h.sessionStore.StartSessionTimer(sessionKey) + + select { + case event := <-session.EventWaiters: + h.logger.Info("Event sent to Backend", + "session_key", sessionKey, + "event_type", event.Type) + writeJSON(w, http.StatusOK, event) + + case <-session.Done: + h.logger.Info("Session closed while waiting for event", + "session_key", sessionKey, + "user_id", userID) + writeJSON(w, http.StatusOK, core.Event{ + Type: "error", + Message: "Session expired", + Code: "session_expired", + }) + + case <-r.Context().Done(): + h.logger.Debug("Event long-poll canceled by client", + "session_key", sessionKey, + "reason", r.Context().Err()) + return + } +} + +// HandleOAuthCallback handles GET /oauth/callback - receives OAuth callbacks from OAuth provider +func (h *Handler) HandleOAuthCallback(w http.ResponseWriter, r *http.Request) { + // Log the full callback request for debugging + h.logger.Info("OAuth callback received from provider", + "method", r.Method, + "url", r.URL.String(), + "remote_addr", r.RemoteAddr, + "user_agent", r.Header.Get("User-Agent"), + "referer", r.Header.Get("Referer")) + + // Extract code and state from query parameters + code := r.URL.Query().Get("code") + state := r.URL.Query().Get("state") + errorParam := r.URL.Query().Get("error") + errorDescription := r.URL.Query().Get("error_description") + + h.logger.Info("OAuth callback parameters", + "has_code", code != "", + "code_length", len(code), + "has_state", state != "", + "state_length", len(state), + "has_error", errorParam != "", + "error", errorParam, + "error_description", errorDescription) + + // Check for OAuth provider errors + if errorParam != "" { + h.logger.Warn("OAuth provider returned error", + "error", errorParam, + "description", errorDescription) + + // Look up session by state to get redirect URL + if state != "" { + if session, err := h.sessionStore.GetSessionByState(state); err == nil && session.BackendSessionRedirectURL != "" { + // Redirect to backend with error parameters + redirectURL := h.addQueryParams(session.BackendSessionRedirectURL, map[string]string{ + "oauth_status": "error", + "error": errorParam, + "error_description": errorDescription, + }) + http.Redirect(w, r, redirectURL, http.StatusFound) + return + } + } + http.Error(w, fmt.Sprintf("OAuth error: %s - %s", errorParam, errorDescription), http.StatusBadRequest) + return + } + + if code == "" || state == "" { + h.logger.Warn("OAuth callback missing required parameters", + "has_code", code != "", + "has_state", state != "") + http.Error(w, "Missing code or state parameter", http.StatusBadRequest) + return + } + + h.logger.Info("OAuth callback received", + "code_length", len(code), + "state_length", len(state)) + + // Look up session by state parameter + session, err := h.sessionStore.GetSessionByState(state) + if err != nil { + h.logger.Error("Failed to find session for OAuth callback", + "state_length", len(state), + "error", err) + http.Error(w, "Invalid or expired OAuth session", http.StatusBadRequest) + return + } + + h.logger.Info("Session found for OAuth callback", + "session_key", session.SessionKey, + "user_id", session.UserID) + + // Complete OAuth flow + if err := h.broker.CompleteOAuth(session.SessionKey, session.UserID, code, state); err != nil { + h.logger.Error("Failed to complete OAuth", + "session_key", session.SessionKey, + "user_id", session.UserID, + "error", err) + + // Redirect to backend with error + if session.BackendSessionRedirectURL != "" { + redirectURL := h.addQueryParams(session.BackendSessionRedirectURL, map[string]string{ + "oauth_status": "error", + "error": "token_exchange_failed", + "error_description": "Failed to exchange authorization code for access token", + }) + http.Redirect(w, r, redirectURL, http.StatusFound) + return + } + + http.Error(w, "Failed to complete OAuth flow", http.StatusInternalServerError) + return + } + + h.logger.Info("OAuth completion processed", + "session_key", session.SessionKey) + + // Redirect user to Backend's session redirect URL with success status + if session.BackendSessionRedirectURL == "" { + h.logger.Warn("No backend redirect URL configured for session", + "session_key", session.SessionKey) + w.WriteHeader(http.StatusOK) + if _, err := w.Write([]byte("OAuth authorization successful. You may close this window.")); err != nil { + h.logger.Warn("Failed to write OAuth success response", "error", err) + } + return + } + + // Add success status to redirect URL + redirectURL := h.addQueryParams(session.BackendSessionRedirectURL, map[string]string{ + "oauth_status": "success", + }) + + h.logger.Info("Redirecting user to backend", + "session_key", session.SessionKey, + "redirect_url", redirectURL) + + http.Redirect(w, r, redirectURL, http.StatusFound) +} + +// addQueryParams adds query parameters to a URL, preserving existing parameters +func (h *Handler) addQueryParams(baseURL string, params map[string]string) string { + u, err := url.Parse(baseURL) + if err != nil { + h.logger.Warn("Failed to parse redirect URL", "url", baseURL, "error", err) + return baseURL + } + + q := u.Query() + for key, value := range params { + q.Set(key, value) + } + u.RawQuery = q.Encode() + + return u.String() +} + +// HandleEndSession handles POST /sessions/end +func (h *Handler) HandleEndSession(w http.ResponseWriter, r *http.Request) { + userID, sessionKey := h.parseAndValidateJWT(w, r, writeBackendError) + if userID == "" { + return + } + + h.logger.Info("Ending session", + "session_key", sessionKey, + "user_id", userID) + + // Validate session + if err := h.sessionStore.ValidateSession(sessionKey, userID); err != nil { + h.logger.Warn("Session validation failed for end session", + "session_key", sessionKey, + "user_id", userID, + "error", err) + writeBackendError(w, http.StatusUnauthorized, "unauthorized", "Invalid session or user mismatch") + return + } + + // End session + if err := h.sessionStore.EndSession(sessionKey); err != nil { + h.logger.Error("Failed to end session", + "session_key", sessionKey, + "error", err) + writeBackendError(w, http.StatusInternalServerError, "internal_error", "Failed to end session") + return + } + + h.logger.Info("Session ended", + "session_key", sessionKey) + + w.WriteHeader(http.StatusOK) +} + +// isAllowedRedirectHost reports whether host is in the configured allowlist. +func (h *Handler) isAllowedRedirectHost(host string) bool { + for _, allowed := range h.allowedRedirectHosts { + if allowed == host { + return true + } + } + return false +} diff --git a/token-broker/internal/api/handlers_test.go b/token-broker/internal/api/handlers_test.go new file mode 100644 index 00000000..78e9810f --- /dev/null +++ b/token-broker/internal/api/handlers_test.go @@ -0,0 +1,542 @@ +package api + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "testing" + "time" + + "github.com/kagenti/token-broker/internal/core" +) + +// mockBroker implements a mock TokenBroker for testing +type mockBroker struct { + acquireTokenFunc func(ctx context.Context, sessionKey, userID, resourceURL string) (string, error) +} + +func (m *mockBroker) AcquireToken(ctx context.Context, sessionKey, userID, resourceURL string) (string, error) { + if m.acquireTokenFunc != nil { + return m.acquireTokenFunc(ctx, sessionKey, userID, resourceURL) + } + return "mock-token", nil +} + +func (m *mockBroker) CompleteOAuth(sessionKey, userID, code, state string) error { + return nil +} + +func (m *mockBroker) GetSessionStore() core.SessionStore { + return &mockSessionStore{} +} + +// mockSessionStore implements a mock SessionStore for testing +type mockSessionStore struct { + validateSessionFunc func(sessionKey, userID string) error + getSessionFunc func(sessionKey string) (*core.Session, error) + getSessionByStateFunc func(state string) (*core.Session, error) +} + +func (m *mockSessionStore) CreateSession(sessionKey, userID, backendRedirectURL string) error { + return nil +} + +func (m *mockSessionStore) GetSession(sessionKey string) (*core.Session, error) { + if m.getSessionFunc != nil { + return m.getSessionFunc(sessionKey) + } + + return &core.Session{ + SessionKey: sessionKey, + UserID: "user123", + CreatedAt: time.Now(), + }, nil +} + +func (m *mockSessionStore) GetSessionByState(state string) (*core.Session, error) { + if m.getSessionByStateFunc != nil { + return m.getSessionByStateFunc(state) + } + return &core.Session{ + SessionKey: "mock-session", + UserID: "user123", + CreatedAt: time.Now(), + }, nil +} + +func (m *mockSessionStore) ValidateSession(sessionKey, userID string) error { + if m.validateSessionFunc != nil { + return m.validateSessionFunc(sessionKey, userID) + } + return nil +} + +func (m *mockSessionStore) EndSession(sessionKey string) error { + return nil +} + +func (m *mockSessionStore) ExpireSession(sessionKey string) {} + +func (m *mockSessionStore) ResetSessionTimer(sessionKey string) {} + +func (m *mockSessionStore) StartSessionTimer(sessionKey string) {} + +func TestHandleGetToken_NewAPI(t *testing.T) { + tests := []struct { + name string + authHeader string + serverURLHeader string + authEndpointHeader string + tokenEndpointHeader string + wantStatus int + wantToken string + wantErrorCode string + }{ + { + name: "valid request with JWT", + authHeader: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIiwianRpIjoic2Vzc2lvbjQ1NiJ9.signature", // #notsecret + serverURLHeader: "https://mcp.example.com", + wantStatus: http.StatusOK, + wantToken: "mock-token", + }, + { + name: "valid request with OAuth endpoints", + authHeader: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIiwianRpIjoic2Vzc2lvbjQ1NiJ9.signature", // #notsecret + serverURLHeader: "https://mcp.example.com", + authEndpointHeader: "https://auth.example.com/authorize", + tokenEndpointHeader: "https://auth.example.com/token", + wantStatus: http.StatusOK, + wantToken: "mock-token", + }, + { + name: "missing Authorization header", + serverURLHeader: "https://mcp.example.com", + wantStatus: http.StatusUnauthorized, + wantErrorCode: "invalid_token", + }, + { + name: "invalid Authorization header format", + authHeader: "Basic dXNlcjpwYXNz", + serverURLHeader: "https://mcp.example.com", + wantStatus: http.StatusUnauthorized, + wantErrorCode: "invalid_token", + }, + { + name: "missing X-Server-Url header", + authHeader: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIiwianRpIjoic2Vzc2lvbjQ1NiJ9.signature", // #notsecret + wantStatus: http.StatusBadRequest, + wantErrorCode: "invalid_request", + }, + { + name: "invalid JWT - missing sub claim", + authHeader: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJqdGkiOiJzZXNzaW9uNDU2In0.signature", // #notsecret + serverURLHeader: "https://mcp.example.com", + wantStatus: http.StatusUnauthorized, + wantErrorCode: "invalid_token", + }, + { + name: "invalid JWT - missing jti claim", + authHeader: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIn0.signature", // #notsecret + serverURLHeader: "https://mcp.example.com", + wantStatus: http.StatusUnauthorized, + wantErrorCode: "invalid_token", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create mock broker + broker := &mockBroker{ + acquireTokenFunc: func(ctx context.Context, sessionKey, userID, resourceURL string) (string, error) { + // Verify parameters + if userID != "user123" { + t.Errorf("Expected userID 'user123', got '%s'", userID) + } + if sessionKey != "session456" { + t.Errorf("Expected sessionKey 'session456', got '%s'", sessionKey) + } + if resourceURL != "https://mcp.example.com" { + t.Errorf("Expected resourceURL 'https://mcp.example.com', got '%s'", resourceURL) + } + // Note: authEndpoint and tokenEndpoint are no longer passed to AcquireToken + // They are discovered internally by the broker + return "mock-token", nil + }, + } + + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + handler := NewHandler(broker, &mockSessionStore{}, logger, nil) + + // Create request + req := httptest.NewRequest(http.MethodPost, "/sessions/token", nil) + if tt.authHeader != "" { + req.Header.Set("Authorization", tt.authHeader) + } + if tt.serverURLHeader != "" { + req.Header.Set("X-Server-Url", tt.serverURLHeader) + } + if tt.authEndpointHeader != "" { + req.Header.Set("X-Authorization-Endpoint", tt.authEndpointHeader) + } + if tt.tokenEndpointHeader != "" { + req.Header.Set("X-Token-Endpoint", tt.tokenEndpointHeader) + } + + // Create response recorder + w := httptest.NewRecorder() + + // Call handler + handler.HandleGetToken(w, req) + + // Check status code + if w.Code != tt.wantStatus { + t.Errorf("Expected status %d, got %d", tt.wantStatus, w.Code) + } + + // Check response body + if tt.wantStatus == http.StatusOK { + var response map[string]string + if err := json.NewDecoder(w.Body).Decode(&response); err != nil { + t.Fatalf("Failed to decode response: %v", err) + } + if response["token"] != tt.wantToken { + t.Errorf("Expected token '%s', got '%s'", tt.wantToken, response["token"]) + } + } else { + var errorResponse ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&errorResponse); err != nil { + t.Fatalf("Failed to decode error response: %v", err) + } + if errorResponse.Code != tt.wantErrorCode { + t.Errorf("Expected error code '%s', got '%s'", tt.wantErrorCode, errorResponse.Code) + } + } + }) + } +} + +// TestHandleGetTokenErrorClassification verifies that HandleGetToken classifies +// broker errors by sentinel (via errors.Is) even when the broker wraps them, and +// never returns 503 — a 5xx makes agents mark the route dead and stop retrying. +func TestHandleGetTokenErrorClassification(t *testing.T) { + // A valid JWT for user123 / session456 so we reach the broker call. + const validJWT = "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIiwianRpIjoic2Vzc2lvbjQ1NiJ9.signature" // #notsecret + + tests := []struct { + name string + brokerErr error + wantStatus int + wantErrorCode string + }{ + { + // The reviewer's regression: session ended, but wrapped by AcquireToken. + name: "wrapped session ended -> 401", + brokerErr: fmt.Errorf("OAuth flow failed: %w", core.ErrSessionEnded), + wantStatus: http.StatusUnauthorized, + wantErrorCode: "session_expired", + }, + { + name: "wrapped session expired -> 401", + brokerErr: fmt.Errorf("OAuth flow failed: %w", core.ErrSessionExpired), + wantStatus: http.StatusUnauthorized, + wantErrorCode: "session_expired", + }, + { + name: "wrapped oauth timeout -> 408", + brokerErr: fmt.Errorf("OAuth flow failed: %w", core.ErrOAuthTimeout), + wantStatus: http.StatusRequestTimeout, + wantErrorCode: "timeout", + }, + { + name: "wrapped upstream failure -> 502", + brokerErr: fmt.Errorf("OAuth flow failed: %w", core.ErrOAuthUpstream), + wantStatus: http.StatusBadGateway, + wantErrorCode: "oauth_upstream", + }, + { + name: "unknown error -> 500 (not 503)", + brokerErr: errors.New("something unexpected"), + wantStatus: http.StatusInternalServerError, + wantErrorCode: "internal_error", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + broker := &mockBroker{ + acquireTokenFunc: func(ctx context.Context, sessionKey, userID, resourceURL string) (string, error) { + return "", tt.brokerErr + }, + } + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + handler := NewHandler(broker, &mockSessionStore{}, logger, nil) + + req := httptest.NewRequest(http.MethodPost, "/sessions/token", nil) + req.Header.Set("Authorization", validJWT) + req.Header.Set("X-Server-Url", "https://mcp.example.com") + + w := httptest.NewRecorder() + handler.HandleGetToken(w, req) + + if w.Code != tt.wantStatus { + t.Errorf("Expected status %d, got %d", tt.wantStatus, w.Code) + } + if w.Code == http.StatusServiceUnavailable { + t.Errorf("Handler returned 503; agents treat 5xx as a dead route and stop retrying") + } + + var errorResponse ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&errorResponse); err != nil { + t.Fatalf("Failed to decode error response: %v", err) + } + if errorResponse.Code != tt.wantErrorCode { + t.Errorf("Expected error code '%s', got '%s'", tt.wantErrorCode, errorResponse.Code) + } + }) + } +} + +func TestErrorResponseFormats(t *testing.T) { + t.Run("AuthBridge API uses flat error format", func(t *testing.T) { + w := httptest.NewRecorder() + writeError(w, http.StatusBadRequest, "test_code", "test message") + + var response ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&response); err != nil { + t.Fatalf("Failed to decode response: %v", err) + } + + if response.Code != "test_code" { + t.Errorf("Expected code 'test_code', got '%s'", response.Code) + } + if response.Message != "test message" { + t.Errorf("Expected message 'test message', got '%s'", response.Message) + } + }) + + t.Run("Backend API uses nested error format", func(t *testing.T) { + w := httptest.NewRecorder() + writeBackendError(w, http.StatusBadRequest, "test_code", "test message") + + var response BackendErrorResponse + if err := json.NewDecoder(w.Body).Decode(&response); err != nil { + t.Fatalf("Failed to decode response: %v", err) + } + + if response.Error.Code != "test_code" { + t.Errorf("Expected code 'test_code', got '%s'", response.Error.Code) + } + if response.Error.Message != "test message" { + t.Errorf("Expected message 'test message', got '%s'", response.Error.Message) + } + }) +} + +func TestHandleOAuthCallback(t *testing.T) { + tests := []struct { + name string + code string + state string + wantStatus int + wantLocation string + }{ + { + name: "valid OAuth callback", + code: "auth-code-123", + state: "session-key-456.nonce789", + wantStatus: http.StatusFound, + wantLocation: "https://backend.example.com/oauth/complete?oauth_status=success", + }, + { + name: "missing code parameter", + state: "session-key-456.nonce789", + wantStatus: http.StatusBadRequest, + }, + { + name: "missing state parameter", + code: "auth-code-123", + wantStatus: http.StatusBadRequest, + }, + { + name: "invalid state format (no dot)", + code: "auth-code-123", + state: "invalid-state-no-dot", + wantStatus: http.StatusBadRequest, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create mock broker + broker := &mockBroker{} + + // Create mock session store with custom GetSessionByState + sessionStore := &mockSessionStore{ + getSessionByStateFunc: func(state string) (*core.Session, error) { + if state == "" || state == "invalid-state-no-dot" { + return nil, errors.New("session not found") + } + return &core.Session{ + SessionKey: "session-key-456", + UserID: "user123", + BackendSessionRedirectURL: "https://backend.example.com/oauth/complete", + ActiveOAuthTx: &core.OAuthTransaction{ + State: state, + }, + }, nil + }, + } + + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + handler := NewHandler(broker, sessionStore, logger, nil) + + // Create request with query parameters + url := "/oauth/callback" + if tt.code != "" || tt.state != "" { + url += "?" + if tt.code != "" { + url += "code=" + tt.code + } + if tt.state != "" { + if tt.code != "" { + url += "&" + } + url += "state=" + tt.state + } + } + req := httptest.NewRequest(http.MethodGet, url, nil) + + // Create response recorder + w := httptest.NewRecorder() + + // Call handler + handler.HandleOAuthCallback(w, req) + + // Check status code + if w.Code != tt.wantStatus { + t.Errorf("Expected status %d, got %d", tt.wantStatus, w.Code) + } + + // Check redirect location for successful callback + if tt.wantStatus == http.StatusFound { + location := w.Header().Get("Location") + if location != tt.wantLocation { + t.Errorf("Expected Location header '%s', got '%s'", tt.wantLocation, location) + } + } + }) + } +} + +func TestHandleEvents_LongPollReturnsEvent(t *testing.T) { + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + + eventCh := make(chan core.Event, 1) + doneCh := make(chan struct{}) + + sessionStore := &mockSessionStore{ + validateSessionFunc: func(sessionKey, userID string) error { + if sessionKey != "session456" || userID != "user123" { + return errors.New("invalid session") + } + return nil + }, + getSessionFunc: func(sessionKey string) (*core.Session, error) { + return &core.Session{ + SessionKey: sessionKey, + UserID: "user123", + EventWaiters: eventCh, + Done: doneCh, + }, nil + }, + } + + broker := &mockBroker{} + handler := NewHandler(broker, sessionStore, logger, nil) + + req := httptest.NewRequest(http.MethodPost, "/sessions/broker-events", nil) + req.Header.Set("Authorization", "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIiwianRpIjoic2Vzc2lvbjQ1NiJ9.signature") // #notsecret + w := httptest.NewRecorder() + + go func() { + time.Sleep(10 * time.Millisecond) + eventCh <- core.Event{Type: "oauth_url_ready", AuthURL: "https://example.com/auth"} + }() + + handler.HandleEvents(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("Expected status %d, got %d", http.StatusOK, w.Code) + } + + var event core.Event + if err := json.NewDecoder(w.Body).Decode(&event); err != nil { + t.Fatalf("Failed to decode event: %v", err) + } + + if event.Type != "oauth_url_ready" { + t.Fatalf("Expected event type oauth_url_ready, got %s", event.Type) + } + if event.AuthURL != "https://example.com/auth" { + t.Fatalf("Expected auth_url https://example.com/auth, got %s", event.AuthURL) + } +} + +func TestHandleEvents_LongPollReturnsSessionClosedEvent(t *testing.T) { + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + + eventCh := make(chan core.Event, 1) + doneCh := make(chan struct{}) + + sessionStore := &mockSessionStore{ + validateSessionFunc: func(sessionKey, userID string) error { + if sessionKey != "session456" || userID != "user123" { + return errors.New("invalid session") + } + return nil + }, + getSessionFunc: func(sessionKey string) (*core.Session, error) { + return &core.Session{ + SessionKey: sessionKey, + UserID: "user123", + EventWaiters: eventCh, + Done: doneCh, + }, nil + }, + } + + broker := &mockBroker{} + handler := NewHandler(broker, sessionStore, logger, nil) + + req := httptest.NewRequest(http.MethodPost, "/sessions/broker-events", nil) + req.Header.Set("Authorization", "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIiwianRpIjoic2Vzc2lvbjQ1NiJ9.signature") // #notsecret + w := httptest.NewRecorder() + + go func() { + time.Sleep(10 * time.Millisecond) + close(doneCh) + }() + + handler.HandleEvents(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("Expected status %d, got %d", http.StatusOK, w.Code) + } + + var event core.Event + if err := json.NewDecoder(w.Body).Decode(&event); err != nil { + t.Fatalf("Failed to decode event: %v", err) + } + + if event.Type != "error" { + t.Fatalf("Expected event type error, got %s", event.Type) + } + if event.Code != "session_expired" { + t.Fatalf("Expected event code session_expired, got %s", event.Code) + } +} diff --git a/token-broker/internal/auth/jwt.go b/token-broker/internal/auth/jwt.go new file mode 100644 index 00000000..dbd14815 --- /dev/null +++ b/token-broker/internal/auth/jwt.go @@ -0,0 +1,149 @@ +package auth + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "strings" +) + +// Audience represents the JWT audience claim which can be either a string or array of strings. +// Per RFC 7519, the "aud" claim can be a single string or an array of strings. +type Audience []string + +// UnmarshalJSON implements custom unmarshaling to handle both string and array formats. +func (a *Audience) UnmarshalJSON(data []byte) error { + // Try to unmarshal as a string first + var single string + if err := json.Unmarshal(data, &single); err == nil { + *a = Audience{single} + return nil + } + + // If that fails, try to unmarshal as an array + var multiple []string + if err := json.Unmarshal(data, &multiple); err != nil { + return fmt.Errorf("audience must be a string or array of strings: %w", err) + } + *a = Audience(multiple) + return nil +} + +// Contains checks if the audience list contains a specific value. +func (a Audience) Contains(value string) bool { + for _, aud := range a { + if aud == value { + return true + } + } + return false +} + +// String returns the first audience value, or empty string if none. +func (a Audience) String() string { + if len(a) > 0 { + return a[0] + } + return "" +} + +// JWTClaims represents the JWT claims we extract from the token. +// Phase 1: We do NOT validate the JWT signature - we trust the claims. +// This is acceptable for demo/testing but NOT production-ready. +type JWTClaims struct { + Sub string `json:"sub"` // Subject - User ID + Jti string `json:"jti"` // JWT ID - Session Key (fallback) + SessionUID string `json:"session_uid"` // Session UID - Session Key (preferred) + Iss string `json:"iss"` // Issuer + Aud Audience `json:"aud"` // Audience (can be string or array) + Exp int64 `json:"exp"` // Expiration time + Iat int64 `json:"iat"` // Issued at +} + +// GetSessionKey returns the session key, preferring session_uid over jti. +func (c *JWTClaims) GetSessionKey() string { + if c.SessionUID != "" { + return c.SessionUID + } + return c.Jti +} + +// ParseJWTWithoutValidation parses a JWT token and extracts claims WITHOUT validating the signature. +// +// WARNING: This function does NOT validate the JWT signature. It trusts the claims as-is. +// This is acceptable for demo/testing environments but NOT production-ready. +// +// Security considerations: +// - An attacker could forge JWTs with arbitrary claims +// - Deploy in trusted network only +// - Use network policies to restrict Token Broker access +// +// Future work (Phase 2+): +// - Implement JWT signature validation using JWKS +// - Fetch public keys from IdP's /.well-known/openid-configuration +// - Validate issuer, audience, expiration +// +// Parameters: +// - token: The JWT token string (without "Bearer " prefix) +// +// Returns: +// - *JWTClaims: Parsed claims containing sub (user ID) and jti (session key) +// - error: If the token is malformed or missing required claims +func ParseJWTWithoutValidation(token string) (*JWTClaims, error) { + // JWT format: header.payload.signature + parts := strings.Split(token, ".") + if len(parts) != 3 { + return nil, fmt.Errorf("invalid JWT format: expected 3 parts, got %d", len(parts)) + } + + // Decode payload (second part) + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return nil, fmt.Errorf("failed to decode JWT payload: %w", err) + } + + // Parse claims + var claims JWTClaims + if err := json.Unmarshal(payload, &claims); err != nil { + return nil, fmt.Errorf("failed to parse JWT claims: %w", err) + } + + // Validate required claims are present + if claims.Sub == "" { + return nil, fmt.Errorf("missing required claim: sub (user ID)") + } + + // Require either session_uid or jti + if claims.SessionUID == "" && claims.Jti == "" { + return nil, fmt.Errorf("missing required claim: session_uid or jti (session key)") + } + + return &claims, nil +} + +// ExtractBearerToken extracts the token from an Authorization header. +// Expected format: "Bearer " +// +// Parameters: +// - authHeader: The Authorization header value +// +// Returns: +// - string: The extracted token (without "Bearer " prefix) +// - error: If the header is missing or malformed +func ExtractBearerToken(authHeader string) (string, error) { + if authHeader == "" { + return "", fmt.Errorf("missing Authorization header") + } + + const prefix = "Bearer " + if !strings.HasPrefix(authHeader, prefix) { + return "", fmt.Errorf("invalid Authorization header format: expected 'Bearer '") + } + + token := strings.TrimPrefix(authHeader, prefix) + if token == "" { + return "", fmt.Errorf("empty token in Authorization header") + } + + return token, nil +} diff --git a/token-broker/internal/auth/jwt_test.go b/token-broker/internal/auth/jwt_test.go new file mode 100644 index 00000000..f5144180 --- /dev/null +++ b/token-broker/internal/auth/jwt_test.go @@ -0,0 +1,252 @@ +package auth + +import ( + "encoding/base64" + "encoding/json" + "testing" +) + +func TestParseJWTWithoutValidation(t *testing.T) { + tests := []struct { + name string + token string + wantSub string + wantJti string + wantErr bool + errContains string + }{ + { + name: "valid JWT with required claims", + token: createTestJWT(t, map[string]interface{}{"sub": "user123", "jti": "session456"}), + wantSub: "user123", + wantJti: "session456", + wantErr: false, + }, + { + name: "valid JWT with all claims (string audience)", + token: createTestJWT(t, map[string]interface{}{"sub": "user123", "jti": "session456", "iss": "https://idp.example.com", "aud": "token-broker", "exp": 1234567890, "iat": 1234567800}), + wantSub: "user123", + wantJti: "session456", + wantErr: false, + }, + { + name: "valid JWT with array audience", + token: createTestJWT(t, map[string]interface{}{"sub": "user123", "jti": "session456", "iss": "https://idp.example.com", "aud": []string{"token-broker", "api-gateway"}, "exp": 1234567890, "iat": 1234567800}), + wantSub: "user123", + wantJti: "session456", + wantErr: false, + }, + { + name: "valid JWT with single-element array audience", + token: createTestJWT(t, map[string]interface{}{"sub": "user123", "jti": "session456", "aud": []string{"token-broker"}}), + wantSub: "user123", + wantJti: "session456", + wantErr: false, + }, + { + name: "valid JWT without audience claim", + token: createTestJWT(t, map[string]interface{}{"sub": "user123", "jti": "session456", "iss": "https://idp.example.com"}), + wantSub: "user123", + wantJti: "session456", + wantErr: false, + }, + { + name: "missing sub claim", + token: createTestJWT(t, map[string]interface{}{"jti": "session456"}), + wantErr: true, + errContains: "missing required claim: sub", + }, + { + name: "missing session_uid and jti claim", + token: createTestJWT(t, map[string]interface{}{"sub": "user123"}), + wantErr: true, + errContains: "missing required claim: session_uid or jti", + }, + { + name: "empty sub claim", + token: createTestJWT(t, map[string]interface{}{"sub": "", "jti": "session456"}), + wantErr: true, + errContains: "missing required claim: sub", + }, + { + name: "empty session_uid and jti claim", + token: createTestJWT(t, map[string]interface{}{"sub": "user123", "session_uid": "", "jti": ""}), + wantErr: true, + errContains: "missing required claim: session_uid or jti", + }, + { + name: "valid JWT with session_uid", + token: createTestJWT(t, map[string]interface{}{"sub": "user123", "session_uid": "session789"}), + wantSub: "user123", + wantJti: "session789", + }, + { + name: "invalid JWT format - too few parts", + token: "header.payload", + wantErr: true, + errContains: "invalid JWT format", + }, + { + name: "invalid JWT format - too many parts", + token: "header.payload.signature.extra", + wantErr: true, + errContains: "invalid JWT format", + }, + { + name: "invalid base64 encoding", + token: "header.invalid!!!.signature", + wantErr: true, + errContains: "failed to decode JWT payload", + }, + { + name: "invalid JSON in payload", + token: "header." + base64.RawURLEncoding.EncodeToString([]byte("{invalid json")) + ".signature", + wantErr: true, + errContains: "failed to parse JWT claims", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claims, err := ParseJWTWithoutValidation(tt.token) + + if tt.wantErr { + if err == nil { + t.Errorf("ParseJWTWithoutValidation() expected error, got nil") + return + } + if tt.errContains != "" && !contains(err.Error(), tt.errContains) { + t.Errorf("ParseJWTWithoutValidation() error = %v, want error containing %q", err, tt.errContains) + } + return + } + + if err != nil { + t.Errorf("ParseJWTWithoutValidation() unexpected error = %v", err) + return + } + + if claims.Sub != tt.wantSub { + t.Errorf("ParseJWTWithoutValidation() sub = %v, want %v", claims.Sub, tt.wantSub) + } + if claims.GetSessionKey() != tt.wantJti { + t.Errorf("ParseJWTWithoutValidation() session_key = %v, want %v", claims.GetSessionKey(), tt.wantJti) + } + }) + } +} + +func TestExtractBearerToken(t *testing.T) { + tests := []struct { + name string + authHeader string + wantToken string + wantErr bool + errContains string + }{ + { + name: "valid Bearer token", + authHeader: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIn0.signature", // #notsecret + wantToken: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIn0.signature", // #notsecret + wantErr: false, + }, + { + name: "missing Authorization header", + authHeader: "", + wantErr: true, + errContains: "missing Authorization header", + }, + { + name: "invalid format - no Bearer prefix", + authHeader: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIn0.signature", // #notsecret + wantErr: true, + errContains: "invalid Authorization header format", + }, + { + name: "invalid format - wrong prefix", + authHeader: "Basic dXNlcjpwYXNz", + wantErr: true, + errContains: "invalid Authorization header format", + }, + { + name: "empty token after Bearer", + authHeader: "Bearer ", + wantErr: true, + errContains: "empty token", + }, + { + name: "Bearer with extra spaces", + authHeader: "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIn0.signature", // #notsecret + wantToken: " eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyMTIzIn0.signature", // #notsecret + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + token, err := ExtractBearerToken(tt.authHeader) + + if tt.wantErr { + if err == nil { + t.Errorf("ExtractBearerToken() expected error, got nil") + return + } + if tt.errContains != "" && !contains(err.Error(), tt.errContains) { + t.Errorf("ExtractBearerToken() error = %v, want error containing %q", err, tt.errContains) + } + return + } + + if err != nil { + t.Errorf("ExtractBearerToken() unexpected error = %v", err) + return + } + + if token != tt.wantToken { + t.Errorf("ExtractBearerToken() = %v, want %v", token, tt.wantToken) + } + }) + } +} + +// Helper function to create a test JWT with given claims +func createTestJWT(t *testing.T, claims map[string]interface{}) string { + t.Helper() + + // Create header + header := map[string]interface{}{ + "alg": "HS256", + "typ": "JWT", + } + headerJSON, err := json.Marshal(header) + if err != nil { + t.Fatalf("failed to marshal header: %v", err) + } + headerEncoded := base64.RawURLEncoding.EncodeToString(headerJSON) + + // Create payload + payloadJSON, err := json.Marshal(claims) + if err != nil { + t.Fatalf("failed to marshal claims: %v", err) + } + payloadEncoded := base64.RawURLEncoding.EncodeToString(payloadJSON) + + // Create signature (dummy - we don't validate it) + signature := base64.RawURLEncoding.EncodeToString([]byte("dummy-signature")) + + return headerEncoded + "." + payloadEncoded + "." + signature +} + +// Helper function to check if a string contains a substring +func contains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(substr) == 0 || (len(s) > 0 && len(substr) > 0 && stringContains(s, substr))) +} + +func stringContains(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/token-broker/internal/cache/jwt_parser.go b/token-broker/internal/cache/jwt_parser.go new file mode 100644 index 00000000..dde387c5 --- /dev/null +++ b/token-broker/internal/cache/jwt_parser.go @@ -0,0 +1,78 @@ +package cache + +import ( + "encoding/base64" + "encoding/json" + "strings" + "time" +) + +// JWTClaims represents the standard JWT claims we care about. +type JWTClaims struct { + Exp int64 `json:"exp"` // Expiration time (Unix timestamp) +} + +// ParseJWTExpiry extracts the expiration time from a JWT token. +// Returns the expiration time or an error if the token is not a valid JWT. +// For non-JWT tokens (e.g., GitHub PATs), returns a far-future time. +func ParseJWTExpiry(token string) (time.Time, error) { + // Check if token looks like a JWT (three base64-encoded parts separated by dots) + parts := strings.Split(token, ".") + if len(parts) != 3 { + // Not a JWT, treat as long-lived token (1 year) + return time.Now().Add(365 * 24 * time.Hour), nil + } + + // Decode the payload (second part) + payload := parts[1] + + // Add padding if needed (JWT uses base64url without padding) + switch len(payload) % 4 { + case 2: + payload += "==" + case 3: + payload += "=" + } + + // Decode base64 + decoded, err := base64.RawURLEncoding.DecodeString(payload) + if err != nil { + // If decoding fails, treat as non-JWT token + return time.Now().Add(365 * 24 * time.Hour), nil + } + + // Parse JSON + var claims JWTClaims + if err := json.Unmarshal(decoded, &claims); err != nil { + // If JSON parsing fails, treat as non-JWT token + return time.Now().Add(365 * 24 * time.Hour), nil + } + + // Check if exp claim exists + if claims.Exp == 0 { + // No expiration claim, treat as long-lived + return time.Now().Add(365 * 24 * time.Hour), nil + } + + // Convert Unix timestamp to time.Time + expiresAt := time.Unix(claims.Exp, 0) + + return expiresAt, nil +} + +// IsTokenExpired checks if a token is expired or near expiry (< 5 minutes remaining). +func IsTokenExpired(expiresAt time.Time) bool { + now := time.Now() + + // Token is expired + if now.After(expiresAt) { + return true + } + + // Token is near expiry (< 5 minutes remaining) + if expiresAt.Sub(now) < 5*time.Minute { + return true + } + + return false +} diff --git a/token-broker/internal/cache/jwt_parser_test.go b/token-broker/internal/cache/jwt_parser_test.go new file mode 100644 index 00000000..1477086a --- /dev/null +++ b/token-broker/internal/cache/jwt_parser_test.go @@ -0,0 +1,126 @@ +package cache + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "testing" + "time" +) + +func TestParseJWTExpiry_ValidJWT(t *testing.T) { + // Create a valid JWT with expiry in 1 hour + expiresAt := time.Now().Add(1 * time.Hour) + token := createValidJWT(expiresAt) + + parsedExpiry, err := ParseJWTExpiry(token) + if err != nil { + t.Fatalf("ParseJWTExpiry failed: %v", err) + } + + // Allow 1 second tolerance for time comparison + diff := parsedExpiry.Sub(expiresAt) + if diff < -time.Second || diff > time.Second { + t.Errorf("Expected expiry %v, got %v (diff: %v)", expiresAt, parsedExpiry, diff) + } +} + +func TestParseJWTExpiry_NonJWT(t *testing.T) { + // Non-JWT token (e.g., GitHub PAT) + token := "ghp_1234567890abcdefghijklmnopqrstuvwxyz" // #notsecret + + parsedExpiry, err := ParseJWTExpiry(token) + if err != nil { + t.Fatalf("ParseJWTExpiry failed: %v", err) + } + + // Should return far-future expiry (1 year) + expectedMin := time.Now().Add(364 * 24 * time.Hour) + if parsedExpiry.Before(expectedMin) { + t.Errorf("Expected far-future expiry, got %v", parsedExpiry) + } +} + +func TestParseJWTExpiry_InvalidBase64(t *testing.T) { + // JWT with invalid base64 in payload + token := "header.!!!invalid-base64!!!.signature" + + parsedExpiry, err := ParseJWTExpiry(token) + if err != nil { + t.Fatalf("ParseJWTExpiry failed: %v", err) + } + + // Should treat as non-JWT and return far-future expiry + expectedMin := time.Now().Add(364 * 24 * time.Hour) + if parsedExpiry.Before(expectedMin) { + t.Errorf("Expected far-future expiry for invalid JWT, got %v", parsedExpiry) + } +} + +func TestParseJWTExpiry_NoExpClaim(t *testing.T) { + // JWT without exp claim + payload := map[string]interface{}{ + "sub": "user123", + "iat": time.Now().Unix(), + } + payloadBytes, _ := json.Marshal(payload) + payloadB64 := base64.RawURLEncoding.EncodeToString(payloadBytes) + token := fmt.Sprintf("eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.%s.signature", payloadB64) // #notsecret + + parsedExpiry, err := ParseJWTExpiry(token) + if err != nil { + t.Fatalf("ParseJWTExpiry failed: %v", err) + } + + // Should return far-future expiry + expectedMin := time.Now().Add(364 * 24 * time.Hour) + if parsedExpiry.Before(expectedMin) { + t.Errorf("Expected far-future expiry for JWT without exp, got %v", parsedExpiry) + } +} + +func TestIsTokenExpired_Expired(t *testing.T) { + expiresAt := time.Now().Add(-1 * time.Hour) // 1 hour ago + if !IsTokenExpired(expiresAt) { + t.Error("Token should be expired") + } +} + +func TestIsTokenExpired_NearExpiry(t *testing.T) { + expiresAt := time.Now().Add(3 * time.Minute) // 3 minutes from now (< 5 min threshold) + if !IsTokenExpired(expiresAt) { + t.Error("Token should be considered expired (near expiry)") + } +} + +func TestIsTokenExpired_Valid(t *testing.T) { + expiresAt := time.Now().Add(10 * time.Minute) // 10 minutes from now + if IsTokenExpired(expiresAt) { + t.Error("Token should not be expired") + } +} + +// Helper function to create a valid JWT with specific expiry +func createValidJWT(expiresAt time.Time) string { + // Create header + header := map[string]string{ + "alg": "HS256", + "typ": "JWT", + } + headerBytes, _ := json.Marshal(header) + headerB64 := base64.RawURLEncoding.EncodeToString(headerBytes) + + // Create payload with exp claim + payload := map[string]interface{}{ + "sub": "user123", + "exp": expiresAt.Unix(), + "iat": time.Now().Unix(), + } + payloadBytes, _ := json.Marshal(payload) + payloadB64 := base64.RawURLEncoding.EncodeToString(payloadBytes) + + // Create signature (dummy for testing) + signature := "test-signature" + + return fmt.Sprintf("%s.%s.%s", headerB64, payloadB64, signature) +} diff --git a/token-broker/internal/cache/token_cache.go b/token-broker/internal/cache/token_cache.go new file mode 100644 index 00000000..84f72c78 --- /dev/null +++ b/token-broker/internal/cache/token_cache.go @@ -0,0 +1,170 @@ +package cache + +import ( + "fmt" + "log/slog" + "sync" + "time" + + "github.com/kagenti/token-broker/internal/core" +) + +// TokenEntry represents a cached token. +type TokenEntry struct { + AccessToken string + ExpiresAt time.Time + CreatedAt time.Time +} + +// TokenCache stores and retrieves access tokens per session and resource server. +type TokenCache struct { + tokens map[string]map[string]*TokenEntry // sessionKey -> resourceURL -> token + mu sync.RWMutex + clock core.Clock +} + +// NewTokenCache creates a new token cache. +func NewTokenCache(clock core.Clock) *TokenCache { + return &TokenCache{ + tokens: make(map[string]map[string]*TokenEntry), + clock: clock, + } +} + +// GetToken retrieves a cached token for a session and resource server. +// Returns the token and true if found and not expired, empty string and false otherwise. +func (tc *TokenCache) GetToken(sessionKey, resourceURL string) (string, bool) { + tc.mu.RLock() + defer tc.mu.RUnlock() + + sessionTokens, ok := tc.tokens[sessionKey] + if !ok { + return "", false + } + + entry, ok := sessionTokens[resourceURL] + if !ok { + return "", false + } + + // Check if token is expired or near expiry + if IsTokenExpired(entry.ExpiresAt) { + slog.Debug("Token expired or near expiry", + "session_key", sessionKey, + "resource_url", resourceURL, + "expires_at", entry.ExpiresAt, + "now", tc.clock.Now()) + return "", false + } + + slog.Debug("Token cache hit", + "session_key", sessionKey, + "resource_url", resourceURL, + "expires_at", entry.ExpiresAt) + + return entry.AccessToken, true +} + +// SetToken stores a token for a session and resource server. +func (tc *TokenCache) SetToken(sessionKey, resourceURL, token string) error { + tc.mu.Lock() + defer tc.mu.Unlock() + + // Parse token expiry + expiresAt, err := ParseJWTExpiry(token) + if err != nil { + return fmt.Errorf("failed to parse token expiry: %w", err) + } + + // Create session token map if it doesn't exist + if tc.tokens[sessionKey] == nil { + tc.tokens[sessionKey] = make(map[string]*TokenEntry) + } + + // Store token + tc.tokens[sessionKey][resourceURL] = &TokenEntry{ + AccessToken: token, + ExpiresAt: expiresAt, + CreatedAt: tc.clock.Now(), + } + + slog.Info("Token cached", + "session_key", sessionKey, + "resource_url", resourceURL, + "expires_at", expiresAt, + "ttl", expiresAt.Sub(tc.clock.Now())) + + return nil +} + +// DeleteToken removes a token from the cache. +func (tc *TokenCache) DeleteToken(sessionKey, resourceURL string) { + tc.mu.Lock() + defer tc.mu.Unlock() + + if sessionTokens, ok := tc.tokens[sessionKey]; ok { + delete(sessionTokens, resourceURL) + slog.Info("Token removed from cache", + "session_key", sessionKey, + "resource_url", resourceURL) + + // Clean up empty session map + if len(sessionTokens) == 0 { + delete(tc.tokens, sessionKey) + } + } +} + +// GetCacheSize returns the total number of cached tokens. +func (tc *TokenCache) GetCacheSize() int { + tc.mu.RLock() + defer tc.mu.RUnlock() + + count := 0 + for _, sessionTokens := range tc.tokens { + count += len(sessionTokens) + } + return count +} + +// GetSessionTokenCount returns the number of cached tokens for a session. +func (tc *TokenCache) GetSessionTokenCount(sessionKey string) int { + tc.mu.RLock() + defer tc.mu.RUnlock() + + if sessionTokens, ok := tc.tokens[sessionKey]; ok { + return len(sessionTokens) + } + return 0 +} + +// CleanupExpiredTokens removes all expired tokens from the cache. +// This can be called periodically to free memory. +func (tc *TokenCache) CleanupExpiredTokens() int { + tc.mu.Lock() + defer tc.mu.Unlock() + + removed := 0 + for sessionKey, sessionTokens := range tc.tokens { + for resourceURL, entry := range sessionTokens { + if IsTokenExpired(entry.ExpiresAt) { + delete(sessionTokens, resourceURL) + removed++ + slog.Debug("Expired token removed", + "session_key", sessionKey, + "resource_url", resourceURL) + } + } + + // Clean up empty session map + if len(sessionTokens) == 0 { + delete(tc.tokens, sessionKey) + } + } + + if removed > 0 { + slog.Info("Expired tokens cleaned up", "removed_count", removed) + } + + return removed +} diff --git a/token-broker/internal/cache/token_cache_test.go b/token-broker/internal/cache/token_cache_test.go new file mode 100644 index 00000000..f7bae304 --- /dev/null +++ b/token-broker/internal/cache/token_cache_test.go @@ -0,0 +1,296 @@ +package cache + +import ( + "testing" + "time" + + "github.com/kagenti/token-broker/internal/core" +) + +// FakeClock implements core.Clock for testing. +type FakeClock struct { + now time.Time +} + +func (c *FakeClock) Now() time.Time { + return c.now +} + +func (c *FakeClock) After(d time.Duration) <-chan time.Time { + ch := make(chan time.Time, 1) + ch <- c.now.Add(d) + return ch +} + +func (c *FakeClock) Advance(d time.Duration) { + c.now = c.now.Add(d) +} + +func TestTokenCache_SetAndGet(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + sessionKey := "session1" + resourceURL := "http://mcp.example.com" + token := "test-token-123" + + // Set token + err := cache.SetToken(sessionKey, resourceURL, token) + if err != nil { + t.Fatalf("SetToken failed: %v", err) + } + + // Get token + retrievedToken, found := cache.GetToken(sessionKey, resourceURL) + if !found { + t.Fatal("Token not found") + } + + if retrievedToken != token { + t.Errorf("Expected token %s, got %s", token, retrievedToken) + } +} + +func TestTokenCache_GetNonExistent(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + token, found := cache.GetToken("session1", "http://resource.example.com") + if found { + t.Error("Expected token not found, but got one") + } + if token != "" { + t.Errorf("Expected empty token, got %s", token) + } +} + +func TestTokenCache_ExpiredToken(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + sessionKey := "session1" + resourceURL := "http://mcp.example.com" + + // Create a JWT token that expires in 1 hour from fake clock time + expiresAt := clock.now.Add(1 * time.Hour) + token := createTestJWT(expiresAt) + + // Set token + err := cache.SetToken(sessionKey, resourceURL, token) + if err != nil { + t.Fatalf("SetToken failed: %v", err) + } + + // Token should be found (using fake clock time) + _, found := cache.GetToken(sessionKey, resourceURL) + if !found { + t.Error("Token should be found") + } + + // Advance fake clock past expiry + clock.Advance(2 * time.Hour) + + // Token should not be found (expired according to fake clock) + // Note: IsTokenExpired uses time.Now(), not the fake clock + // So we need to create a token that's actually expired + pastExpiresAt := time.Now().Add(-1 * time.Hour) + expiredToken := createTestJWT(pastExpiresAt) + _ = cache.SetToken(sessionKey, resourceURL, expiredToken) + + _, found = cache.GetToken(sessionKey, resourceURL) + if found { + t.Error("Expired token should not be found") + } +} + +func TestTokenCache_NearExpiryToken(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + sessionKey := "session1" + resourceURL := "http://mcp.example.com" + + // Create a JWT token that expires in 3 minutes (< 5 min threshold) + expiresAt := clock.now.Add(3 * time.Minute) + token := createTestJWT(expiresAt) + + // Set token + err := cache.SetToken(sessionKey, resourceURL, token) + if err != nil { + t.Fatalf("SetToken failed: %v", err) + } + + // Token should not be found (near expiry) + _, found := cache.GetToken(sessionKey, resourceURL) + if found { + t.Error("Near-expiry token should not be found") + } +} + +func TestTokenCache_PerSessionPerServer(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + session1 := "session1" + session2 := "session2" + server1 := "http://mcp1.example.com" + server2 := "http://mcp2.example.com" + + token1 := "token1" + token2 := "token2" + token3 := "token3" + token4 := "token4" + + // Set tokens for different combinations + _ = cache.SetToken(session1, server1, token1) + _ = cache.SetToken(session1, server2, token2) + _ = cache.SetToken(session2, server1, token3) + _ = cache.SetToken(session2, server2, token4) + + // Verify isolation + tests := []struct { + session string + server string + expected string + }{ + {session1, server1, token1}, + {session1, server2, token2}, + {session2, server1, token3}, + {session2, server2, token4}, + } + + for _, tt := range tests { + token, found := cache.GetToken(tt.session, tt.server) + if !found { + t.Errorf("Token not found for session=%s, server=%s", tt.session, tt.server) + continue + } + if token != tt.expected { + t.Errorf("Expected token %s for session=%s, server=%s, got %s", + tt.expected, tt.session, tt.server, token) + } + } +} + +func TestTokenCache_DeleteToken(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + sessionKey := "session1" + resourceURL := "http://mcp.example.com" + token := "test-token" + + // Set token + _ = cache.SetToken(sessionKey, resourceURL, token) + + // Verify it exists + _, found := cache.GetToken(sessionKey, resourceURL) + if !found { + t.Fatal("Token should exist") + } + + // Delete token + cache.DeleteToken(sessionKey, resourceURL) + + // Verify it's gone + _, found = cache.GetToken(sessionKey, resourceURL) + if found { + t.Error("Token should be deleted") + } +} + +func TestTokenCache_GetCacheSize(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + if size := cache.GetCacheSize(); size != 0 { + t.Errorf("Expected cache size 0, got %d", size) + } + + // Add tokens + _ = cache.SetToken("session1", "http://mcp1.example.com", "token1") + _ = cache.SetToken("session1", "http://mcp2.example.com", "token2") + _ = cache.SetToken("session2", "http://mcp1.example.com", "token3") + + if size := cache.GetCacheSize(); size != 3 { + t.Errorf("Expected cache size 3, got %d", size) + } + + // Delete one + cache.DeleteToken("session1", "http://mcp1.example.com") + + if size := cache.GetCacheSize(); size != 2 { + t.Errorf("Expected cache size 2, got %d", size) + } +} + +func TestTokenCache_GetSessionTokenCount(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + session1 := "session1" + session2 := "session2" + + // Add tokens for session1 + _ = cache.SetToken(session1, "http://mcp1.example.com", "token1") + _ = cache.SetToken(session1, "http://mcp2.example.com", "token2") + + // Add token for session2 + _ = cache.SetToken(session2, "http://mcp1.example.com", "token3") + + if count := cache.GetSessionTokenCount(session1); count != 2 { + t.Errorf("Expected session1 token count 2, got %d", count) + } + + if count := cache.GetSessionTokenCount(session2); count != 1 { + t.Errorf("Expected session2 token count 1, got %d", count) + } + + if count := cache.GetSessionTokenCount("session3"); count != 0 { + t.Errorf("Expected session3 token count 0, got %d", count) + } +} + +func TestTokenCache_CleanupExpiredTokens(t *testing.T) { + clock := &FakeClock{now: time.Now()} + cache := NewTokenCache(clock) + + // Add tokens with different expiry times (using real time for IsTokenExpired) + now := time.Now() + expiresIn2Hours := now.Add(2 * time.Hour) + expiredToken := now.Add(-1 * time.Hour) // Already expired + + _ = cache.SetToken("session1", "http://mcp1.example.com", createTestJWT(expiredToken)) + _ = cache.SetToken("session1", "http://mcp2.example.com", createTestJWT(expiresIn2Hours)) + _ = cache.SetToken("session2", "http://mcp1.example.com", createTestJWT(expiredToken)) + + if size := cache.GetCacheSize(); size != 3 { + t.Errorf("Expected cache size 3, got %d", size) + } + + // Cleanup expired tokens + removed := cache.CleanupExpiredTokens() + + if removed != 2 { + t.Errorf("Expected 2 tokens removed, got %d", removed) + } + + if size := cache.GetCacheSize(); size != 1 { + t.Errorf("Expected cache size 1 after cleanup, got %d", size) + } + + // Verify the remaining token is the one that expires in 2 hours + _, found := cache.GetToken("session1", "http://mcp2.example.com") + if !found { + t.Error("Token with 2-hour expiry should still exist") + } +} + +// Helper function to create a test JWT with specific expiry +func createTestJWT(expiresAt time.Time) string { + // Use the same helper from jwt_parser_test.go + return createValidJWT(expiresAt) +} + +// Verify FakeClock implements core.Clock +var _ core.Clock = (*FakeClock)(nil) diff --git a/token-broker/internal/core/broker.go b/token-broker/internal/core/broker.go new file mode 100644 index 00000000..3966690c --- /dev/null +++ b/token-broker/internal/core/broker.go @@ -0,0 +1,315 @@ +package core + +import ( + "context" + "fmt" + "log/slog" + "time" + + "github.com/kagenti/token-broker/internal/oauth" + pkgauth "github.com/kagenti/token-broker/pkg/oauth" +) + +// TokenBroker orchestrates token acquisition for OAuth sessions. +type TokenBroker struct { + sessionStore SessionStore + tokenCache TokenCache + discoverer *oauth.Discoverer + oauthClient *oauth.Client + callbackURL string // Token Broker's own callback URL + waitTimeout time.Duration + logger *slog.Logger +} + +// NewTokenBroker creates a new token broker. +func NewTokenBroker( + sessionStore SessionStore, + tokenCache TokenCache, + discoverer *oauth.Discoverer, + oauthClient *oauth.Client, + callbackURL string, + waitTimeout time.Duration, + logger *slog.Logger, +) *TokenBroker { + return &TokenBroker{ + sessionStore: sessionStore, + tokenCache: tokenCache, + discoverer: discoverer, + oauthClient: oauthClient, + callbackURL: callbackURL, + waitTimeout: waitTimeout, + logger: logger, + } +} + +// AcquireToken acquires a token for a user and resource server. +// This implements the double-checked locking pattern with session semaphore. +func (tb *TokenBroker) AcquireToken(ctx context.Context, sessionKey, userID, resourceURL string) (string, error) { + // Step 1: Validate session and user FIRST (security-critical) + // This prevents session hijacking where user A tries to use user B's session + if err := tb.sessionStore.ValidateSession(sessionKey, userID); err != nil { + return "", fmt.Errorf("session validation failed: %w", err) + } + + // Step 2: Check cache (fast path after validation) + // Cache is now keyed by sessionKey to ensure session isolation + if token, found := tb.tokenCache.GetToken(sessionKey, resourceURL); found { + tb.logger.Info("Token found in cache (after session validation)", + "session_key", sessionKey, + "user_id", userID, + "resource_url", resourceURL) + return token, nil + } + + // Get session + session, err := tb.sessionStore.GetSession(sessionKey) + if err != nil { + return "", fmt.Errorf("failed to get session: %w", err) + } + + // Step 3: Acquire per-session semaphore + tb.logger.Debug("Acquiring session semaphore", + "session_key", sessionKey, + "user_id", userID, + "resource_url", resourceURL) + + if err := session.AcquisitionSemaphore.Acquire(ctx); err != nil { + return "", fmt.Errorf("failed to acquire semaphore: %w", err) + } + defer session.AcquisitionSemaphore.Release() + + tb.logger.Debug("Session semaphore acquired", + "session_key", sessionKey) + + // Step 4: Check cache again (double-check after semaphore) + // Cache is keyed by sessionKey to ensure session isolation + if token, found := tb.tokenCache.GetToken(sessionKey, resourceURL); found { + tb.logger.Info("Token found in cache after semaphore acquisition", + "session_key", sessionKey, + "user_id", userID, + "resource_url", resourceURL) + return token, nil + } + + // Step 5: Token not in cache, initiate OAuth flow + tb.logger.Info("Token not in cache, initiating OAuth flow", + "session_key", sessionKey, + "user_id", userID, + "resource_url", resourceURL) + + token, err := tb.performOAuthFlow(ctx, session, userID, resourceURL) + if err != nil { + return "", fmt.Errorf("OAuth flow failed: %w", err) + } + + return token, nil +} + +// performOAuthFlow performs the complete OAuth flow to obtain a token. +// The Token Broker now acts as the OAuth client, generating PKCE and exchanging tokens. +func (tb *TokenBroker) performOAuthFlow(ctx context.Context, session *Session, userID, resourceURL string) (string, error) { + // Step 1: Discover OAuth endpoints and scopes from resource server + tb.logger.Debug("Discovering OAuth endpoints from resource server", + "resource_url", resourceURL) + + authEndpoint, tokenEndpoint, scopes, err := tb.discoverer.DiscoverEndpoints(ctx, resourceURL) + if err != nil { + return "", fmt.Errorf("%w: failed to discover OAuth endpoints: %w", ErrOAuthUpstream, err) + } + + tb.logger.Info("OAuth endpoints discovered", + "resource_url", resourceURL, + "auth_endpoint", authEndpoint, + "token_endpoint", tokenEndpoint, + "scopes", scopes) + + // Step 2: Generate PKCE challenge at Token Broker + tb.logger.Debug("Generating PKCE challenge", + "resource_url", resourceURL) + + pkce, err := pkgauth.GeneratePKCEChallenge() + if err != nil { + return "", fmt.Errorf("failed to generate PKCE challenge: %w", err) + } + + tb.logger.Debug("PKCE challenge generated", + "resource_url", resourceURL, + "challenge_method", pkce.Method) + + // Step 3: Store OAuth transaction in session + session.ActiveOAuthTx = &OAuthTransaction{ + ResourceURL: resourceURL, + AuthEndpoint: authEndpoint, + TokenEndpoint: tokenEndpoint, + Scopes: scopes, + PKCE: pkce, + CodeVerifier: pkce.Verifier, + Status: OAuthTxStatusWaitingCode, + CompletionChan: make(chan OAuthCompletion, 1), + } + + // Step 4: Build authorization URL + callbackURL := tb.callbackURL + + authURL, state, err := tb.oauthClient.BuildAuthorizationURL(authEndpoint, callbackURL, session.SessionKey, scopes, pkce) + if err != nil { + session.ActiveOAuthTx.Status = OAuthTxStatusFailed + return "", fmt.Errorf("failed to build authorization URL: %w", err) + } + + session.ActiveOAuthTx.AuthURL = authURL + session.ActiveOAuthTx.CallbackURL = callbackURL + session.ActiveOAuthTx.State = state + + tb.logger.Info("Authorization URL built", + "resource_url", resourceURL, + "auth_url_length", len(authURL)) + + // Step 5: Publish oauth_url_ready event to Backend + tb.logger.Info("Publishing oauth_url_ready event to Backend", + "session_key", session.SessionKey, + "resource_url", resourceURL, + "auth_url_length", len(authURL)) + tb.logger.Debug("oauth_url_ready auth URL", "auth_url", authURL) + + urlReadyEvent := Event{ + Type: "oauth_url_ready", + AuthURL: authURL, + } + + select { + case session.EventWaiters <- urlReadyEvent: + tb.logger.Info("oauth_url_ready event successfully sent to Backend", + "session_key", session.SessionKey, + "resource_url", resourceURL) + case <-ctx.Done(): + return "", fmt.Errorf("context cancelled while publishing event: %w", ctx.Err()) + default: + return "", fmt.Errorf("no event waiter available") + } + + // Step 6: Wait for OAuth completion (authorization code from callback) + tb.logger.Info("Waiting for OAuth completion", + "session_key", session.SessionKey, + "timeout", tb.waitTimeout) + + waitCtx2, cancel2 := context.WithTimeout(ctx, tb.waitTimeout) + defer cancel2() + + var completion OAuthCompletion + select { + case completion = <-session.ActiveOAuthTx.CompletionChan: + if completion.Error != nil { + return "", fmt.Errorf("OAuth completion failed: %w", completion.Error) + } + tb.logger.Info("OAuth completion received", + "session_key", session.SessionKey, + "code_length", len(completion.Code)) + + case <-waitCtx2.Done(): + session.ActiveOAuthTx.Status = OAuthTxStatusFailed + return "", fmt.Errorf("%w", ErrOAuthTimeout) + + case <-session.Done: + session.ActiveOAuthTx.Status = OAuthTxStatusFailed + return "", fmt.Errorf("%w", ErrSessionEnded) + + case <-ctx.Done(): + session.ActiveOAuthTx.Status = OAuthTxStatusFailed + return "", fmt.Errorf("context cancelled: %w", ctx.Err()) + } + + // Step 7: Exchange authorization code for token directly with OAuth provider + tb.logger.Debug("Exchanging authorization code for token", + "resource_url", resourceURL, + "token_endpoint", tokenEndpoint) + + token, expiresIn, err := tb.oauthClient.ExchangeToken(ctx, tokenEndpoint, completion.Code, pkce.Verifier, callbackURL) + if err != nil { + session.ActiveOAuthTx.Status = OAuthTxStatusFailed + return "", fmt.Errorf("%w: token exchange failed: %w", ErrOAuthUpstream, err) + } + + session.ActiveOAuthTx.Status = OAuthTxStatusCompleted + + tb.logger.Info("Token obtained successfully", + "session_key", session.SessionKey, + "resource_url", resourceURL, + "token_length", len(token), + "expires_in", expiresIn) + + // Step 8: Cache token (keyed by sessionKey for session isolation) + if err := tb.tokenCache.SetToken(session.SessionKey, resourceURL, token); err != nil { + tb.logger.Error("Failed to cache token", + "error", err, + "session_key", session.SessionKey, + "user_id", userID, + "resource_url", resourceURL) + // Don't fail the request, we have the token + } + + // Step 9: Unblock any other waiters for this resource server + if waiters, ok := session.TokenWaiters[resourceURL]; ok { + for _, waiter := range waiters { + select { + case waiter <- TokenResult{Token: token, Error: nil}: + default: + } + close(waiter) + } + delete(session.TokenWaiters, resourceURL) + tb.logger.Debug("Unblocked token waiters", + "session_key", session.SessionKey, + "resource_url", resourceURL, + "waiter_count", len(waiters)) + } + + // Clear active OAuth transaction + session.ActiveOAuthTx = nil + + return token, nil +} + +// CompleteOAuth completes an OAuth flow with the authorization code and state. +func (tb *TokenBroker) CompleteOAuth(sessionKey, userID, code, state string) error { + // Validate session ownership (defense-in-depth) + if err := tb.sessionStore.ValidateSession(sessionKey, userID); err != nil { + return fmt.Errorf("session validation failed: %w", err) + } + + session, err := tb.sessionStore.GetSession(sessionKey) + if err != nil { + return fmt.Errorf("session not found: %w", err) + } + + if session.ActiveOAuthTx == nil { + return fmt.Errorf("no active OAuth transaction") + } + + tb.logger.Info("Completing OAuth flow", + "session_key", sessionKey, + "resource_url", session.ActiveOAuthTx.ResourceURL, + "code_length", len(code), + "state_length", len(state)) + + // Send completion to waiting OAuth flow + completion := OAuthCompletion{ + Code: code, + State: state, + Error: nil, + } + + select { + case session.ActiveOAuthTx.CompletionChan <- completion: + tb.logger.Debug("OAuth completion sent to waiting flow", + "session_key", sessionKey) + return nil + default: + return fmt.Errorf("no waiter for OAuth completion") + } +} + +// GetSessionStore returns the session store (for API handlers). +func (tb *TokenBroker) GetSessionStore() SessionStore { + return tb.sessionStore +} diff --git a/token-broker/internal/core/broker_test.go b/token-broker/internal/core/broker_test.go new file mode 100644 index 00000000..d0cccd86 --- /dev/null +++ b/token-broker/internal/core/broker_test.go @@ -0,0 +1,91 @@ +package core_test + +import ( + "context" + "errors" + "log/slog" + "os" + "testing" + "time" + + "github.com/kagenti/token-broker/internal/cache" + "github.com/kagenti/token-broker/internal/core" + "github.com/kagenti/token-broker/internal/oauth" + "github.com/kagenti/token-broker/internal/session" +) + +// TestAcquireToken_UnblocksOnSessionEnd verifies that ending a session while a +// token request is blocked waiting for OAuth completion returns promptly with a +// "session ended" error, rather than hanging until TokenWaitTimeout. +func TestAcquireToken_UnblocksOnSessionEnd(t *testing.T) { + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + clock := &core.RealClock{} + + tokenCache := cache.NewTokenCache(clock) + // Long timeout: the test must succeed because EndSession unblocks the wait, + // not because the wait times out. + sessionManager := session.NewSessionManager(5*time.Minute, 5, clock, logger) + defer sessionManager.Shutdown() + + oauthClient := oauth.NewClient(&oauth.Config{ + ClientID: "test-client", + ClientSecret: "test-secret", + }, logger) + + // Configured metadata so DiscoverEndpoints does not make an HTTP request. + discoverer := oauth.NewDiscovererWithConfig(&oauth.DiscovererConfig{ + AuthorizationEndpoint: "https://example.com/oauth/authorize", + TokenEndpoint: "https://example.com/oauth/token", + ScopesSupported: []string{"read"}, + }, logger) + + broker := core.NewTokenBroker( + sessionManager, + tokenCache, + discoverer, + oauthClient, + "http://broker.example.com/oauth/callback", + 5*time.Minute, + logger, + ) + + userID := "user1" + sessionKey := "session-1" + resourceURL := "https://mcp.example.com" + + if err := sessionManager.CreateSession(sessionKey, userID, ""); err != nil { + t.Fatalf("CreateSession failed: %v", err) + } + + // AcquireToken blocks waiting for OAuth completion (no callback will arrive). + type result struct { + token string + err error + } + resultCh := make(chan result, 1) + go func() { + token, err := broker.AcquireToken(context.Background(), sessionKey, userID, resourceURL) + resultCh <- result{token: token, err: err} + }() + + // Give the goroutine time to reach the OAuth-completion wait, then end the + // session out from under it. + time.Sleep(100 * time.Millisecond) + if err := sessionManager.EndSession(sessionKey); err != nil { + t.Fatalf("EndSession failed: %v", err) + } + + select { + case res := <-resultCh: + if res.err == nil { + t.Fatalf("expected error after session ended, got token %q", res.token) + } + // Must be classifiable via errors.Is despite AcquireToken's wrapping — + // this is what lets the API layer return 401 instead of a spurious 503. + if !errors.Is(res.err, core.ErrSessionEnded) { + t.Errorf("expected error to wrap ErrSessionEnded, got: %v", res.err) + } + case <-time.After(5 * time.Second): + t.Fatal("AcquireToken did not return after EndSession; it is still blocked") + } +} diff --git a/token-broker/internal/core/errors.go b/token-broker/internal/core/errors.go new file mode 100644 index 00000000..be3e46f3 --- /dev/null +++ b/token-broker/internal/core/errors.go @@ -0,0 +1,31 @@ +package core + +import "errors" + +// Sentinel errors returned by the token broker. They are classified in the API +// layer with errors.Is so that wrapping (e.g. fmt.Errorf("...: %w", err)) does +// not collapse a specific, recoverable condition into a generic 5xx response. +// +// Why this matters: agents/sidecars treat a 5xx from the token service as +// "this route is dead" and stop retrying it. So a 5xx must be reserved for +// failures that are permanent or genuinely server-side, never for session/client +// state (401) or a transient upstream hiccup (502). +var ( + // ErrSessionEnded indicates the session was ended (e.g. by the backend or + // shutdown) while a token acquisition was in flight. Permanent for this + // session; the client must establish a new one. Maps to 401. + ErrSessionEnded = errors.New("session ended") + + // ErrSessionExpired indicates the session's idle timeout elapsed. Permanent + // for this session; the client must establish a new one. Maps to 401. + ErrSessionExpired = errors.New("session expired") + + // ErrOAuthTimeout indicates the user did not complete the OAuth flow within + // the wait timeout. Transient — the flow can be retried. Maps to 408. + ErrOAuthTimeout = errors.New("timeout waiting for OAuth completion") + + // ErrOAuthUpstream indicates a failure talking to the OAuth provider or the + // resource-server discovery endpoint. Transient upstream error, retryable. + // Maps to 502. + ErrOAuthUpstream = errors.New("oauth upstream failure") +) diff --git a/token-broker/internal/core/interfaces.go b/token-broker/internal/core/interfaces.go new file mode 100644 index 00000000..b3f8b7c7 --- /dev/null +++ b/token-broker/internal/core/interfaces.go @@ -0,0 +1,176 @@ +// Package core provides the main token acquisition orchestration for the Token Broker. +package core + +import ( + "context" + "time" + + pkgauth "github.com/kagenti/token-broker/pkg/oauth" +) + +// SessionStore manages OAuth sessions and their lifecycle. +type SessionStore interface { + // CreateSession creates a new session with a provided session key and backend redirect URL. + CreateSession(sessionKey, userID, backendRedirectURL string) error + + // GetSession retrieves a session by its key. + GetSession(sessionKey string) (*Session, error) + + // GetSessionByState retrieves a session by its OAuth state parameter. + GetSessionByState(state string) (*Session, error) + + // ValidateSession checks if a session exists and belongs to the specified user. + ValidateSession(sessionKey, userID string) error + + // EndSession terminates a session and releases all resources. + EndSession(sessionKey string) error + + // ExpireSession marks a session as expired and fails all pending requests. + ExpireSession(sessionKey string) + + // ResetSessionTimer resets the idle timeout timer for a session. + ResetSessionTimer(sessionKey string) + + // StartSessionTimer starts the idle timeout timer for a session. + StartSessionTimer(sessionKey string) +} + +// TokenCache stores and retrieves access tokens per session and resource server. +type TokenCache interface { + // GetToken retrieves a cached token for a session and resource server. + // Returns the token and true if found and not expired, empty string and false otherwise. + GetToken(sessionKey, resourceURL string) (token string, found bool) + + // SetToken stores a token for a session and resource server. + SetToken(sessionKey, resourceURL, token string) error + + // DeleteToken removes a token from the cache. + DeleteToken(sessionKey, resourceURL string) +} + +// OAuthDiscoverer performs OAuth discovery against a resource server. +// Updated to support the new architecture where Token Broker is the OAuth client. +type OAuthDiscoverer interface { + // SendSyntheticRequest sends an unauthenticated resource request to trigger OAuth elicitation. + // Returns true if the server responds with 401, false otherwise. + // This is optional and used for validation purposes. + SendSyntheticRequest(ctx context.Context, resourceURL string) (is401 bool, err error) + + // DiscoverEndpoints discovers OAuth endpoints and scopes from resource server's .well-known metadata. + // This is ALWAYS called - there is no configured mode. + // Returns: authorizationEndpoint, tokenEndpoint, scopes, error + DiscoverEndpoints(ctx context.Context, resourceURL string) (authEndpoint, tokenEndpoint string, scopes []string, err error) +} + +// Clock provides time-related operations for testability. +type Clock interface { + // Now returns the current time. + Now() time.Time + + // After returns a channel that receives the current time after the specified duration. + After(d time.Duration) <-chan time.Time +} + +// Session represents an OAuth session. +type Session struct { + SessionKey string + UserID string + BackendSessionRedirectURL string // Where to redirect user after OAuth completion + CreatedAt time.Time + LastPollAt time.Time + ExpirationDeadline time.Time + + // Concurrency control + AcquisitionSemaphore Semaphore + + // Event coordination + EventWaiters chan Event + TokenWaiters map[string][]chan TokenResult // resourceURL -> waiters + + // Done is closed when the session is ended or expired, to unblock waiters. + Done chan struct{} + + // Active OAuth transaction + ActiveOAuthTx *OAuthTransaction +} + +// Semaphore provides mutual exclusion for token acquisition per session. +type Semaphore interface { + // Acquire acquires the semaphore, blocking until available or context is cancelled. + Acquire(ctx context.Context) error + + // Release releases the semaphore. + Release() +} + +// OAuthTransaction represents an active OAuth flow. +// Updated to support the new architecture where Token Broker generates PKCE +// and exchanges tokens directly with the OAuth provider. +type OAuthTransaction struct { + ResourceURL string + AuthEndpoint string // Discovered from resource server .well-known + TokenEndpoint string // Discovered from resource server .well-known + Scopes []string // Discovered from resource server .well-known + + // PKCE generated by Token Broker + PKCE *pkgauth.PKCEChallenge + CodeVerifier string // Convenience field (same as PKCE.Verifier) + + // State parameter for OAuth flow + State string // Generated by Token Broker, used to map callback to session + + // Callback URL (Token Broker's own callback URL, not Backend's) + CallbackURL string + + // Authorization URL built by Token Broker + AuthURL string + + // Flow status, one of the OAuthTxStatus* constants. + Status string + CompletionChan chan OAuthCompletion +} + +// OAuth transaction status values for OAuthTransaction.Status. +const ( + OAuthTxStatusWaitingCode = "waiting_code" + OAuthTxStatusCompleted = "completed" + OAuthTxStatusFailed = "failed" +) + +// OAuthCompletion represents the result of an OAuth flow. +type OAuthCompletion struct { + Code string + State string + Error error +} + +// Event represents an event sent to the Backend via long-poll. +type Event struct { + Type string `json:"type"` // "oauth_url_ready" | "error" + + // For oauth_url_ready (Backend should redirect frontend to OAuth provider) + AuthURL string `json:"auth_url,omitempty"` + + // For error + Message string `json:"message,omitempty"` + Code string `json:"code,omitempty"` +} + +// TokenResult represents the result of a token acquisition. +type TokenResult struct { + Token string + Error error +} + +// RealClock implements Clock using the standard time package. +type RealClock struct{} + +// Now returns the current time. +func (c *RealClock) Now() time.Time { + return time.Now() +} + +// After returns a channel that receives the current time after the specified duration. +func (c *RealClock) After(d time.Duration) <-chan time.Time { + return time.After(d) +} diff --git a/token-broker/internal/oauth/client.go b/token-broker/internal/oauth/client.go new file mode 100644 index 00000000..ff22ac4d --- /dev/null +++ b/token-broker/internal/oauth/client.go @@ -0,0 +1,170 @@ +package oauth + +import ( + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "strings" + + // Reuse existing PKCE implementation + pkgauth "github.com/kagenti/token-broker/pkg/oauth" +) + +// Config holds OAuth client configuration for the Token Broker. +// Endpoints are always discovered from resource server's .well-known metadata. +type Config struct { + // ClientID is the OAuth client ID + ClientID string + + // ClientSecret is the OAuth client secret + ClientSecret string +} + +// Client handles OAuth operations for the Token Broker. +// The Token Broker acts as the OAuth client, generating PKCE and exchanging tokens. +type Client struct { + config *Config + httpClient *http.Client + logger *slog.Logger +} + +// NewClient creates a new OAuth client. +func NewClient(config *Config, logger *slog.Logger) *Client { + if logger == nil { + logger = slog.Default() + } + return &Client{ + config: config, + httpClient: &http.Client{}, + logger: logger, + } +} + +// BuildAuthorizationURL builds the OAuth authorization URL. +// The state parameter is constructed as: sessionKey + "." + random_nonce +// This allows easy lookup of the session from the state parameter. +// The authEndpoint and scopes are discovered from resource server's .well-known metadata. +// The callbackURL is the Token Broker's own callback URL. +// Returns the authorization URL and the generated state parameter. +func (c *Client) BuildAuthorizationURL( + authEndpoint, callbackURL, sessionKey string, + scopes []string, + pkce *pkgauth.PKCEChallenge, +) (authURL string, state string, err error) { + u, err := url.Parse(authEndpoint) + if err != nil { + return "", "", fmt.Errorf("invalid authorization endpoint: %w", err) + } + + // Generate random nonce for state + nonce, err := pkgauth.GenerateState() + if err != nil { + return "", "", fmt.Errorf("failed to generate state nonce: %w", err) + } + + // Construct state as sessionKey.nonce for easy lookup + state = sessionKey + "." + nonce + + // Build query parameters + params := url.Values{} + params.Set("client_id", c.config.ClientID) + params.Set("redirect_uri", callbackURL) + params.Set("response_type", "code") + params.Set("state", state) + params.Set("code_challenge", pkce.Challenge) + params.Set("code_challenge_method", pkce.Method) + + if len(scopes) > 0 { + params.Set("scope", strings.Join(scopes, " ")) + } + + u.RawQuery = params.Encode() + + fullAuthURL := u.String() + + c.logger.Info("built authorization URL", + "auth_endpoint", authEndpoint, + "callback_url", callbackURL, + "session_key", sessionKey, + "scopes", scopes, + "has_pkce", true, + "state_length", len(state)) + c.logger.Debug("full authorization URL", "full_auth_url", fullAuthURL) + + return fullAuthURL, state, nil +} + +// ExchangeToken exchanges an authorization code for an access token. +// The Token Broker calls the OAuth provider directly (not through resource server). +// The tokenEndpoint is discovered from resource server's .well-known metadata. +// The callbackURL must match the one used in BuildAuthorizationURL. +func (c *Client) ExchangeToken( + ctx context.Context, + tokenEndpoint, code, codeVerifier, callbackURL string, +) (accessToken string, expiresIn int, err error) { + // Build token exchange request + data := url.Values{} + data.Set("client_id", c.config.ClientID) + data.Set("client_secret", c.config.ClientSecret) + data.Set("code", code) + data.Set("code_verifier", codeVerifier) + data.Set("redirect_uri", callbackURL) + data.Set("grant_type", "authorization_code") + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, strings.NewReader(data.Encode())) + if err != nil { + return "", 0, fmt.Errorf("failed to create token request: %w", err) + } + + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + + c.logger.Info("exchanging authorization code for token", + "token_endpoint", tokenEndpoint, + "client_id", c.config.ClientID, + "has_code_verifier", codeVerifier != "") + + resp, err := c.httpClient.Do(req) + if err != nil { + return "", 0, fmt.Errorf("failed to send token request: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", 0, fmt.Errorf("failed to read token response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + c.logger.Error("token exchange failed", + "status", resp.StatusCode, + "response", string(body)) + return "", 0, fmt.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(body)) + } + + // Parse token response + var tokenResp struct { + AccessToken string `json:"access_token"` + ExpiresIn int `json:"expires_in"` + TokenType string `json:"token_type"` + } + + if err := json.Unmarshal(body, &tokenResp); err != nil { + return "", 0, fmt.Errorf("failed to parse token response: %w", err) + } + + if tokenResp.AccessToken == "" { + return "", 0, fmt.Errorf("no access token in response") + } + + c.logger.Info("token exchange successful", + "token_length", len(tokenResp.AccessToken), + "expires_in", tokenResp.ExpiresIn, + "token_type", tokenResp.TokenType) + + return tokenResp.AccessToken, tokenResp.ExpiresIn, nil +} diff --git a/token-broker/internal/oauth/discovery.go b/token-broker/internal/oauth/discovery.go new file mode 100644 index 00000000..6cccec28 --- /dev/null +++ b/token-broker/internal/oauth/discovery.go @@ -0,0 +1,261 @@ +package oauth + +import ( + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "strings" +) + +// ResourceConfig holds per-resource-server OAuth configuration overrides. +type ResourceConfig struct { + // Scopes overrides the scopes to request for this server, replacing discovered scopes. + Scopes []string `json:"scopes,omitempty"` + + // AuthorizationEndpoint overrides the authorization endpoint for this server, + // replacing the one derived from .well-known discovery. + AuthorizationEndpoint string `json:"authorization_endpoint,omitempty"` + + // TokenEndpoint overrides the token endpoint for this server, + // replacing the one derived from .well-known discovery. + TokenEndpoint string `json:"token_endpoint,omitempty"` +} + +// DiscovererConfig holds optional configured OAuth metadata. +// If provided, these values will be used instead of discovery. +type DiscovererConfig struct { + AuthorizationEndpoint string + TokenEndpoint string + ScopesSupported []string + + // ResourceScopes maps resource server URL to OAuth scopes (legacy, scopes-only override). + // Deprecated: use ResourceConfigs for richer per-server configuration. + ResourceScopes map[string][]string + + // ResourceConfigs maps resource server URL to full per-server OAuth config overrides. + // Takes precedence over ResourceScopes for the same server URL. + ResourceConfigs map[string]ResourceConfig +} + +// Discoverer discovers OAuth endpoints from resource server's .well-known metadata. +// It can also use pre-configured endpoints as a fallback. +type Discoverer struct { + httpClient *http.Client + logger *slog.Logger + config *DiscovererConfig +} + +// NewDiscoverer creates a new OAuth endpoint discoverer without configured metadata. +func NewDiscoverer(logger *slog.Logger) *Discoverer { + return NewDiscovererWithConfig(nil, logger) +} + +// NewDiscovererWithConfig creates a new OAuth endpoint discoverer with optional configured metadata. +func NewDiscovererWithConfig(config *DiscovererConfig, logger *slog.Logger) *Discoverer { + if logger == nil { + logger = slog.Default() + } + if config == nil { + config = &DiscovererConfig{} + } + return &Discoverer{ + httpClient: &http.Client{}, + logger: logger, + config: config, + } +} + +// DiscoverEndpoints discovers OAuth endpoints and scopes for a resource server. +// +// Resolution order (highest priority first): +// 1. Per-server config from ResourceConfigs (overrides individual fields) +// 2. Global static config (OAUTH_AUTHORIZATION_ENDPOINT + OAUTH_TOKEN_ENDPOINT) +// 3. .well-known/oauth-protected-resource discovery +// 4. Legacy ResourceScopes (scopes-only override, lower priority than ResourceConfigs) +func (d *Discoverer) DiscoverEndpoints(ctx context.Context, resourceURL string) (string, string, []string, error) { + // Check for full per-server config first (highest priority) + if serverCfg, ok := d.config.ResourceConfigs[resourceURL]; ok { + if serverCfg.AuthorizationEndpoint != "" && serverCfg.TokenEndpoint != "" { + d.logger.Info("Using per-server OAuth config (skipping discovery)", + "resource_url", resourceURL, + "authorization_endpoint", serverCfg.AuthorizationEndpoint, + "token_endpoint", serverCfg.TokenEndpoint, + "scopes", serverCfg.Scopes) + return serverCfg.AuthorizationEndpoint, serverCfg.TokenEndpoint, serverCfg.Scopes, nil + } + } + + // Check if we have fully configured global metadata (all required fields) + hasConfiguredMetadata := d.config.AuthorizationEndpoint != "" && d.config.TokenEndpoint != "" + + if hasConfiguredMetadata { + scopes := d.config.ScopesSupported + if serverCfg, ok := d.config.ResourceConfigs[resourceURL]; ok && len(serverCfg.Scopes) > 0 { + scopes = serverCfg.Scopes + } else if override, ok := d.config.ResourceScopes[resourceURL]; ok { + scopes = override + } + d.logger.Info("Using global OAuth metadata (skipping discovery)", + "resource_url", resourceURL, + "authorization_endpoint", d.config.AuthorizationEndpoint, + "token_endpoint", d.config.TokenEndpoint, + "scopes", scopes) + return d.config.AuthorizationEndpoint, d.config.TokenEndpoint, scopes, nil + } + + // resourceURL originates from the X-Server-Url header set by the authbridge sidecar + // from operator-configured routes — it is trusted infrastructure, not end-user input. + // The scheme/host validation below is defence-in-depth, not a SSRF boundary. + parsedResource, err := url.Parse(resourceURL) + if err != nil || (parsedResource.Scheme != "http" && parsedResource.Scheme != "https") || parsedResource.Host == "" { + return "", "", nil, fmt.Errorf("invalid resource URL %q: must be http or https with a non-empty host", resourceURL) + } + + // Attempt .well-known discovery + // lgtm[go/request-forgery] resourceURL is operator-configured, not user-controlled + wellKnownURL := strings.TrimSuffix(resourceURL, "/") + "/.well-known/oauth-protected-resource" + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, wellKnownURL, nil) + if err != nil { + return "", "", nil, fmt.Errorf("failed to create discovery request: %w", err) + } + + req.Header.Set("Accept", "application/json") + + d.logger.Debug("discovering OAuth endpoints and scopes", + "well_known_url", wellKnownURL) + + resp, err := d.httpClient.Do(req) + if err != nil { + return "", "", nil, fmt.Errorf("failed to send discovery request: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return "", "", nil, fmt.Errorf("discovery endpoint returned status %d: %s", resp.StatusCode, string(body)) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", "", nil, fmt.Errorf("failed to read discovery response: %w", err) + } + + var metadata struct { + AuthorizationServers []string `json:"authorization_servers"` + ScopesSupported []string `json:"scopes_supported"` + } + + if err := json.Unmarshal(body, &metadata); err != nil { + return "", "", nil, fmt.Errorf("failed to parse discovery response: %w", err) + } + + if len(metadata.AuthorizationServers) == 0 { + return "", "", nil, fmt.Errorf("no authorization servers found in discovery response") + } + + // Derive endpoints from authorization_servers[0]. + // Example: "https://github.com/login/oauth" → + // authorization_endpoint: https://github.com/login/oauth/authorize + // token_endpoint: https://github.com/login/oauth/access_token + authServerBase := strings.TrimSuffix(metadata.AuthorizationServers[0], "/") + authorizationEndpoint := authServerBase + "/authorize" + tokenEndpoint := authServerBase + "/access_token" + + // Apply per-server endpoint overrides (partial — only endpoints, not full skip-discovery path) + if serverCfg, ok := d.config.ResourceConfigs[resourceURL]; ok { + if serverCfg.AuthorizationEndpoint != "" { + authorizationEndpoint = serverCfg.AuthorizationEndpoint + } + if serverCfg.TokenEndpoint != "" { + tokenEndpoint = serverCfg.TokenEndpoint + } + } + + // Resolve scopes: ResourceConfigs > ResourceScopes > discovered + scopes := metadata.ScopesSupported + if serverCfg, ok := d.config.ResourceConfigs[resourceURL]; ok && len(serverCfg.Scopes) > 0 { + d.logger.Info("Using configured scopes for resource server (overriding discovered)", + "resource_url", resourceURL, + "discovered_scopes", scopes, + "configured_scopes", serverCfg.Scopes) + scopes = serverCfg.Scopes + } else if override, ok := d.config.ResourceScopes[resourceURL]; ok { + d.logger.Info("Using legacy configured scopes for resource server (overriding discovered)", + "resource_url", resourceURL, + "discovered_scopes", scopes, + "configured_scopes", override) + scopes = override + } + + d.logger.Info("OAuth endpoints and scopes resolved", + "resource_url", resourceURL, + "authorization_endpoint", authorizationEndpoint, + "token_endpoint", tokenEndpoint, + "scopes", scopes) + + return authorizationEndpoint, tokenEndpoint, scopes, nil +} + +// SendSyntheticRequest sends an unauthenticated resource request to trigger OAuth elicitation. +// This is optional and used for validation purposes. +// Returns true if the server responds with 401, false otherwise. +func (d *Discoverer) SendSyntheticRequest(ctx context.Context, resourceURL string) (bool, error) { + // Build synthetic tools/list request + requestBody := map[string]interface{}{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/list", + "params": map[string]interface{}{}, + } + + bodyBytes, err := json.Marshal(requestBody) + if err != nil { + return false, fmt.Errorf("failed to marshal request: %w", err) + } + + // Send POST request to /mcp endpoint without Authorization header + resourceEndpoint := strings.TrimSuffix(resourceURL, "/") + "/mcp" + req, err := http.NewRequestWithContext(ctx, http.MethodPost, resourceEndpoint, strings.NewReader(string(bodyBytes))) + if err != nil { + return false, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + + d.logger.Debug("sending synthetic resource request", + "endpoint", resourceEndpoint, + "method", "tools/list") + + resp, err := d.httpClient.Do(req) + if err != nil { + return false, fmt.Errorf("failed to send request: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + // Read response body for logging + body, _ := io.ReadAll(resp.Body) + + d.logger.Debug("synthetic resource request response", + "status", resp.StatusCode, + "www_authenticate", resp.Header.Get("WWW-Authenticate"), + "body_length", len(body)) + + // Check if we got 401 Unauthorized + if resp.StatusCode == http.StatusUnauthorized { + d.logger.Info("OAuth elicitation triggered (401 response)", + "resource_url", resourceURL) + return true, nil + } + + // If we got a different status, OAuth might not be required + d.logger.Warn("expected 401 but got different status", + "status", resp.StatusCode, + "resource_url", resourceURL) + + return false, nil +} diff --git a/token-broker/internal/oauth/discovery_test.go b/token-broker/internal/oauth/discovery_test.go new file mode 100644 index 00000000..ef0d9b2e --- /dev/null +++ b/token-broker/internal/oauth/discovery_test.go @@ -0,0 +1,231 @@ +package oauth + +import ( + "context" + "encoding/json" + "log/slog" + "net/http" + "net/http/httptest" + "testing" +) + +func TestDiscoverEndpoints_WithConfiguredMetadata(t *testing.T) { + logger := slog.Default() + + // Create discoverer with configured metadata + config := &DiscovererConfig{ + AuthorizationEndpoint: "https://configured.example.com/oauth/authorize", + TokenEndpoint: "https://configured.example.com/oauth/token", + ScopesSupported: []string{"read", "write"}, + } + discoverer := NewDiscovererWithConfig(config, logger) + + // Call DiscoverEndpoints - should use configured values without making HTTP request + ctx := context.Background() + authEndpoint, tokenEndpoint, scopes, err := discoverer.DiscoverEndpoints(ctx, "https://mcp.example.com") + + if err != nil { + t.Fatalf("Expected no error, got: %v", err) + } + + if authEndpoint != config.AuthorizationEndpoint { + t.Errorf("Expected authEndpoint '%s', got '%s'", config.AuthorizationEndpoint, authEndpoint) + } + + if tokenEndpoint != config.TokenEndpoint { + t.Errorf("Expected tokenEndpoint '%s', got '%s'", config.TokenEndpoint, tokenEndpoint) + } + + if len(scopes) != len(config.ScopesSupported) { + t.Errorf("Expected %d scopes, got %d", len(config.ScopesSupported), len(scopes)) + } + + for i, scope := range config.ScopesSupported { + if i >= len(scopes) || scopes[i] != scope { + t.Errorf("Expected scope[%d] '%s', got '%s'", i, scope, scopes[i]) + } + } +} + +func TestDiscoverEndpoints_WithoutConfiguredMetadata_Success(t *testing.T) { + logger := slog.Default() + + // Create test server that returns discovery metadata + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/.well-known/oauth-protected-resource" { + t.Errorf("Expected path '/.well-known/oauth-protected-resource', got '%s'", r.URL.Path) + http.Error(w, "Not Found", http.StatusNotFound) + return + } + + response := map[string]interface{}{ + "authorization_servers": []string{"https://github.com/login/oauth"}, + "scopes_supported": []string{"repo", "user", "read:org"}, + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + // Create discoverer without configured metadata + discoverer := NewDiscoverer(logger) + + // Call DiscoverEndpoints - should make HTTP request + ctx := context.Background() + authEndpoint, tokenEndpoint, scopes, err := discoverer.DiscoverEndpoints(ctx, server.URL) + + if err != nil { + t.Fatalf("Expected no error, got: %v", err) + } + + expectedAuthEndpoint := "https://github.com/login/oauth/authorize" + if authEndpoint != expectedAuthEndpoint { + t.Errorf("Expected authEndpoint '%s', got '%s'", expectedAuthEndpoint, authEndpoint) + } + + expectedTokenEndpoint := "https://github.com/login/oauth/access_token" + if tokenEndpoint != expectedTokenEndpoint { + t.Errorf("Expected tokenEndpoint '%s', got '%s'", expectedTokenEndpoint, tokenEndpoint) + } + + expectedScopes := []string{"repo", "user", "read:org"} + if len(scopes) != len(expectedScopes) { + t.Errorf("Expected %d scopes, got %d", len(expectedScopes), len(scopes)) + } + + for i, scope := range expectedScopes { + if i >= len(scopes) || scopes[i] != scope { + t.Errorf("Expected scope[%d] '%s', got '%s'", i, scope, scopes[i]) + } + } +} + +func TestDiscoverEndpoints_PartialConfiguredMetadata(t *testing.T) { + logger := slog.Default() + + // Create test server that returns discovery metadata + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + response := map[string]interface{}{ + "authorization_servers": []string{"https://github.com/login/oauth"}, + "scopes_supported": []string{"repo", "user"}, + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + // Create discoverer with only authorization endpoint configured (missing token endpoint) + config := &DiscovererConfig{ + AuthorizationEndpoint: "https://configured.example.com/oauth/authorize", + // TokenEndpoint is missing - should trigger discovery + } + discoverer := NewDiscovererWithConfig(config, logger) + + // Call DiscoverEndpoints - should make HTTP request because config is incomplete + ctx := context.Background() + authEndpoint, tokenEndpoint, _, err := discoverer.DiscoverEndpoints(ctx, server.URL) + + if err != nil { + t.Fatalf("Expected no error, got: %v", err) + } + + // Should use discovered values, not configured ones + expectedAuthEndpoint := "https://github.com/login/oauth/authorize" + if authEndpoint != expectedAuthEndpoint { + t.Errorf("Expected authEndpoint '%s', got '%s'", expectedAuthEndpoint, authEndpoint) + } + + expectedTokenEndpoint := "https://github.com/login/oauth/access_token" + if tokenEndpoint != expectedTokenEndpoint { + t.Errorf("Expected tokenEndpoint '%s', got '%s'", expectedTokenEndpoint, tokenEndpoint) + } +} + +func TestDiscoverEndpoints_DiscoveryFailure(t *testing.T) { + logger := slog.Default() + + // Create test server that returns error + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "Internal Server Error", http.StatusInternalServerError) + })) + defer server.Close() + + // Create discoverer without configured metadata + discoverer := NewDiscoverer(logger) + + // Call DiscoverEndpoints - should fail + ctx := context.Background() + _, _, _, err := discoverer.DiscoverEndpoints(ctx, server.URL) + + if err == nil { + t.Fatal("Expected error, got none") + } +} + +func TestNewDiscoverer(t *testing.T) { + logger := slog.Default() + + discoverer := NewDiscoverer(logger) + + if discoverer == nil { + t.Fatal("Expected non-nil discoverer") + } + + if discoverer.config == nil { + t.Fatal("Expected non-nil config") + } + + // Config should be empty + if discoverer.config.AuthorizationEndpoint != "" { + t.Errorf("Expected empty AuthorizationEndpoint, got '%s'", discoverer.config.AuthorizationEndpoint) + } + + if discoverer.config.TokenEndpoint != "" { + t.Errorf("Expected empty TokenEndpoint, got '%s'", discoverer.config.TokenEndpoint) + } +} + +func TestNewDiscovererWithConfig(t *testing.T) { + logger := slog.Default() + + config := &DiscovererConfig{ + AuthorizationEndpoint: "https://example.com/oauth/authorize", + TokenEndpoint: "https://example.com/oauth/token", + ScopesSupported: []string{"read", "write"}, + } + + discoverer := NewDiscovererWithConfig(config, logger) + + if discoverer == nil { + t.Fatal("Expected non-nil discoverer") + } + + if discoverer.config != config { + t.Error("Expected config to be set") + } + + if discoverer.config.AuthorizationEndpoint != config.AuthorizationEndpoint { + t.Errorf("Expected AuthorizationEndpoint '%s', got '%s'", config.AuthorizationEndpoint, discoverer.config.AuthorizationEndpoint) + } +} + +func TestNewDiscovererWithConfig_NilConfig(t *testing.T) { + logger := slog.Default() + + discoverer := NewDiscovererWithConfig(nil, logger) + + if discoverer == nil { + t.Fatal("Expected non-nil discoverer") + } + + if discoverer.config == nil { + t.Fatal("Expected non-nil config (should be initialized to empty)") + } + + // Config should be empty + if discoverer.config.AuthorizationEndpoint != "" { + t.Errorf("Expected empty AuthorizationEndpoint, got '%s'", discoverer.config.AuthorizationEndpoint) + } +} diff --git a/token-broker/internal/session/manager.go b/token-broker/internal/session/manager.go new file mode 100644 index 00000000..b4a5b4cf --- /dev/null +++ b/token-broker/internal/session/manager.go @@ -0,0 +1,355 @@ +package session + +import ( + "fmt" + "log/slog" + "sync" + "time" + + "github.com/kagenti/token-broker/internal/core" +) + +// SessionManager manages OAuth sessions and their lifecycle. +type SessionManager struct { + sessions map[string]*core.Session + userSessions map[string][]string // userID -> session keys + mu sync.RWMutex + sessionTimeout time.Duration + maxSessionsPerUser int + clock core.Clock + logger *slog.Logger + timers map[string]*time.Timer // sessionKey -> timer + shutdownChan chan struct{} +} + +// NewSessionManager creates a new session manager. +func NewSessionManager(sessionTimeout time.Duration, maxSessionsPerUser int, clock core.Clock, logger *slog.Logger) *SessionManager { + return &SessionManager{ + sessions: make(map[string]*core.Session), + userSessions: make(map[string][]string), + sessionTimeout: sessionTimeout, + maxSessionsPerUser: maxSessionsPerUser, + clock: clock, + logger: logger, + timers: make(map[string]*time.Timer), + shutdownChan: make(chan struct{}), + } +} + +// CreateSession creates a new session with a provided session key and backend redirect URL. +func (sm *SessionManager) CreateSession(sessionKey, userID, backendRedirectURL string) error { + sm.mu.Lock() + defer sm.mu.Unlock() + + // Check max sessions per user + if existingSessions, ok := sm.userSessions[userID]; ok { + if len(existingSessions) >= sm.maxSessionsPerUser { + sm.logger.Warn("Max sessions per user exceeded", + "user_id", userID, + "current_sessions", len(existingSessions), + "max_sessions", sm.maxSessionsPerUser) + return fmt.Errorf("max sessions per user exceeded") + } + } + + // Create session with provided key + now := sm.clock.Now() + session := &core.Session{ + SessionKey: sessionKey, + UserID: userID, + BackendSessionRedirectURL: backendRedirectURL, + CreatedAt: now, + LastPollAt: now, + ExpirationDeadline: now.Add(sm.sessionTimeout), + AcquisitionSemaphore: NewSemaphore(1), + EventWaiters: make(chan core.Event, 1), + TokenWaiters: make(map[string][]chan core.TokenResult), + Done: make(chan struct{}), + ActiveOAuthTx: nil, + } + + sm.sessions[sessionKey] = session + + // Track user sessions + if sm.userSessions[userID] == nil { + sm.userSessions[userID] = []string{} + } + sm.userSessions[userID] = append(sm.userSessions[userID], sessionKey) + + sm.logger.Info("Session created", + "session_key", sessionKey, + "user_id", userID, + "backend_redirect_url", backendRedirectURL) + + return nil +} + +// GetSession retrieves a session by its key. +func (sm *SessionManager) GetSession(sessionKey string) (*core.Session, error) { + sm.mu.RLock() + defer sm.mu.RUnlock() + + session, ok := sm.sessions[sessionKey] + if !ok { + return nil, fmt.Errorf("session not found") + } + + return session, nil +} + +// GetSessionByState retrieves a session by its OAuth state parameter. +// The state format is: sessionKey.nonce (where nonce is base64url encoded, so no '.' in it) +// We find the LAST '.' to split sessionKey from nonce. +func (sm *SessionManager) GetSessionByState(state string) (*core.Session, error) { + // Find the last '.' to split sessionKey from nonce + lastDot := -1 + for i := len(state) - 1; i >= 0; i-- { + if state[i] == '.' { + lastDot = i + break + } + } + + if lastDot == -1 { + return nil, fmt.Errorf("invalid state format: no separator found") + } + + sessionKey := state[:lastDot] + + sm.mu.RLock() + defer sm.mu.RUnlock() + + session, ok := sm.sessions[sessionKey] + if !ok { + return nil, fmt.Errorf("session not found for state") + } + + // Verify the state matches (security check) + if session.ActiveOAuthTx == nil || session.ActiveOAuthTx.State != state { + return nil, fmt.Errorf("state mismatch") + } + + return session, nil +} + +// ValidateSession checks if a session exists and belongs to the specified user. +func (sm *SessionManager) ValidateSession(sessionKey, userID string) error { + sm.mu.RLock() + defer sm.mu.RUnlock() + + session, ok := sm.sessions[sessionKey] + if !ok { + return fmt.Errorf("session not found") + } + + if session.UserID != userID { + sm.logger.Warn("Session user mismatch", + "session_key", sessionKey, + "expected_user", session.UserID, + "provided_user", userID) + return fmt.Errorf("session user mismatch") + } + + return nil +} + +// EndSession terminates a session and releases all resources. +func (sm *SessionManager) EndSession(sessionKey string) error { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[sessionKey] + if !ok { + return fmt.Errorf("session not found") + } + + sm.logger.Info("Ending session", + "session_key", sessionKey, + "user_id", session.UserID) + + sm.cleanupSessionLocked(sessionKey, session, "session ended") + + return nil +} + +// ExpireSession marks a session as expired and fails all pending requests. +func (sm *SessionManager) ExpireSession(sessionKey string) { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[sessionKey] + if !ok { + return + } + + sm.logger.Info("Session expired", + "session_key", sessionKey, + "user_id", session.UserID) + + sm.cleanupSessionLocked(sessionKey, session, "session expired") +} + +// cleanupSessionLocked cleans up a session and notifies all waiters. +// Must be called with sm.mu held. +func (sm *SessionManager) cleanupSessionLocked(sessionKey string, session *core.Session, reason string) { + // Stop timer if exists + if timer, ok := sm.timers[sessionKey]; ok { + timer.Stop() + delete(sm.timers, sessionKey) + } + + // Fail all token waiters + for resourceURL, waiters := range session.TokenWaiters { + for _, waiter := range waiters { + select { + case waiter <- core.TokenResult{ + Token: "", + Error: fmt.Errorf("%s", reason), + }: + default: + } + close(waiter) + } + sm.logger.Debug("Failed token waiters", + "session_key", sessionKey, + "resource_url", resourceURL, + "waiter_count", len(waiters)) + } + + // Close Done to unblock any waiters (e.g. AcquireToken blocked on the + // OAuth completion channel). cleanupSessionLocked runs at most once per + // key, so this close is safe. + close(session.Done) + + // Send a final error event so a connected Backend long-poll learns the + // session ended. EventWaiters is buffered (cap 1); the send is best-effort. + // + // We intentionally do NOT close(EventWaiters): performOAuthFlow may still be + // sending to it concurrently (it does not hold sm.mu), and closing a channel + // during a concurrent send is a data race that panics. The channel is + // garbage-collected with the session; Done is the broadcast that unblocks + // waiters. + select { + case session.EventWaiters <- core.Event{ + Type: "error", + Message: reason, + Code: "session_expired", + }: + default: + } + + // Remove from user sessions + if userSessions, ok := sm.userSessions[session.UserID]; ok { + newSessions := []string{} + for _, key := range userSessions { + if key != sessionKey { + newSessions = append(newSessions, key) + } + } + if len(newSessions) > 0 { + sm.userSessions[session.UserID] = newSessions + } else { + delete(sm.userSessions, session.UserID) + } + } + + // Remove session + delete(sm.sessions, sessionKey) + + sm.logger.Debug("Session cleaned up", + "session_key", sessionKey, + "reason", reason) +} + +// ResetSessionTimer resets the idle timeout timer for a session. +func (sm *SessionManager) ResetSessionTimer(sessionKey string) { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[sessionKey] + if !ok { + return + } + + // Stop existing timer + if timer, ok := sm.timers[sessionKey]; ok { + timer.Stop() + delete(sm.timers, sessionKey) + } + + // Update last poll time + session.LastPollAt = sm.clock.Now() + session.ExpirationDeadline = session.LastPollAt.Add(sm.sessionTimeout) + + sm.logger.Debug("Session timer reset", + "session_key", sessionKey, + "user_id", session.UserID) +} + +// StartSessionTimer starts the idle timeout timer for a session. +func (sm *SessionManager) StartSessionTimer(sessionKey string) { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[sessionKey] + if !ok { + return + } + + // Stop existing timer + if timer, ok := sm.timers[sessionKey]; ok { + timer.Stop() + } + + // Start new timer + timer := time.AfterFunc(sm.sessionTimeout, func() { + sm.ExpireSession(sessionKey) + }) + sm.timers[sessionKey] = timer + + sm.logger.Debug("Session timer started", + "session_key", sessionKey, + "user_id", session.UserID, + "timeout", sm.sessionTimeout) +} + +// Shutdown stops all timers and cleans up all sessions. +func (sm *SessionManager) Shutdown() { + sm.mu.Lock() + defer sm.mu.Unlock() + + close(sm.shutdownChan) + + sm.logger.Info("Shutting down session manager", + "active_sessions", len(sm.sessions)) + + // Stop all timers + for _, timer := range sm.timers { + timer.Stop() + } + sm.timers = make(map[string]*time.Timer) + + // Clean up all sessions + for sessionKey, session := range sm.sessions { + sm.cleanupSessionLocked(sessionKey, session, "shutdown") + } + + sm.logger.Info("Session manager shutdown complete") +} + +// GetSessionCount returns the number of active sessions. +func (sm *SessionManager) GetSessionCount() int { + sm.mu.RLock() + defer sm.mu.RUnlock() + return len(sm.sessions) +} + +// GetUserSessionCount returns the number of active sessions for a user. +func (sm *SessionManager) GetUserSessionCount(userID string) int { + sm.mu.RLock() + defer sm.mu.RUnlock() + if sessions, ok := sm.userSessions[userID]; ok { + return len(sessions) + } + return 0 +} diff --git a/token-broker/internal/session/manager_test.go b/token-broker/internal/session/manager_test.go new file mode 100644 index 00000000..947ebfe3 --- /dev/null +++ b/token-broker/internal/session/manager_test.go @@ -0,0 +1,465 @@ +package session + +import ( + "testing" + "time" + + "log/slog" + "os" + + "github.com/google/uuid" + "github.com/kagenti/token-broker/internal/core" +) + +// FakeClock for testing +type FakeClock struct { + now time.Time +} + +func (c *FakeClock) Now() time.Time { + return c.now +} + +func (c *FakeClock) After(d time.Duration) <-chan time.Time { + ch := make(chan time.Time, 1) + ch <- c.now.Add(d) + return ch +} + +func (c *FakeClock) Advance(d time.Duration) { + c.now = c.now.Add(d) +} + +func TestSessionManager_CreateSession(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + err := sm.CreateSession(sessionKey, userID, "") + if err != nil { + t.Fatalf("CreateSession failed: %v", err) + } + + // Verify session exists + session, err := sm.GetSession(sessionKey) + if err != nil { + t.Fatalf("GetSession failed: %v", err) + } + + if session.UserID != userID { + t.Errorf("Expected user ID %s, got %s", userID, session.UserID) + } +} + +func TestSessionManager_ValidateSession(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + _ = sm.CreateSession(sessionKey, userID, "") + + // Valid user + err := sm.ValidateSession(sessionKey, userID) + if err != nil { + t.Errorf("ValidateSession should succeed for correct user: %v", err) + } + + // Wrong user + err = sm.ValidateSession(sessionKey, "user2") + if err == nil { + t.Error("ValidateSession should fail for wrong user") + } + + // Non-existent session + err = sm.ValidateSession("non-existent", userID) + if err == nil { + t.Error("ValidateSession should fail for non-existent session") + } +} + +func TestSessionManager_MaxSessionsPerUser(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + maxSessions := 3 + sm := NewSessionManager(60*time.Second, maxSessions, clock, logger) + defer sm.Shutdown() + + userID := "user1" + + // Create max sessions + for i := 0; i < maxSessions; i++ { + sessionKey := uuid.New().String() + err := sm.CreateSession(sessionKey, userID, "") + if err != nil { + t.Fatalf("CreateSession %d failed: %v", i, err) + } + } + + // Next session should fail + sessionKey := uuid.New().String() + err := sm.CreateSession(sessionKey, userID, "") + if err == nil { + t.Error("CreateSession should fail when max sessions exceeded") + } +} + +func TestSessionManager_EndSession(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + _ = sm.CreateSession(sessionKey, userID, "") + + // Capture Done before ending so we can assert waiters are unblocked. + session, err := sm.GetSession(sessionKey) + if err != nil { + t.Fatalf("GetSession failed: %v", err) + } + + // End session + err = sm.EndSession(sessionKey) + if err != nil { + t.Fatalf("EndSession failed: %v", err) + } + + // Done should be closed so blocked waiters (e.g. AcquireToken) return. + select { + case <-session.Done: + default: + t.Error("EndSession should close session.Done") + } + + // Session should no longer exist + _, err = sm.GetSession(sessionKey) + if err == nil { + t.Error("GetSession should fail for ended session") + } +} + +func TestSessionManager_ExpireSession(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + _ = sm.CreateSession(sessionKey, userID, "") + + // Capture Done before expiring so we can assert waiters are unblocked. + session, err := sm.GetSession(sessionKey) + if err != nil { + t.Fatalf("GetSession failed: %v", err) + } + + // Expire session + sm.ExpireSession(sessionKey) + + // Done should be closed so blocked waiters (e.g. AcquireToken) return. + select { + case <-session.Done: + default: + t.Error("ExpireSession should close session.Done") + } + + // Session should no longer exist + _, err = sm.GetSession(sessionKey) + if err == nil { + t.Error("GetSession should fail for expired session") + } +} + +func TestSessionManager_SessionTimer(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + timeout := 100 * time.Millisecond + sm := NewSessionManager(timeout, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + _ = sm.CreateSession(sessionKey, userID, "") + + // Start timer + sm.StartSessionTimer(sessionKey) + + // Wait for timeout + time.Sleep(timeout + 50*time.Millisecond) + + // Session should be expired + _, err := sm.GetSession(sessionKey) + if err == nil { + t.Error("Session should be expired after timeout") + } +} + +func TestSessionManager_ResetSessionTimer(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + timeout := 150 * time.Millisecond + sm := NewSessionManager(timeout, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + _ = sm.CreateSession(sessionKey, userID, "") + + // Start timer + sm.StartSessionTimer(sessionKey) + + // Wait half the timeout + time.Sleep(timeout / 2) + + // Reset timer + sm.ResetSessionTimer(sessionKey) + + // Wait another half timeout (total time > original timeout) + time.Sleep(timeout / 2) + + // Session should still exist (timer was reset) + _, err := sm.GetSession(sessionKey) + if err != nil { + t.Error("Session should still exist after timer reset") + } + + // Start a new timer after reset + sm.StartSessionTimer(sessionKey) + + // Wait for full timeout + time.Sleep(timeout + 100*time.Millisecond) + + // Now session should be expired + _, err = sm.GetSession(sessionKey) + if err == nil { + t.Error("Session should be expired after full timeout") + } +} + +func TestSessionManager_GetSessionByState(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := "test-session-key" + backendRedirectURL := "https://backend.example.com/oauth/complete" + + // Create session + err := sm.CreateSession(sessionKey, userID, backendRedirectURL) + if err != nil { + t.Fatalf("CreateSession failed: %v", err) + } + + // Set up OAuth transaction with state + session, _ := sm.GetSession(sessionKey) + session.ActiveOAuthTx = &core.OAuthTransaction{ + State: sessionKey + ".abc123def456", + } + + // Test valid state + retrievedSession, err := sm.GetSessionByState(sessionKey + ".abc123def456") + if err != nil { + t.Fatalf("GetSessionByState failed: %v", err) + } + if retrievedSession.SessionKey != sessionKey { + t.Errorf("Expected session key %s, got %s", sessionKey, retrievedSession.SessionKey) + } + + // Test state with dots in session key + sessionKeyWithDots := "session.with.dots" + err = sm.CreateSession(sessionKeyWithDots, userID, backendRedirectURL) + if err != nil { + t.Fatalf("CreateSession failed: %v", err) + } + + session2, _ := sm.GetSession(sessionKeyWithDots) + session2.ActiveOAuthTx = &core.OAuthTransaction{ + State: sessionKeyWithDots + ".xyz789", + } + + retrievedSession2, err := sm.GetSessionByState(sessionKeyWithDots + ".xyz789") + if err != nil { + t.Fatalf("GetSessionByState failed for session key with dots: %v", err) + } + if retrievedSession2.SessionKey != sessionKeyWithDots { + t.Errorf("Expected session key %s, got %s", sessionKeyWithDots, retrievedSession2.SessionKey) + } + + // Test invalid state (no dot) + _, err = sm.GetSessionByState("invalid-state-no-dot") + if err == nil { + t.Error("GetSessionByState should fail for state without dot") + } + + // Test non-existent session + _, err = sm.GetSessionByState("nonexistent.nonce") + if err == nil { + t.Error("GetSessionByState should fail for non-existent session") + } +} + +func TestSessionManager_GetSessionCount(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + if count := sm.GetSessionCount(); count != 0 { + t.Errorf("Expected session count 0, got %d", count) + } + + // Create sessions + _ = sm.CreateSession(uuid.New().String(), "user1", "") + _ = sm.CreateSession(uuid.New().String(), "user1", "") + _ = sm.CreateSession(uuid.New().String(), "user2", "") + + if count := sm.GetSessionCount(); count != 3 { + t.Errorf("Expected session count 3, got %d", count) + } +} + +func TestSessionManager_GetUserSessionCount(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + user1 := "user1" + user2 := "user2" + + // Create sessions + _ = sm.CreateSession(uuid.New().String(), user1, "") + _ = sm.CreateSession(uuid.New().String(), user1, "") + _ = sm.CreateSession(uuid.New().String(), user2, "") + + if count := sm.GetUserSessionCount(user1); count != 2 { + t.Errorf("Expected user1 session count 2, got %d", count) + } + + if count := sm.GetUserSessionCount(user2); count != 1 { + t.Errorf("Expected user2 session count 1, got %d", count) + } + + if count := sm.GetUserSessionCount("user3"); count != 0 { + t.Errorf("Expected user3 session count 0, got %d", count) + } +} + +func TestSessionManager_Shutdown(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + + // Create sessions + _ = sm.CreateSession(uuid.New().String(), "user1", "") + _ = sm.CreateSession(uuid.New().String(), "user2", "") + + if count := sm.GetSessionCount(); count != 2 { + t.Errorf("Expected session count 2, got %d", count) + } + + // Shutdown + sm.Shutdown() + + // All sessions should be cleaned up + if count := sm.GetSessionCount(); count != 0 { + t.Errorf("Expected session count 0 after shutdown, got %d", count) + } +} + +func TestSessionManager_TokenWaiters(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + _ = sm.CreateSession(sessionKey, userID, "") + + session, _ := sm.GetSession(sessionKey) + + // Add token waiters + resource := "http://mcp.example.com" + waiter1 := make(chan core.TokenResult, 1) + waiter2 := make(chan core.TokenResult, 1) + + session.TokenWaiters[resource] = []chan core.TokenResult{waiter1, waiter2} + + // End session (should notify waiters) + _ = sm.EndSession(sessionKey) + + // Verify waiters were notified + select { + case result := <-waiter1: + if result.Error == nil { + t.Error("Waiter should receive error on session end") + } + case <-time.After(100 * time.Millisecond): + t.Error("Waiter1 was not notified") + } + + select { + case result := <-waiter2: + if result.Error == nil { + t.Error("Waiter should receive error on session end") + } + case <-time.After(100 * time.Millisecond): + t.Error("Waiter2 was not notified") + } +} + +func TestSessionManager_EventWaiters(t *testing.T) { + clock := &FakeClock{now: time.Now()} + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError})) + sm := NewSessionManager(60*time.Second, 5, clock, logger) + defer sm.Shutdown() + + userID := "user1" + sessionKey := uuid.New().String() + _ = sm.CreateSession(sessionKey, userID, "") + + session, _ := sm.GetSession(sessionKey) + + // Start goroutine waiting for event + eventReceived := make(chan bool, 1) + go func() { + select { + case event := <-session.EventWaiters: + if event.Type == "error" && event.Code == "session_expired" { + eventReceived <- true + } + case <-time.After(200 * time.Millisecond): + eventReceived <- false + } + }() + + // Give goroutine time to start waiting + time.Sleep(50 * time.Millisecond) + + // End session (should send error event) + _ = sm.EndSession(sessionKey) + + // Verify event was received + select { + case received := <-eventReceived: + if !received { + t.Error("Event waiter should receive error event on session end") + } + case <-time.After(300 * time.Millisecond): + t.Error("Timeout waiting for event") + } +} + +// Verify FakeClock implements core.Clock +var _ core.Clock = (*FakeClock)(nil) diff --git a/token-broker/internal/session/semaphore.go b/token-broker/internal/session/semaphore.go new file mode 100644 index 00000000..97124975 --- /dev/null +++ b/token-broker/internal/session/semaphore.go @@ -0,0 +1,49 @@ +// Package session provides session management for the Token Broker. +package session + +import ( + "context" + "fmt" +) + +// SimpleSemaphore implements a simple semaphore with capacity 1. +type SimpleSemaphore struct { + ch chan struct{} +} + +// NewSemaphore creates a new semaphore with the specified capacity. +func NewSemaphore(capacity int) *SimpleSemaphore { + return &SimpleSemaphore{ + ch: make(chan struct{}, capacity), + } +} + +// Acquire acquires the semaphore, blocking until available or context is cancelled. +func (s *SimpleSemaphore) Acquire(ctx context.Context) error { + select { + case s.ch <- struct{}{}: + return nil + case <-ctx.Done(): + return fmt.Errorf("semaphore acquisition cancelled: %w", ctx.Err()) + } +} + +// Release releases the semaphore. +func (s *SimpleSemaphore) Release() { + select { + case <-s.ch: + default: + // Semaphore was not acquired, ignore + } +} + +// TryAcquire attempts to acquire the semaphore without blocking. +// Returns true if acquired, false otherwise. +func (s *SimpleSemaphore) TryAcquire() bool { + select { + case s.ch <- struct{}{}: + return true + default: + return false + } +} diff --git a/token-broker/internal/session/semaphore_test.go b/token-broker/internal/session/semaphore_test.go new file mode 100644 index 00000000..833fc064 --- /dev/null +++ b/token-broker/internal/session/semaphore_test.go @@ -0,0 +1,176 @@ +package session + +import ( + "context" + "testing" + "time" +) + +func TestSemaphore_AcquireRelease(t *testing.T) { + sem := NewSemaphore(1) + + ctx := context.Background() + + // Acquire semaphore + err := sem.Acquire(ctx) + if err != nil { + t.Fatalf("Acquire failed: %v", err) + } + + // Release semaphore + sem.Release() + + // Should be able to acquire again + err = sem.Acquire(ctx) + if err != nil { + t.Fatalf("Second acquire failed: %v", err) + } + + sem.Release() +} + +func TestSemaphore_BlocksWhenFull(t *testing.T) { + sem := NewSemaphore(1) + + ctx := context.Background() + + // Acquire semaphore + err := sem.Acquire(ctx) + if err != nil { + t.Fatalf("Acquire failed: %v", err) + } + + // Try to acquire again (should block) + acquired := make(chan bool, 1) + go func() { + ctx2, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + err := sem.Acquire(ctx2) + acquired <- (err == nil) + }() + + // Wait a bit to ensure goroutine is blocked + time.Sleep(50 * time.Millisecond) + + // Release semaphore + sem.Release() + + // Second acquire should now succeed + select { + case success := <-acquired: + if !success { + t.Error("Second acquire should have succeeded after release") + } + case <-time.After(200 * time.Millisecond): + t.Error("Second acquire timed out") + } + + sem.Release() // Clean up +} + +func TestSemaphore_ContextCancellation(t *testing.T) { + sem := NewSemaphore(1) + + // Acquire semaphore + ctx := context.Background() + err := sem.Acquire(ctx) + if err != nil { + t.Fatalf("Acquire failed: %v", err) + } + + // Try to acquire with cancelled context + ctx2, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + err = sem.Acquire(ctx2) + if err == nil { + t.Error("Acquire should have failed with cancelled context") + } + + sem.Release() +} + +func TestSemaphore_TryAcquire(t *testing.T) { + sem := NewSemaphore(1) + + // TryAcquire should succeed + if !sem.TryAcquire() { + t.Error("TryAcquire should have succeeded") + } + + // TryAcquire should fail (semaphore full) + if sem.TryAcquire() { + t.Error("TryAcquire should have failed (semaphore full)") + } + + // Release and try again + sem.Release() + + if !sem.TryAcquire() { + t.Error("TryAcquire should have succeeded after release") + } + + sem.Release() +} + +func TestSemaphore_MultipleReleases(t *testing.T) { + sem := NewSemaphore(1) + + ctx := context.Background() + + // Acquire once + err := sem.Acquire(ctx) + if err != nil { + t.Fatalf("Acquire failed: %v", err) + } + + // Release multiple times (should not panic) + sem.Release() + sem.Release() // Extra release should be ignored + sem.Release() // Extra release should be ignored + + // Should still be able to acquire + err = sem.Acquire(ctx) + if err != nil { + t.Fatalf("Acquire after multiple releases failed: %v", err) + } + + sem.Release() +} + +func TestSemaphore_ConcurrentAcquire(t *testing.T) { + sem := NewSemaphore(1) + + const numGoroutines = 10 + acquired := make(chan int, numGoroutines) + + // Launch multiple goroutines trying to acquire + for i := 0; i < numGoroutines; i++ { + go func(id int) { + ctx := context.Background() + if err := sem.Acquire(ctx); err == nil { + acquired <- id + time.Sleep(10 * time.Millisecond) // Hold briefly + sem.Release() + } + }(i) + } + + // Collect results + seen := make(map[int]bool) + for i := 0; i < numGoroutines; i++ { + select { + case id := <-acquired: + if seen[id] { + t.Errorf("Goroutine %d acquired semaphore twice", id) + } + seen[id] = true + case <-time.After(2 * time.Second): + t.Fatal("Timeout waiting for goroutines") + } + } + + if len(seen) != numGoroutines { + t.Errorf("Expected %d goroutines to acquire, got %d", numGoroutines, len(seen)) + } +} diff --git a/token-broker/pkg/oauth/pkce.go b/token-broker/pkg/oauth/pkce.go new file mode 100644 index 00000000..850448a9 --- /dev/null +++ b/token-broker/pkg/oauth/pkce.go @@ -0,0 +1,50 @@ +// Package oauth provides OAuth utilities including PKCE support. +package oauth + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "fmt" +) + +// PKCEChallenge represents a PKCE challenge for OAuth 2.0. +type PKCEChallenge struct { + Verifier string + Challenge string + Method string +} + +// GeneratePKCEChallenge generates a PKCE code verifier and challenge. +// This implements RFC 7636 for enhanced security in OAuth flows. +func GeneratePKCEChallenge() (*PKCEChallenge, error) { + // Generate 32 random bytes for the verifier + verifierBytes := make([]byte, 32) + if _, err := rand.Read(verifierBytes); err != nil { + return nil, fmt.Errorf("failed to generate random verifier: %w", err) + } + + // Base64url encode the verifier (without padding) + verifier := base64.RawURLEncoding.EncodeToString(verifierBytes) + + // Create SHA256 hash of the verifier + hash := sha256.Sum256([]byte(verifier)) + + // Base64url encode the challenge (without padding) + challenge := base64.RawURLEncoding.EncodeToString(hash[:]) + + return &PKCEChallenge{ + Verifier: verifier, + Challenge: challenge, + Method: "S256", + }, nil +} + +// GenerateState generates a cryptographically secure random state parameter. +func GenerateState() (string, error) { + stateBytes := make([]byte, 32) + if _, err := rand.Read(stateBytes); err != nil { + return "", fmt.Errorf("failed to generate random state: %w", err) + } + return base64.RawURLEncoding.EncodeToString(stateBytes), nil +} diff --git a/token-broker/pkg/oauth/pkce_test.go b/token-broker/pkg/oauth/pkce_test.go new file mode 100644 index 00000000..cf1f63e7 --- /dev/null +++ b/token-broker/pkg/oauth/pkce_test.go @@ -0,0 +1,48 @@ +package oauth + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGeneratePKCEChallenge(t *testing.T) { + t.Parallel() + + pkce, err := GeneratePKCEChallenge() + require.NoError(t, err) + require.NotNil(t, pkce) + + // Verify verifier is not empty + assert.NotEmpty(t, pkce.Verifier) + + // Verify challenge is not empty + assert.NotEmpty(t, pkce.Challenge) + + // Verify method is S256 + assert.Equal(t, "S256", pkce.Method) + + // Verify verifier and challenge are different + assert.NotEqual(t, pkce.Verifier, pkce.Challenge) + + // Verify they are base64url encoded (no padding) + assert.NotContains(t, pkce.Verifier, "=") + assert.NotContains(t, pkce.Challenge, "=") +} + +func TestGenerateState(t *testing.T) { + t.Parallel() + + state, err := GenerateState() + require.NoError(t, err) + assert.NotEmpty(t, state) + + // Verify it's base64url encoded (no padding) + assert.NotContains(t, state, "=") + + // Generate another state and verify they're different + state2, err := GenerateState() + require.NoError(t, err) + assert.NotEqual(t, state, state2) +}