homelab/k8s/apps/key-rotation-controller/controller/internal/zitadel/client.go

package zitadel

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// Config holds everything the controller needs to call ZITADEL.
type Config struct {
	// Issuer is the instance URL, for example https://zitadel.riotpiao.com.
	Issuer string
	// ClientID is the API app that this controller authenticates as.
	ClientID string
	// UserID is the machine user the assertion acts for. Empty means ClientID.
	UserID string
	// KeyID and PrivateKeyPEM identify the controller's own machine key.
	KeyID         string
	PrivateKeyPEM []byte
	// HTTPClient is optional; a sane default is used when nil.
	HTTPClient *http.Client
}

// Client is a thin ZITADEL management client.
type Client struct {
	cfg   Config
	http  *http.Client
	mu    sync.Mutex
	token string
	until time.Time
}

// New builds a client.
func New(cfg Config) *Client {
	client := cfg.HTTPClient
	if client == nil {
		client = &http.Client{Timeout: 20 * time.Second}
	}
	return &Client{cfg: cfg, http: client}
}

// AccessToken returns a cached access token, minting one when needed.
func (c *Client) AccessToken() (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.token != "" && time.Now().Before(c.until) {
		return c.token, nil
	}

	assertion, err := SignAssertion(c.cfg.PrivateKeyPEM, c.cfg.KeyID, c.cfg.ClientID, c.cfg.UserID, c.cfg.Issuer, time.Now())
	if err != nil {
		return "", err
	}

	form := url.Values{}
	form.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
	form.Set("assertion", assertion)
	// The management API rejects a token without the ZITADEL project audience:
	// the request returns 401 Token.Invalid even though the assertion was accepted.
	form.Set("scope", "openid urn:zitadel:iam:org:project:id:zitadel:aud")

	req, err := http.NewRequest(http.MethodPost, strings.TrimRight(c.cfg.Issuer, "/")+"/oauth/v2/token", strings.NewReader(form.Encode()))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := c.http.Do(req)
	if err != nil {
		return "", fmt.Errorf("token request failed: %w", err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16))
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("token endpoint returned %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}

	var parsed struct {
		AccessToken string `json:"access_token"`
		ExpiresIn   int    `json:"expires_in"`
	}
	if err := json.Unmarshal(body, &parsed); err != nil {
		return "", fmt.Errorf("token response did not parse: %w", err)
	}
	if parsed.AccessToken == "" {
		return "", fmt.Errorf("token response carried no access_token")
	}

	lifetime := parsed.ExpiresIn
	if lifetime <= 0 {
		lifetime = 300
	}
	c.token = parsed.AccessToken
	// Refresh at 80 % of the lifetime so a slow call never presents an expired token.
	c.until = time.Now().Add(time.Duration(lifetime) * time.Second * 4 / 5)
	return c.token, nil
}

// projectAppPath builds the management path for one OIDC app under a project.
func (c *Client) projectAppPath(projectID, appID string) string {
	return fmt.Sprintf("/management/v1/projects/%s/apps/%s", projectID, appID)
}

// MachineKey is one registered key for a service user.
type MachineKey struct {
	KeyID      string `json:"keyId"`
	Expiration string `json:"expirationDate,omitempty"`
	// Details holds the PEM bundle when the key was just created.
	Details string `json:"keyDetails,omitempty"`
}

// NewMachineKeyJSON is the JSON document ZITADEL returns for a new JSON key.
type NewMachineKeyJSON struct {
	Type     string `json:"type"`
	KeyID    string `json:"keyId"`
	Key      string `json:"key"`
	UserID   string `json:"userId"`
	ExpireAt string `json:"expirationDate,omitempty"`
}

// CreateMachineKey adds a key to a service user and returns its material.
//
// ZITADEL keeps many keys per user, so this never invalidates the key that is
// currently deployed. That is the whole reason the machine-key provider can
// rotate with no gap.
func (c *Client) CreateMachineKey(userID string) (*NewMachineKeyJSON, error) {
	payload := map[string]any{"type": "KEY_TYPE_JSON"}
	var out struct {
		KeyID      string `json:"keyId"`
		KeyDetails string `json:"keyDetails"`
	}
	if err := c.post("/management/v1/users/"+userID+"/keys", payload, &out); err != nil {
		return nil, err
	}
	if out.KeyDetails == "" {
		return nil, fmt.Errorf("create machine key returned no key material")
	}

	raw, err := base64Decode(out.KeyDetails)
	if err != nil {
		return nil, fmt.Errorf("key material is not base64: %w", err)
	}

	var key NewMachineKeyJSON
	if err := json.Unmarshal(raw, &key); err != nil {
		return nil, fmt.Errorf("key material is not the expected JSON: %w", err)
	}
	return &key, nil
}

// DeleteMachineKey removes one key. Called only after the new key is proven.
func (c *Client) DeleteMachineKey(userID, keyID string) error {
	return c.delete("/management/v1/users/" + userID + "/keys/" + keyID)
}

// GenerateClientSecret replaces the client secret of an OIDC app.
//
// ZITADEL stores exactly one secret per app, so the previous secret stops working
// at this call. The caller must publish the new value immediately.
func (c *Client) GenerateClientSecret(projectID, appID string) (string, error) {
	var out struct {
		ClientSecret string `json:"clientSecret"`
	}
	path := c.projectAppPath(projectID, appID) + "/oidc_config/_generate_client_secret"
	if err := c.post(path, map[string]any{}, &out); err != nil {
		return "", err
	}
	if out.ClientSecret == "" {
		return "", fmt.Errorf("generate client secret returned an empty value")
	}
	return out.ClientSecret, nil
}

// Introspect checks a client credential against the token endpoint.
//
// A successful call returns an object with `active`. The credentials are valid
// when the endpoint answers 200, whatever `active` says, because `active`
// describes the token rather than the client.
func (c *Client) Introspect(clientID, clientSecret string) error {
	form := url.Values{}
	form.Set("token", "probe")

	req, err := http.NewRequest(http.MethodPost, strings.TrimRight(c.cfg.Issuer, "/")+"/oauth/v2/introspect", strings.NewReader(form.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth(clientID, clientSecret)

	resp, err := c.http.Do(req)
	if err != nil {
		return fmt.Errorf("introspection request failed: %w", err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16))

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("introspection returned %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	return nil
}

// ExchangeMachineKey obtains a token with a specific key, proving that key works.
func (c *Client) ExchangeMachineKey(clientID, userID, keyID string, keyPEM []byte, audienceProject string) error {
	assertion, err := SignAssertion(keyPEM, keyID, clientID, userID, c.cfg.Issuer, time.Now())
	if err != nil {
		return err
	}

	scope := "openid"
	if audienceProject != "" {
		scope += " urn:zitadel:iam:org:project:id:" + audienceProject + ":aud"
	}

	form := url.Values{}
	form.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
	form.Set("assertion", assertion)
	form.Set("scope", scope)

	req, err := http.NewRequest(http.MethodPost, strings.TrimRight(c.cfg.Issuer, "/")+"/oauth/v2/token", strings.NewReader(form.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := c.http.Do(req)
	if err != nil {
		return fmt.Errorf("key exchange failed: %w", err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16))
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("key exchange returned %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	return nil
}

func (c *Client) do(method, path string, payload any, out any) error {
	token, err := c.AccessToken()
	if err != nil {
		return err
	}

	var body io.Reader
	if payload != nil {
		encoded, err := json.Marshal(payload)
		if err != nil {
			return err
		}
		body = bytes.NewReader(encoded)
	}

	req, err := http.NewRequest(method, strings.TrimRight(c.cfg.Issuer, "/")+path, body)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+token)
	if payload != nil {
		req.Header.Set("Content-Type", "application/json")
	}

	resp, err := c.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("%s %s returned %d: %s", method, path, resp.StatusCode, strings.TrimSpace(string(raw)))
	}
	if out != nil && len(raw) > 0 {
		if err := json.Unmarshal(raw, out); err != nil {
			return fmt.Errorf("%s %s response did not parse: %w", method, path, err)
		}
	}
	return nil
}

func (c *Client) post(path string, payload, out any) error {
	return c.do(http.MethodPost, path, payload, out)
}

func (c *Client) delete(path string) error {
	return c.do(http.MethodDelete, path, nil, nil)
}

func base64Decode(s string) ([]byte, error) {
	if raw, err := base64StdDecode(s); err == nil {
		return raw, nil
	}
	return base64URLDecode(s)
}

func base64StdDecode(s string) ([]byte, error) {
	return base64.StdEncoding.DecodeString(s)
}

func base64URLDecode(s string) ([]byte, error) {
	return base64.RawURLEncoding.DecodeString(s)
}