zitadel-k8s-operator/src/pkg/zitadel/zitadel.go

package zitadel

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"time"

	zitadelv1alpha1 "bitbucket.org/topmanage-software-engineering/zitadel-k8s-operator/src/api/v1alpha1"
	"bitbucket.org/topmanage-software-engineering/zitadel-k8s-operator/src/pkg/deployment"
	"bitbucket.org/topmanage-software-engineering/zitadel-k8s-operator/src/pkg/machinekey"
	"github.com/gorilla/schema"
	"google.golang.org/grpc"

	"github.com/zitadel/oidc/pkg/client"
	httphelper "github.com/zitadel/oidc/pkg/http"
	"github.com/zitadel/oidc/pkg/oidc"
	"github.com/zitadel/zitadel-go/v3/pkg/client/admin"
	"github.com/zitadel/zitadel-go/v3/pkg/client/management"
	"github.com/zitadel/zitadel-go/v3/pkg/client/zitadel"
	"golang.org/x/oauth2"
	"gopkg.in/square/go-jose.v2"
	corev1 "k8s.io/api/core/v1"
)

type MachineKey struct {
	Type   string `json:"type"`
	KeyID  string `json:"keyId"`
	Key    string `json:"key"`
	UserID string `json:"userId"`
}

func NewClient(ctx context.Context, zitadelCluster *zitadelv1alpha1.ZitadelCluster, refresolver zitadelv1alpha1.RefResolver) (*management.Client, error) {
	machineKeyData, err := refresolver.SecretKeyRef(ctx, corev1.SecretKeySelector{LocalObjectReference: corev1.LocalObjectReference{Name: machinekey.MachineKeySecretName(zitadelCluster)}, Key: machinekey.Key}, zitadelCluster.Namespace)
	if err != nil {
		return nil, err
	}
	api, err := management.NewClient(ctx, GetIssuer(zitadelCluster), GetAPI(zitadelCluster), []string{oidc.ScopeOpenID, zitadel.ScopeZitadelAPI()}, zitadel.WithInsecure(), zitadel.WithJWTProfileTokenSource(Discover([]byte(machineKeyData), GetAPIUrl(zitadelCluster), GetAuthority(zitadelCluster), GetAPI(zitadelCluster))),
		zitadel.WithDialOptions(grpc.WithAuthority(GetAuthority(zitadelCluster))),
	)

	if err != nil {
		return nil, fmt.Errorf("ERROR CREATING CLIENT: %v", err)
	}
	return api, nil
}

func NewAdminClient(ctx context.Context, zitadelCluster *zitadelv1alpha1.ZitadelCluster, refresolver zitadelv1alpha1.RefResolver) (*admin.Client, error) {
	machineKeyData, err := refresolver.SecretKeyRef(ctx, corev1.SecretKeySelector{LocalObjectReference: corev1.LocalObjectReference{Name: machinekey.MachineKeySecretName(zitadelCluster)}, Key: machinekey.Key}, zitadelCluster.Namespace)
	if err != nil {
		return nil, err
	}
	api, err := admin.NewClient(ctx, GetIssuer(zitadelCluster), GetAPI(zitadelCluster), []string{oidc.ScopeOpenID, zitadel.ScopeZitadelAPI()}, zitadel.WithInsecure(), zitadel.WithJWTProfileTokenSource(Discover([]byte(machineKeyData), GetAPIUrl(zitadelCluster), GetAuthority(zitadelCluster), GetAPI(zitadelCluster))),
		zitadel.WithDialOptions(grpc.WithAuthority(GetAuthority(zitadelCluster))),
	)
	if err != nil {
		return nil, fmt.Errorf("ERROR CREATING CLIENT: %v", err)
	}
	return api, nil
}

func GetAuthority(zitadel *zitadelv1alpha1.ZitadelCluster) string {
	return fmt.Sprintf("%s:%d", zitadel.Spec.Host, zitadel.Spec.ExternalPort)
}

func GetIssuer(zitadel *zitadelv1alpha1.ZitadelCluster) string {
	scheme := "http"
	if zitadel.Spec.ExternalSecure {
		scheme = "https"
	}
	return fmt.Sprintf("%s://%s:%d", scheme, zitadel.Spec.Host, zitadel.Spec.ExternalPort)
}

func GetAPI(zitadel *zitadelv1alpha1.ZitadelCluster) string {
	return fmt.Sprintf("%s:%d", deployment.ServiceFQDN(zitadel.ObjectMeta), deployment.ZitadelPort)
}

func GetAPIUrl(zitadel *zitadelv1alpha1.ZitadelCluster) string {
	return fmt.Sprintf("http://%s:%d", deployment.ServiceFQDN(zitadel.ObjectMeta), deployment.ZitadelPort)
}

type jwtProfileTokenSource struct {
	clientID      string
	audience      []string
	signer        jose.Signer
	scopes        []string
	httpClient    *http.Client
	tokenEndpoint string
	host          string
}

func Discover(key []byte, discoverUrl string, host string, api string) func(issuer string, scopes []string) (oauth2.TokenSource, error) {
	return func(issuer string, scopes []string) (oauth2.TokenSource, error) {
		var machineKeyData MachineKey
		if err := json.Unmarshal(key, &machineKeyData); err != nil {
			return nil, err
		}
		signer, err := client.NewSignerFromPrivateKeyByte([]byte(machineKeyData.Key), machineKeyData.KeyID)
		if err != nil {
			return nil, err
		}
		source := &jwtProfileTokenSource{
			host:       host,
			clientID:   machineKeyData.UserID,
			audience:   []string{issuer},
			signer:     signer,
			scopes:     scopes,
			httpClient: http.DefaultClient,
		}
		config, err := GetDiscoveryConfig(discoverUrl, http.DefaultClient, host, api)
		if err != nil {
			return nil, err
		}
		source.tokenEndpoint = config.TokenEndpoint
		return source, nil
	}
}

func GetDiscoveryConfig(issuer string, httpClient *http.Client, host string, api string, wellKnownUrl ...string) (*oidc.DiscoveryConfiguration, error) {
	wellKnown := strings.TrimSuffix(issuer, "/") + oidc.DiscoveryEndpoint
	if len(wellKnownUrl) == 1 && wellKnownUrl[0] != "" {
		wellKnown = wellKnownUrl[0]
	}
	req, err := http.NewRequest("GET", wellKnown, nil)
	if err != nil {
		return nil, err
	}
	req.Host = host
	discoveryConfig := new(oidc.DiscoveryConfiguration)
	err = httphelper.HttpRequest(httpClient, req, &discoveryConfig)
	discoveryConfig.TokenEndpoint = replaceEndpoint(discoveryConfig.TokenEndpoint, host, api)
	discoveryConfig.AuthorizationEndpoint = replaceEndpoint(discoveryConfig.AuthorizationEndpoint, host, api)
	discoveryConfig.IntrospectionEndpoint = replaceEndpoint(discoveryConfig.IntrospectionEndpoint, host, api)
	discoveryConfig.EndSessionEndpoint = replaceEndpoint(discoveryConfig.EndSessionEndpoint, host, api)
	discoveryConfig.RevocationEndpoint = replaceEndpoint(discoveryConfig.RevocationEndpoint, host, api)
	discoveryConfig.UserinfoEndpoint = replaceEndpoint(discoveryConfig.UserinfoEndpoint, host, api)

	if err != nil {
		return nil, err
	}
	return discoveryConfig, nil
}

func replaceEndpoint(endpoint string, host string, api string) string {
	return strings.ReplaceAll(strings.ReplaceAll(endpoint, host, api), "https", "http")
}

func (j *jwtProfileTokenSource) TokenEndpoint() string {
	return j.tokenEndpoint
}

func (j *jwtProfileTokenSource) HttpClient() *http.Client {
	return j.httpClient
}

func (j *jwtProfileTokenSource) Token() (*oauth2.Token, error) {
	assertion, err := client.SignedJWTProfileAssertion(j.clientID, j.audience, time.Hour, j.signer)
	if err != nil {
		return nil, err
	}
	token, err := callTokenEndpoint(oidc.NewJWTProfileGrantRequest(assertion, j.scopes...), nil, j, j.host)
	if err != nil {
		return nil, err
	}
	return token, err
}

var Encoder = func() httphelper.Encoder {
	e := schema.NewEncoder()
	e.RegisterEncoder(oidc.SpaceDelimitedArray{}, func(value reflect.Value) string {
		return value.Interface().(oidc.SpaceDelimitedArray).Encode()
	})
	return e
}()

func callTokenEndpoint(request interface{}, authFn interface{}, caller client.TokenEndpointCaller, host string) (newToken *oauth2.Token, err error) {
	req, err := httphelper.FormRequest(caller.TokenEndpoint(), request, Encoder, authFn)
	if err != nil {
		return nil, err
	}
	tokenRes := new(oidc.AccessTokenResponse)
	req.Host = host
	if err := httphelper.HttpRequest(caller.HttpClient(), req, &tokenRes); err != nil {
		return nil, fmt.Errorf("Error calling token endpoint: %v", err)
	}
	return &oauth2.Token{
		AccessToken:  tokenRes.AccessToken,
		TokenType:    tokenRes.TokenType,
		RefreshToken: tokenRes.RefreshToken,
		Expiry:       time.Now().UTC().Add(time.Duration(tokenRes.ExpiresIn) * time.Second),
	}, nil
}