ITWorx-Pulse-Public/internal/auth/oidc.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/coreos/go-oidc/v3/oidc"
	"golang.org/x/oauth2"
)

const (
	defaultFlowLifetime = 10 * time.Minute
)

type OIDCConfig struct {
	Issuer       string
	ClientID     string
	ClientSecret string
	RedirectURL  string
	Scopes       []string
}

type Authorization struct {
	URL          string
	State        string
	Nonce        string
	CodeVerifier string
	ExpiresAt    time.Time
}

func BeginAuthorization(endpoint oauth2.Endpoint, config OIDCConfig, now time.Time) (Authorization, error) {
	if endpoint.AuthURL == "" || config.ClientID == "" || config.RedirectURL == "" {
		return Authorization{}, errors.New("OIDC authorization configuration is incomplete")
	}
	state, err := randomToken()
	if err != nil {
		return Authorization{}, errors.New("generate authorization state")
	}
	nonce, err := randomToken()
	if err != nil {
		return Authorization{}, errors.New("generate authorization nonce")
	}
	verifier, err := randomToken()
	if err != nil {
		return Authorization{}, errors.New("generate PKCE verifier")
	}
	scopes := config.Scopes
	if len(scopes) == 0 {
		scopes = []string{oidc.ScopeOpenID, "profile", "email"}
	}
	oauthConfig := oauth2.Config{
		ClientID:     config.ClientID,
		ClientSecret: config.ClientSecret,
		Endpoint:     endpoint,
		RedirectURL:  config.RedirectURL,
		Scopes:       scopes,
	}
	authURL := oauthConfig.AuthCodeURL(state,
		oauth2.SetAuthURLParam("nonce", nonce),
		oauth2.SetAuthURLParam("code_challenge", pkceChallenge(verifier)),
		oauth2.SetAuthURLParam("code_challenge_method", "S256"),
	)
	return Authorization{URL: authURL, State: state, Nonce: nonce, CodeVerifier: verifier, ExpiresAt: now.Add(defaultFlowLifetime)}, nil
}

func ValidateCallback(flow Authorization, state, code string, now time.Time) error {
	if flow.State == "" || subtle.ConstantTimeCompare([]byte(flow.State), []byte(state)) != 1 {
		return errors.New("OIDC state validation failed")
	}
	if flow.CodeVerifier == "" || flow.Nonce == "" {
		return errors.New("OIDC flow is incomplete")
	}
	if now.After(flow.ExpiresAt) {
		return errors.New("OIDC authorization expired")
	}
	if strings.TrimSpace(code) == "" {
		return errors.New("OIDC authorization code is required")
	}
	return nil
}

func Exchange(ctx context.Context, flow Authorization, config OIDCConfig, endpoint oauth2.Endpoint, state, code string) (*oauth2.Token, error) {
	if err := ValidateCallback(flow, state, code, time.Now()); err != nil {
		return nil, err
	}
	oauthConfig := oauth2.Config{ClientID: config.ClientID, ClientSecret: config.ClientSecret, Endpoint: endpoint, RedirectURL: config.RedirectURL}
	return oauthConfig.Exchange(ctx, code, oauth2.SetAuthURLParam("code_verifier", flow.CodeVerifier))
}

// Discovery is the provider metadata required to run one authorization code flow:
// the authorization/token endpoints for BeginAuthorization and Exchange, and the
// ID token verifier for VerifyIDToken. Resolve it once and reuse it.
type Discovery struct {
	Endpoint oauth2.Endpoint
	Verifier *oidc.IDTokenVerifier
}

func Discover(ctx context.Context, config OIDCConfig) (Discovery, error) {
	if config.Issuer == "" || config.ClientID == "" {
		return Discovery{}, errors.New("OIDC issuer and client ID are required")
	}
	provider, err := oidc.NewProvider(ctx, config.Issuer)
	if err != nil {
		return Discovery{}, fmt.Errorf("OIDC discovery failed")
	}
	return Discovery{Endpoint: provider.Endpoint(), Verifier: provider.Verifier(&oidc.Config{ClientID: config.ClientID})}, nil
}

func NewVerifier(ctx context.Context, config OIDCConfig) (*oidc.IDTokenVerifier, error) {
	discovery, err := Discover(ctx, config)
	if err != nil {
		return nil, err
	}
	return discovery.Verifier, nil
}

func VerifyIDToken(ctx context.Context, verifier *oidc.IDTokenVerifier, rawToken, expectedNonce string) (*oidc.IDToken, error) {
	if verifier == nil || strings.TrimSpace(rawToken) == "" || expectedNonce == "" {
		return nil, errors.New("OIDC token verification input is incomplete")
	}
	token, err := verifier.Verify(ctx, rawToken)
	if err != nil {
		return nil, errors.New("OIDC token verification failed")
	}
	var claims struct {
		Nonce string `json:"nonce"`
	}
	if err := token.Claims(&claims); err != nil || subtle.ConstantTimeCompare([]byte(claims.Nonce), []byte(expectedNonce)) != 1 {
		return nil, errors.New("OIDC nonce validation failed")
	}
	return token, nil
}

const (
	defaultGroupsClaim = "groups"
	maxIdentityGroups  = 128
)

// Identity is the bounded subset of verified ID token claims Pulse consumes.
type Identity struct {
	Subject string
	Groups  []string
}

// ExtractIdentity reads the subject and the configured role claim from an already
// verified ID token. The claim may be a list of strings or a single string; values
// are trimmed, empty values dropped and the list bounded.
func ExtractIdentity(token *oidc.IDToken, groupsClaim string) (Identity, error) {
	if token == nil {
		return Identity{}, errors.New("OIDC identity token is required")
	}
	if groupsClaim == "" {
		groupsClaim = defaultGroupsClaim
	}
	subject := strings.TrimSpace(token.Subject)
	if subject == "" {
		return Identity{}, errors.New("OIDC subject claim is required")
	}
	var claims map[string]json.RawMessage
	if err := token.Claims(&claims); err != nil {
		return Identity{}, errors.New("OIDC claims could not be read")
	}
	raw, ok := claims[groupsClaim]
	if !ok {
		return Identity{Subject: subject}, nil
	}
	groups, err := normalizeGroupClaim(raw)
	if err != nil {
		return Identity{}, err
	}
	return Identity{Subject: subject, Groups: groups}, nil
}

func normalizeGroupClaim(raw json.RawMessage) ([]string, error) {
	var values []string
	if err := json.Unmarshal(raw, &values); err != nil {
		var single string
		if err := json.Unmarshal(raw, &single); err != nil {
			return nil, errors.New("OIDC role claim is malformed")
		}
		values = []string{single}
	}
	groups := make([]string, 0, len(values))
	for _, value := range values {
		trimmed := strings.TrimSpace(value)
		if trimmed == "" || len(groups) >= maxIdentityGroups {
			continue
		}
		groups = append(groups, trimmed)
	}
	return groups, nil
}

func randomToken() (string, error) {
	bytes := make([]byte, 32)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(bytes), nil
}

func pkceChallenge(verifier string) string {
	digest := sha256.Sum256([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(digest[:])
}

type Role string

const (
	RoleViewer        Role = "viewer"
	RoleOperator      Role = "operator"
	RoleEditor        Role = "editor"
	RoleAdministrator Role = "administrator"
)

type Permission string

const (
	PermissionView    Permission = "view"
	PermissionOperate Permission = "operate"
	PermissionEdit    Permission = "edit"
	PermissionAdmin   Permission = "admin"
)

type Principal struct {
	Subject string
	Role    Role
}

func MapRoles(claims []string, mapping map[string]Role) (Role, error) {
	priority := map[Role]int{RoleViewer: 1, RoleOperator: 2, RoleEditor: 3, RoleAdministrator: 4}
	var selected Role
	for _, claim := range claims {
		role, ok := mapping[claim]
		if !ok || priority[role] <= priority[selected] {
			continue
		}
		selected = role
	}
	if selected == "" {
		return "", errors.New("no authorized Pulse role")
	}
	return selected, nil
}

func Allows(role Role, permission Permission) bool {
	level := map[Role]int{RoleViewer: 1, RoleOperator: 2, RoleEditor: 3, RoleAdministrator: 4}[role]
	required := map[Permission]int{PermissionView: 1, PermissionOperate: 2, PermissionEdit: 3, PermissionAdmin: 4}[permission]
	return level > 0 && required > 0 && level >= required
}

func Require(permission Permission, next http.Handler) http.Handler {
	return http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		principal, ok := PrincipalFromContext(request.Context())
		if !ok {
			response.Header().Set("Cache-Control", "private, no-store")
			http.Error(response, "unauthorized", http.StatusUnauthorized)
			return
		}
		if !Allows(principal.Role, permission) {
			response.Header().Set("Cache-Control", "private, no-store")
			http.Error(response, "forbidden", http.StatusForbidden)
			return
		}
		next.ServeHTTP(response, request)
	})
}

type contextKey struct{}

func WithPrincipal(ctx context.Context, principal Principal) context.Context {
	return context.WithValue(ctx, contextKey{}, principal)
}

func PrincipalFromContext(ctx context.Context) (Principal, bool) {
	principal, ok := ctx.Value(contextKey{}).(Principal)
	return principal, ok && principal.Subject != ""
}

type BreakGlassPolicy struct {
	Enabled bool
}

func (policy BreakGlassPolicy) Allows() bool { return policy.Enabled }