edge/pkg/module/auth/option.go

package auth

import (
	"context"
	"fmt"
	"net/http"

	"forge.cadoles.com/arcad/edge/pkg/jwtutil"
	"github.com/pkg/errors"
)

type GetClaimFunc func(ctx context.Context, r *http.Request, name string) (string, error)

type Option struct {
	GetClaim      GetClaimFunc
	ProfileClaims []string
}

type OptionFunc func(*Option)

func defaultOptions() *Option {
	return &Option{
		GetClaim: dummyGetClaim,
		ProfileClaims: []string{
			ClaimSubject,
			ClaimIssuer,
			ClaimEdgeEntrypoint,
			ClaimEdgeRole,
			ClaimPreferredUsername,
			ClaimEdgeTenant,
		},
	}
}

func dummyGetClaim(ctx context.Context, r *http.Request, name string) (string, error) {
	return "", errors.Errorf("dummy getclaim func cannot retrieve claim '%s'", name)
}

func WithGetClaims(fn GetClaimFunc) OptionFunc {
	return func(o *Option) {
		o.GetClaim = fn
	}
}

func WithProfileClaims(claims ...string) OptionFunc {
	return func(o *Option) {
		o.ProfileClaims = claims
	}
}

func WithJWT(getKeySet jwtutil.GetKeySetFunc) OptionFunc {
	funcs := []jwtutil.FindTokenOptionFunc{
		jwtutil.WithFinders(
			jwtutil.FindTokenFromAuthorizationHeader,
			jwtutil.FindTokenFromQueryString(CookieName),
			jwtutil.FindTokenFromCookie(CookieName),
		),
	}

	return func(o *Option) {
		o.GetClaim = func(ctx context.Context, r *http.Request, name string) (string, error) {
			token, err := jwtutil.FindToken(r, getKeySet, funcs...)
			if err != nil {
				return "", errors.WithStack(err)
			}

			tokenMap, err := token.AsMap(ctx)
			if err != nil {
				return "", errors.WithStack(err)
			}

			value, exists := tokenMap[name]
			if !exists {
				return "", nil
			}

			return fmt.Sprintf("%v", value), nil
		}
	}
}