Squid as an Authenticated Explicit Forward Proxy: JWT Validation with Centralized Credential Injection

#!/usr/bin/env python3
"""
Squid basic-auth helper for JWT validation.
Reads lines of "username base64(dummy:jwt)" from stdin.
Writes "OK" or "ERR" to stdout, one response per input line.

Required environment variables:
  JWT_JWKS_URL  – URL of the JWKS endpoint (RS256)
      OR
  JWT_SECRET    – HMAC secret (HS256, dev/test only)
  JWT_AUDIENCE  – expected 'aud' claim
  JWT_ISSUER    – expected 'iss' claim
"""
import sys
import os
import base64
import json
import time
import hmac
import hashlib
import struct
import urllib.request

# --------------------------------------------------------------------------- #
# Minimal JWT decode + verify (no third-party deps required in POC container) #
# --------------------------------------------------------------------------- #

def _b64pad(s: str) -> str:
    return s + "=" * (-len(s) % 4)

def _decode_part(part: str) -> dict:
    return json.loads(base64.urlsafe_b64decode(_b64pad(part)))

def _verify_hs256(header_b64: str, payload_b64: str, sig_b64: str, secret: str) -> bool:
    msg = f"{header_b64}.{payload_b64}".encode()
    expected = base64.urlsafe_b64encode(
        hmac.new(secret.encode(), msg, hashlib.sha256).digest()
    ).rstrip(b"=").decode()
    return hmac.compare_digest(expected, sig_b64)

def _fetch_jwks(url: str) -> dict:
    with urllib.request.urlopen(url, timeout=5) as r:
        return json.loads(r.read())

# Simple in-process JWKS cache (keyed by kid)
_jwks_cache: dict = {}
_jwks_fetched_at: float = 0.0
_JWKS_TTL = 300  # seconds

def _get_jwk(kid: str, jwks_url: str) -> dict | None:
    global _jwks_cache, _jwks_fetched_at
    if time.time() - _jwks_fetched_at > _JWKS_TTL:
        jwks = _fetch_jwks(jwks_url)
        _jwks_cache = {k["kid"]: k for k in jwks.get("keys", [])}
        _jwks_fetched_at = time.time()
    return _jwks_cache.get(kid)

def _verify_rs256(header_b64: str, payload_b64: str, sig_b64: str, kid: str, jwks_url: str) -> bool:
    """
    Verify an RS256 JWT using the public key fetched from the JWKS endpoint.
    Uses the cryptography package which is pre-installed in the helper image.
    """
    try:
        from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers
        from cryptography.hazmat.primitives.asymmetric.padding import PKCS1v15
        from cryptography.hazmat.primitives.hashes import SHA256
        from cryptography.exceptions import InvalidSignature
    except ImportError:
        # Fail closed: without the cryptography package, RS256 verification is impossible.
        sys.stderr.write("ERROR: cryptography package not available; rejecting RS256 token\n")
        return False

    jwk = _get_jwk(kid, jwks_url)
    if not jwk:
        return False

    def _b64int(v: str) -> int:
        return int.from_bytes(base64.urlsafe_b64decode(_b64pad(v)), "big")

    pub = RSAPublicNumbers(_b64int(jwk["e"]), _b64int(jwk["n"])).public_key()
    msg = f"{header_b64}.{payload_b64}".encode()
    sig = base64.urlsafe_b64decode(_b64pad(sig_b64))
    try:
        pub.verify(sig, msg, PKCS1v15(), SHA256())
        return True
    except InvalidSignature:
        return False

def validate_jwt(token: str) -> tuple[bool, str]:
    """
    Returns (ok: bool, username: str).
    username is the 'sub' claim on success, empty string on failure.
    """
    parts = token.split(".")
    if len(parts) != 3:
        return False, ""
    header_b64, payload_b64, sig_b64 = parts

    try:
        header = _decode_part(header_b64)
        payload = _decode_part(payload_b64)
    except Exception:
        return False, ""

    # Expiry
    if payload.get("exp", 0) < time.time():
        sys.stderr.write(f"JWT expired for sub={payload.get('sub')}\n")
        return False, ""

    # Audience
    expected_aud = os.environ.get("JWT_AUDIENCE", "squid-proxy")
    aud = payload.get("aud", "")
    if isinstance(aud, list):
        if expected_aud not in aud:
            return False, ""
    elif aud != expected_aud:
        return False, ""

    # Issuer
    expected_iss = os.environ.get("JWT_ISSUER", "")
    if expected_iss and payload.get("iss") != expected_iss:
        return False, ""

    # Signature
    alg = header.get("alg", "")
    if alg == "HS256":
        secret = os.environ.get("JWT_SECRET", "")
        if not secret:
            sys.stderr.write("JWT_SECRET not set for HS256\n")
            return False, ""
        if not _verify_hs256(header_b64, payload_b64, sig_b64, secret):
            return False, ""
    elif alg == "RS256":
        jwks_url = os.environ.get("JWT_JWKS_URL", "")
        if not jwks_url:
            sys.stderr.write("JWT_JWKS_URL not set for RS256\n")
            return False, ""
        kid = header.get("kid", "")
        if not _verify_rs256(header_b64, payload_b64, sig_b64, kid, jwks_url):
            return False, ""
    else:
        sys.stderr.write(f"Unsupported algorithm: {alg}\n")
        return False, ""

    return True, payload.get("sub", "anonymous")

def main():
    sys.stdout.reconfigure(line_buffering=True)  # critical: Squid reads line-by-line
    sys.stderr.reconfigure(line_buffering=True)  # ensure error logs appear immediately
    for line in sys.stdin:
        line = line.rstrip("\n")
        # Squid sends:  <username> <password>
        # where password = base64(dummy:jwt) decoded from the Proxy-Authorization header
        parts = line.split(" ", 1)
        if len(parts) != 2:
            sys.stdout.write("ERR\n")
            continue
        _username_placeholder, password = parts
        # password is the raw token (Squid strips the base64 for us in basic mode)
        token = password.strip()
        ok, sub = validate_jwt(token)
        if ok:
            sys.stdout.write(f"OK user={sub}\n")
        else:
            sys.stdout.write("ERR\n")

if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Squid external ACL helper for destination policy enforcement.

Squid passes:  <user> <dst_host>
  (configured via external_acl_type fmt=%LOGIN %DST)

The helper reads a YAML policy file (POLICY_FILE env var) that looks like:

  rules:
    - users: ["ci-pipeline", "svc-analytics"]
      allow_destinations:
        - "api.github.com"
        - "api.datadoghq.com"
    - users: ["*"]
      allow_destinations:
        - "api.example-internal.com"

'*' in users matches any authenticated user.
"""
import sys
import os
import json

def load_policy(path: str) -> list[dict]:
    try:
        import yaml  # pyyaml installed in helper image
        with open(path) as f:
            return yaml.safe_load(f).get("rules", [])
    except ImportError:
        raise RuntimeError(
            "pyyaml is required to load the policy file. "
            "Install it with: pip install pyyaml"
        )

_policy_path = os.environ.get("POLICY_FILE", "/etc/squid/policy.yaml")
_rules: list[dict] = []

def _reload_policy():
    global _rules
    try:
        _rules = load_policy(_policy_path)
    except Exception as e:
        sys.stderr.write(f"Failed to load policy: {e}\n")
        _rules = []

_reload_policy()

def is_permitted(user: str, dst_host: str) -> bool:
    for rule in _rules:
        users = rule.get("users", [])
        if "*" in users or user in users:
            allowed = rule.get("allow_destinations", [])
            if dst_host in allowed or "*" in allowed:
                return True
    return False

def main():
    sys.stdout.reconfigure(line_buffering=True)
    for line in sys.stdin:
        line = line.rstrip("\n")
        parts = line.split(" ", 1)
        if len(parts) != 2:
            sys.stdout.write("ERR\n")
            continue
        user, dst = parts
        if is_permitted(user, dst):
            sys.stdout.write("OK\n")
        else:
            sys.stderr.write(f"DENY user={user} dst={dst}\n")
            sys.stdout.write("ERR\n")

if __name__ == "__main__":
    main()

mkdir -p squid-jwt-proxy/token-issuer
openssl genrsa -out squid-jwt-proxy/token-issuer/issuer_private_key.pem 2048
openssl rsa -in squid-jwt-proxy/token-issuer/issuer_private_key.pem \
    -pubout -out squid-jwt-proxy/token-issuer/issuer_public_key.pem

version: "3.9"

services:
  # ── Token issuer ─────────────────────────────────────────────────────────
  token-issuer:
    build: ./token-issuer
    ports:
      - "8080:8080"
    environment:
      PRIVATE_KEY_PATH: /app/issuer_private_key.pem
      PUBLIC_KEY_PATH:  /app/issuer_public_key.pem

  # ── Squid explicit forward proxy ─────────────────────────────────────────
  squid:
    build: ./squid
    ports:
      - "3128:3128"
    environment:
      # JWT validation
      JWT_JWKS_URL:  "http://token-issuer:8080/.well-known/jwks.json"
      JWT_AUDIENCE:  "squid-proxy"
      JWT_ISSUER:    "http://token-issuer:8080"
      # Upstream service credential injected by the proxy
      UPSTREAM_SVC_TOKEN: "svc-account-token-abc123"   # replace with real token
      # Policy file path (mounted via volume)
      POLICY_FILE:   "/etc/squid/policy.yaml"
    volumes:
      - ./squid/policy.yaml:/etc/squid/policy.yaml:ro
      - ./helpers:/usr/local/lib/squid/helpers:ro
    depends_on:
      - token-issuer

  # ── Echo server (stands in for the real upstream API) ────────────────────
  echo-server:
    image: ealen/echo-server:latest
    ports:
      - "3000:80"

FROM ubuntu:22.04

RUN apt-get update && apt-get install -y --no-install-recommends \
    squid \
    python3 \
    python3-pip \
    gettext-base \
    && pip3 install --no-cache-dir cryptography pyyaml \
    && rm -rf /var/lib/apt/lists/*

COPY entrypoint.sh /entrypoint.sh
COPY squid.conf.tmpl /etc/squid/squid.conf.tmpl
RUN chmod +x /entrypoint.sh

EXPOSE 3128
ENTRYPOINT ["/entrypoint.sh"]

#!/bin/bash
set -euo pipefail

# Render squid.conf from template, substituting environment variables
envsubst < /etc/squid/squid.conf.tmpl > /etc/squid/squid.conf

# Ensure helper scripts are executable
chmod +x /usr/local/lib/squid/helpers/*.py

# Start Squid in the foreground
exec squid -N -f /etc/squid/squid.conf

# ── Ports ────────────────────────────────────────────────────────────────────
http_port 3128

# ── Access log ───────────────────────────────────────────────────────────────
access_log stdio:/dev/stdout
cache_log  stdio:/dev/stderr
cache_store_log none
cache deny all

# ── Authentication helper ─────────────────────────────────────────────────────
# Reads lines of "<placeholder> <jwt>" from stdin; writes "OK user=<sub>" or "ERR"
auth_param basic program /usr/bin/python3 /usr/local/lib/squid/helpers/jwt_auth.py
auth_param basic children 5 startup=1 idle=1
auth_param basic realm "JWT Proxy"
auth_param basic credentialsttl 2 minutes

# ── External ACL helper ───────────────────────────────────────────────────────
# %LOGIN = authenticated username, %DST = destination host
external_acl_type jwt_policy \
    children=3 \
    %LOGIN %DST \
    /usr/bin/python3 /usr/local/lib/squid/helpers/policy_acl.py

# ── ACL definitions ──────────────────────────────────────────────────────────
acl authenticated  proxy_auth REQUIRED
acl policy_permit  external   jwt_policy
acl CONNECT        method     CONNECT
acl SSL_ports      port       443
acl Safe_ports     port       80 443 8080

# ── Access rules ─────────────────────────────────────────────────────────────
# Deny requests to unsafe ports
http_access deny !Safe_ports
http_access deny CONNECT !SSL_ports

# Require authentication
http_access deny !authenticated

# Require policy approval
http_access deny !policy_permit

# Allow everything that passed auth + policy
http_access allow authenticated policy_permit

# Default deny
http_access deny all

# ── Header rewriting ─────────────────────────────────────────────────────────
# Strip any Authorization header forwarded by the client
request_header_access Authorization deny all

# Inject the centrally managed service credential
# UPSTREAM_SVC_TOKEN is substituted at container start by envsubst
request_header_add Authorization "Bearer ${UPSTREAM_SVC_TOKEN}" all

# Squid automatically removes Proxy-Authorization before forwarding;
# the following makes that explicit and ensures it is always removed.
request_header_access Proxy-Authorization deny all

rules:
  # CI pipelines may reach GitHub API and Datadog
  - users:
      - "ci-pipeline"
      - "ci-staging"
    allow_destinations:
      - "api.github.com"
      - "api.datadoghq.com"

  # Analytics service may reach Snowflake and internal data APIs
  - users:
      - "svc-analytics"
    allow_destinations:
      - "account.snowflakecomputing.com"
      - "data-api.internal.example.com"

  # Any authenticated user may reach the echo server (POC only)
  - users:
      - "*"
    allow_destinations:
      - "echo-server"

FROM python:3.12-slim

WORKDIR /app
RUN pip install --no-cache-dir fastapi uvicorn[standard] PyJWT cryptography

COPY main.py .
COPY issuer_private_key.pem .
COPY issuer_public_key.pem .

EXPOSE 8080
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8080"]

import os
import base64
import datetime

import jwt
from fastapi import FastAPI, Query
from cryptography.hazmat.primitives.serialization import (
    load_pem_private_key, load_pem_public_key
)

app = FastAPI()

_priv = load_pem_private_key(
    open(os.environ["PRIVATE_KEY_PATH"], "rb").read(), password=None
)
_pub  = load_pem_public_key(
    open(os.environ["PUBLIC_KEY_PATH"], "rb").read()
)
_pub_numbers = _pub.public_numbers()

def _b64url(n: int, length: int) -> str:
    return base64.urlsafe_b64encode(
        n.to_bytes(length, "big")
    ).rstrip(b"=").decode()

@app.get("/.well-known/jwks.json")
def jwks():
    return {
        "keys": [{
            "kty": "RSA",
            "use": "sig",
            "alg": "RS256",
            "kid": "poc-v1",
            "n": _b64url(_pub_numbers.n, 256),
            "e": _b64url(_pub_numbers.e, 3),
        }]
    }

@app.post("/token")
def issue_token(
    sub: str = Query(..., description="Subject (service or user identity)"),
    ttl: int = Query(900, description="Token lifetime in seconds"),
):
    now = datetime.datetime.now(datetime.UTC)
    payload = {
        "iss": "http://token-issuer:8080",
        "sub": sub,
        "aud": "squid-proxy",
        "iat": now,
        "exp": now + datetime.timedelta(seconds=ttl),
    }
    token = jwt.encode(payload, _priv, algorithm="RS256", headers={"kid": "poc-v1"})
    return {"access_token": token, "expires_in": ttl, "sub": sub}

cd squid-jwt-proxy
docker compose up --build -d

# Issue a 15-minute token for the ci-pipeline identity
TOKEN=$(curl -s -X POST \
  "http://localhost:8080/token?sub=ci-pipeline&ttl=900" \
  | python3 -c "import sys,json; print(json.load(sys.stdin)['access_token'])")

echo "JWT: ${TOKEN:0:60}..."

# curl encodes the credentials automatically
curl -v \
  -x "http://placeholder:${TOKEN}@localhost:3128" \
  "http://echo-server/anything"

{
  "request": {
    "headers": {
      "Authorization": "Bearer svc-account-token-abc123",
      "Host": "echo-server",
      ...
    }
  }
}

# Attempt to reach a destination not in the ci-pipeline allow-list
curl -v \
  -x "http://placeholder:${TOKEN}@localhost:3128" \
  "http://account.snowflakecomputing.com/session/v1/login-request" \
  2>&1 | grep -E "HTTP|< "

# Craft an obviously invalid token
curl -v \
  -x "http://placeholder:not-a-jwt@localhost:3128" \
  "http://echo-server/anything" \
  2>&1 | grep "407"

rules:
  - users: ["ci-pipeline"]
    allow:
      - host: "api.github.com"
        methods: ["GET"]
      - host: "api.datadoghq.com"
        methods: ["GET", "POST"]

request_header_access X-Api-Key deny all
request_header_access X-Auth-Token deny all
request_header_access Cookie deny all