quorate/services/stalwart/cmd/controller/oidc.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

type oidcDiscoveryDocument struct {
	Issuer        string `json:"issuer"`
	JWKSURI       string `json:"jwks_uri"`
	TokenEndpoint string `json:"token_endpoint"`
}

func readSecret(path string) (string, error) {
	content, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	value := strings.TrimSpace(string(content))
	if value == "" {
		return "", errors.New("secret file is empty")
	}
	return value, nil
}

func fetchJSON(ctx context.Context, client *http.Client, target string, destination any) error {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
	if err != nil {
		return err
	}
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer func() { _ = response.Body.Close() }()
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return responseStatusError("OIDC endpoint", response)
	}
	return json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(destination)
}

func discoverOIDC(ctx context.Context, client *http.Client, issuer string) (oidcDiscoveryDocument, error) {
	issuerURL, err := url.Parse(issuer)
	if err != nil || issuerURL.Scheme == "" || issuerURL.Host == "" {
		return oidcDiscoveryDocument{}, errors.New("OIDC issuer is invalid")
	}
	target := strings.TrimRight(issuer, "/") + "/.well-known/openid-configuration"
	var document oidcDiscoveryDocument
	if err := fetchJSON(ctx, client, target, &document); err != nil {
		return oidcDiscoveryDocument{}, err
	}
	if strings.TrimRight(document.Issuer, "/") != strings.TrimRight(issuer, "/") {
		return oidcDiscoveryDocument{}, errors.New("OIDC discovery issuer does not match")
	}
	for _, endpoint := range []string{document.JWKSURI, document.TokenEndpoint} {
		endpointURL, err := url.Parse(endpoint)
		if err != nil || endpointURL.Scheme != issuerURL.Scheme || endpointURL.Host != issuerURL.Host {
			return oidcDiscoveryDocument{}, errors.New("OIDC endpoint origin does not match issuer")
		}
	}
	return document, nil
}

func oidcAccessToken(ctx context.Context, issuer, audience, clientID, clientSecret string) (string, error) {
	client := &http.Client{Timeout: 15 * time.Second}
	discovery, err := discoverOIDC(ctx, client, issuer)
	if err != nil {
		return "", err
	}
	form := url.Values{
		"grant_type": {"client_credentials"},
		"scope":      {"openid urn:zitadel:iam:org:project:id:" + audience + ":aud"},
	}
	request, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		discovery.TokenEndpoint,
		strings.NewReader(form.Encode()),
	)
	if err != nil {
		return "", err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.SetBasicAuth(clientID, clientSecret)
	response, err := client.Do(request)
	if err != nil {
		return "", err
	}
	defer func() { _ = response.Body.Close() }()
	if response.StatusCode != http.StatusOK {
		return "", responseStatusError("OIDC token endpoint", response, clientSecret)
	}
	var tokenResponse struct {
		AccessToken string `json:"access_token"`
	}
	if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&tokenResponse); err != nil {
		return "", err
	}
	if tokenResponse.AccessToken == "" {
		return "", errors.New("OIDC token response has no access token")
	}
	return tokenResponse.AccessToken, nil
}

func oidcTokenFromFiles(ctx context.Context, issuerFile, audienceFile, clientIDFile, clientSecretFile string) (string, error) {
	issuer, err := readSecret(issuerFile)
	if err != nil {
		return "", fmt.Errorf("read OIDC issuer: %w", err)
	}
	audience, err := readSecret(audienceFile)
	if err != nil {
		return "", fmt.Errorf("read OIDC audience: %w", err)
	}
	clientID, err := readSecret(clientIDFile)
	if err != nil {
		return "", fmt.Errorf("read OIDC client ID: %w", err)
	}
	clientSecret, err := readSecret(clientSecretFile)
	if err != nil {
		return "", fmt.Errorf("read OIDC client secret: %w", err)
	}
	return oidcAccessToken(ctx, issuer, audience, clientID, clientSecret)
}

func oidcEnvironment(ctx context.Context, value options) ([]string, error) {
	if value.oidcIssuerFile == "" {
		return nil, nil
	}
	token, err := oidcTokenFromFiles(
		ctx,
		value.oidcIssuerFile,
		value.oidcAudienceFile,
		value.oidcClientIDFile,
		value.oidcClientSecretFile,
	)
	if err != nil {
		return nil, err
	}
	return []string{"STALWART_TOKEN=" + token}, nil
}

func waitForOIDC(ctx context.Context, issuerFile string, interval time.Duration) error {
	issuer, err := readSecret(issuerFile)
	if err != nil {
		return err
	}
	client := &http.Client{Timeout: 15 * time.Second}
	for {
		discovery, err := discoverOIDC(ctx, client, issuer)
		if err == nil {
			var keys struct {
				Keys []json.RawMessage `json:"keys"`
			}
			if fetchJSON(ctx, client, discovery.JWKSURI, &keys) == nil && len(keys.Keys) != 0 {
				return nil
			}
		}
		timer := time.NewTimer(interval)
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}
	}
}