feat(paperless-ai): mint the LLM token with a machine key, not an init container
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:
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user