riotpiao-poimen/homelab
2
0
Fork 0
Code Issues 23 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

feat(paperless-ai): mint the LLM token with a machine key, not an init container

Browse Source 
The init container fetched a token once at pod start with a client_credentials grant
against a web app. A web app has no such grant, so the fetch could not work; and a
token fetched once is dead within the hour while the pod runs for weeks.

llm-token-vendor runs beside the app. It signs a machine-key assertion for the
paperless-llm service account, exchanges it for a token, and rewrites
/app/data/llm_token.txt atomically at 80% of the token lifetime. The vendor image
keeps using CUSTOM_API_KEY_FILE, so it needs no change and no OIDC support.

The key is re-read on every mint, so a rotation lands without a restart: publish the
new key, the next mint uses it, retire the old one.

paperless-llm is a separate service account from poimen-llm, so revoking one caller
leaves the other working.
This commit is contained in:
rock
2026-10-27 23:48:45 -07:00
parent bd81503e33
commit 4d23c983ab
6 changed files with 310 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
k8s/apps/llm-token-vendor/Dockerfile
+10
View File
@@ -0,0 +1,10 @@
FROM golang:1.26 AS build
WORKDIR /src
COPY go.mod ./
COPY *.go ./
RUN CGO_ENABLED=0 GOOS=linux go build -trimpath -ldflags="-s -w" -o /out/vendor .
FROM gcr.io/distroless/static:nonroot
COPY --from=build /out/vendor /usr/local/bin/vendor
USER 65532:65532
ENTRYPOINT ["/usr/local/bin/vendor"]
k8s/apps/llm-token-vendor/go.mod
+3
View File
@@ -0,0 +1,3 @@
module github.com/riotpiao/llm-token-vendor
go 1.26.0
k8s/apps/llm-token-vendor/llm-token-vendor
Executable
BIN
View File
Binary file not shown.
k8s/apps/llm-token-vendor/main.go
+197
View File
@@ -0,0 +1,197 @@
// llm-token-vendor mints ZITADEL access tokens for a caller that cannot sign.
//
// Some callers are vendor images: paperless-ai can only read a bearer token from a
// file. It cannot sign an assertion, so it cannot use a machine key, and a client
// secret would rotate with a window.
//
// This sidecar closes the gap. It signs with a machine key, exchanges the assertion
// for an access token, and writes the token to the file the vendor image already
// reads. The file is replaced atomically and refreshed at 80% of the token lifetime,
// so the vendor never reads a partial or expired value.
//
// The key is re-read on every mint. A rotation therefore lands without a restart:
// publish the new key, the next mint uses it, retire the old one.
package main
import (
"crypto"
"crypto/rand"
"crypto/rsa"
"crypto/sha256"
"crypto/x509"
"encoding/base64"
"encoding/json"
"encoding/pem"
"flag"
"fmt"
"io"
"log"
"net/http"
"net/url"
"os"
"path/filepath"
"time"
)
type config struct {
tokenURL string
issuer string
userID string
keyID string
keyPath string
scope string
outPath string
}
func main() {
var c config
flag.StringVar(&c.tokenURL, "token-url", env("TOKEN_URL", ""), "ZITADEL token endpoint")
flag.StringVar(&c.issuer, "issuer", env("ISSUER", ""), "ZITADEL issuer, the assertion audience")
flag.StringVar(&c.userID, "user-id", env("USER_ID", ""), "machine user id")
flag.StringVar(&c.keyID, "key-id", env("KEY_ID", ""), "machine key id")
flag.StringVar(&c.keyPath, "key-file", env("KEY_FILE", ""), "PEM private key, re-read on every mint")
flag.StringVar(&c.scope, "scope", env("SCOPE", "openid"), "requested scopes")
flag.StringVar(&c.outPath, "out", env("OUT", ""), "file the vendor image reads")
flag.Parse()
if c.tokenURL == "" || c.issuer == "" || c.userID == "" || c.keyPath == "" || c.outPath == "" {
log.Fatal("token-url, issuer, user-id, key-file and out are required")
}
for {
lifetime, err := mintAndWrite(c)
if err != nil {
// Exit rather than spin: a pod that cannot mint must be visible, and the
// vendor image must not keep serving the caller a stale token.
log.Fatalf("cannot mint a token: %v", err)
}
wait := time.Duration(float64(lifetime) * 0.8)
log.Printf("token written to %s, next mint in %s", c.outPath, wait.Round(time.Second))
time.Sleep(wait)
}
}
func mintAndWrite(c config) (time.Duration, error) {
pemBytes, err := os.ReadFile(c.keyPath)
if err != nil {
return 0, fmt.Errorf("cannot read the key %s: %w", c.keyPath, err)
}
assertion, err := assertion(c, time.Now(), pemBytes)
if err != nil {
return 0, err
}
form := url.Values{}
form.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
form.Set("assertion", assertion)
form.Set("scope", c.scope)
resp, err := http.PostForm(c.tokenURL, form)
if err != nil {
return 0, fmt.Errorf("token request failed: %w", err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
if resp.StatusCode != http.StatusOK {
return 0, fmt.Errorf("token endpoint returned %d: %s", resp.StatusCode, string(body))
}
var out struct {
AccessToken string `json:"access_token"`
ExpiresIn int `json:"expires_in"`
}
if err := json.Unmarshal(body, &out); err != nil {
return 0, fmt.Errorf("cannot decode the token response: %w", err)
}
if out.AccessToken == "" {
return 0, fmt.Errorf("the token response holds no access_token")
}
// Write beside the target, then rename. A rename is atomic, so the vendor either
// reads the previous token or the new one, never a half-written file.
tmp, err := os.CreateTemp(filepath.Dir(c.outPath), ".token-*")
if err != nil {
return 0, fmt.Errorf("cannot create a temporary file: %w", err)
}
defer os.Remove(tmp.Name())
if _, err := tmp.WriteString(out.AccessToken); err != nil {
tmp.Close()
return 0, fmt.Errorf("cannot write the token: %w", err)
}
tmp.Close()
if err := os.Chmod(tmp.Name(), 0o644); err != nil {
return 0, fmt.Errorf("cannot set the file mode: %w", err)
}
if err := os.Rename(tmp.Name(), c.outPath); err != nil {
return 0, fmt.Errorf("cannot replace %s: %w", c.outPath, err)
}
lifetime := time.Duration(out.ExpiresIn) * time.Second
if lifetime <= 0 {
lifetime = 5 * time.Minute
}
return lifetime, nil
}
func assertion(c config, now time.Time, pemBytes []byte) (string, error) {
block, _ := pem.Decode(pemBytes)
if block == nil {
return "", fmt.Errorf("the key file is not PEM")
}
var key *rsa.PrivateKey
if k, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
key = k
} else {
parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
if err != nil {
return "", fmt.Errorf("the key is not an RSA private key: %w", err)
}
rk, ok := parsed.(*rsa.PrivateKey)
if !ok {
return "", fmt.Errorf("the key holds %T, not RSA", parsed)
}
key = rk
}
enc := func(v any) (string, error) {
raw, err := json.Marshal(v)
if err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(raw), nil
}
header := map[string]any{"alg": "RS256", "typ": "JWT"}
if c.keyID != "" {
header["kid"] = c.keyID
}
claims := map[string]any{
"iss": c.userID,
"sub": c.userID,
"aud": c.issuer,
// Backdated, so a clock that is a few seconds fast cannot make the assertion
// "not yet valid".
"iat": now.Add(-5 * time.Second).Unix(),
"exp": now.Add(5 * time.Minute).Unix(),
}
h, err := enc(header)
if err != nil {
return "", err
}
cl, err := enc(claims)
if err != nil {
return "", err
}
signing := h + "." + cl
digest := sha256.Sum256([]byte(signing))
sig, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
if err != nil {
return "", fmt.Errorf("cannot sign the assertion: %w", err)
}
return signing + "." + base64.RawURLEncoding.EncodeToString(sig), nil
}
func env(key, fallback string) string {
if v := os.Getenv(key); v != "" {
return v
}
return fallback
}
k8s/apps/llm-token-vendor/token_vendor_test.go
+64
View File
@@ -0,0 +1,64 @@
package main
import (
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/base64"
"encoding/json"
"encoding/pem"
"strings"
"testing"
"time"
)
func key(t *testing.T) []byte {
t.Helper()
k, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("cannot generate a key: %v", err)
}
return pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k)})
}
func seg(t *testing.T, s string, out any) {
t.Helper()
raw, err := base64.RawURLEncoding.DecodeString(s)
if err != nil {
t.Fatalf("not base64url: %v", err)
}
if err := json.Unmarshal(raw, out); err != nil {
t.Fatalf("not JSON: %v", err)
}
}
func TestAssertionCarriesTheKeyAndTheAudience(t *testing.T) {
c := config{issuer: "https://zitadel.example", userID: "3970", keyID: "3971"}
tok, err := assertion(c, time.Unix(1_700_000_000, 0), key(t))
if err != nil {
t.Fatalf("signing failed: %v", err)
}
parts := strings.Split(tok, ".")
if len(parts) != 3 {
t.Fatalf("a JWT has three segments, got %d", len(parts))
}
var h map[string]any
seg(t, parts[0], &h)
if h["kid"] != "3971" || h["alg"] != "RS256" {
t.Fatalf("header must name the key and the algorithm, got %v", h)
}
var cl map[string]any
seg(t, parts[1], &cl)
if cl["iss"] != "3970" || cl["sub"] != "3970" || cl["aud"] != "https://zitadel.example" {
t.Fatalf("claims are wrong: %v", cl)
}
if int64(cl["iat"].(float64)) != 1_699_999_995 {
t.Fatalf("iat must be backdated five seconds, got %v", cl["iat"])
}
}
func TestAssertionRejectsAKeyThatIsNotPEM(t *testing.T) {
if _, err := assertion(config{userID: "u", issuer: "i"}, time.Now(), []byte("nope")); err == nil {
t.Fatalf("a non-PEM key must be rejected")
}
}
k8s/apps/paperless/paperless-ai.yaml
+36 -46
View File
@@ -36,53 +36,43 @@ spec:
- name: paperless-ai-env-config
configMap:
name: paperless-ai-env
initContainers:
- name: fetch-llm-token
image: curlimages/curl:8.12.0
securityContext:
runAsUser: 0
fsGroup: 0
command:
- sh
- -c
- |
set -e
echo "[init] Fetching JWT token from Authentik for paperless-ai-agent..."
# Get JWT token via client_credentials grant
# Client secret sourced from environment (injected from paperless-ai-config Secret)
TOKEN_RESPONSE=$(curl -s -X POST https://zitadel.riotpiao.com/oauth/v2/token \
-d "grant_type=client_credentials" \
-d "client_id=396757803559027071" \
-d "client_secret=${LLM_AUTH_CLIENT_SECRET}" \
-d "scope=openid llm:inference" 2>/dev/null)
# Extract token from JSON response using sed (handles spaces after colons)
TOKEN=$(echo "$TOKEN_RESPONSE" | sed -n 's/.*"access_token"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p')
if [ -z "$TOKEN" ]; then
echo "[error] Failed to get token. Response: $TOKEN_RESPONSE"
exit 1
fi
# Store token in file for main container to read
# Write to both locations for compatibility
mkdir -p /tmp/llm-token /app/data 2>/dev/null || true
echo "$TOKEN" | tee /tmp/llm-token/llm_token.txt > /dev/null 2>&1
echo "$TOKEN" > /app/data/llm_token.txt 2>/dev/null || true
echo "[init] Token fetched and stored"
[ -f /tmp/llm-token/llm_token.txt ] && echo " -> /tmp/llm-token/llm_token.txt"
[ -f /app/data/llm_token.txt ] && echo " -> /app/data/llm_token.txt"
env:
- name: LLM_AUTH_CLIENT_SECRET
valueFrom:
secretKeyRef:
name: paperless-ai-config
key: LLM_AUTH_CLIENT_SECRET
volumeMounts:
- name: paperless-ai-data
mountPath: /app/data
- name: llm-key
secret:
secretName: paperless-llm-key
defaultMode: 0400
# No init container. The token used to be fetched once at pod start with a
# client_credentials grant against a web app, which cannot work: a web app has
# no such grant, and even a token that was minted would be dead within the hour
# while this pod runs for weeks.
#
# llm-token-vendor runs beside the app instead. It signs a machine-key assertion,
# exchanges it for a token, and rewrites /app/data/llm_token.txt at 80% of the
# token lifetime. The vendor image keeps reading CUSTOM_API_KEY_FILE unchanged.
# The key is re-read on every mint, so a rotation lands without a restart.
containers:
- name: llm-token-vendor
image: forgejo.riotpiao.com/riotpiao-poimen/llm-token-vendor:0.1.0
imagePullPolicy: Always
args:
- --token-url=https://zitadel.riotpiao.com/oauth/v2/token
- --issuer=https://zitadel.riotpiao.com
- --user-id=397072668735570071
- --key-id=397072674272051351
- --key-file=/var/run/llm-key/key
- --scope=openid urn:zitadel:iam:org:project:id:396757796898471937:aud
- --out=/app/data/llm_token.txt
volumeMounts:
- {name: paperless-ai-data, mountPath: /app/data}
- {name: llm-key, mountPath: /var/run/llm-key, readOnly: true}
securityContext:
allowPrivilegeEscalation: false
capabilities: {drop: ["ALL"]}
runAsNonRoot: true
runAsUser: 65532
seccompProfile: {type: RuntimeDefault}
resources:
requests: {cpu: 10m, memory: 32Mi}
limits: {cpu: 100m, memory: 64Mi}
- name: paperless-ai
image: clusterzx/paperless-ai:latest
env: