auth/authmw/token.go

package authmw

import (
	"context"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"net/http"
	"slices"
	"strings"
	"time"

	"git.t-juice.club/microfilm/auth"
	"github.com/golang-jwt/jwt/v5"
	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	"go.opentelemetry.io/otel"
)

type ctxType string

var ctxKeyClaims ctxType = "claims"

var ErrNoClaimsInRequest = fmt.Errorf("no claims in request")

func VerifyToken(authURL string, permittedRoles []string) func(http.Handler) http.Handler {
	fn := func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			ctx, span := otel.GetTracerProvider().Tracer("").Start(r.Context(), "verify-token")
			defer span.End()

			authHeader := r.Header.Get("Authorization")
			if !strings.Contains(authHeader, "Bearer ") {
				// No token, pass if unathorized in permitted
				// else reject
				if slices.Contains[[]string, string](permittedRoles, auth.RoleUnauthorized) {
					next.ServeHTTP(w, r)
					return
				}

				// Reject and write error response
				w.WriteHeader(http.StatusUnauthorized)
				var errResp auth.ErrorResponse
				errResp.Message = fmt.Sprintf("Authorization required: %s", strings.Join(permittedRoles, ","))
				errResp.Status = http.StatusUnauthorized

				encoder := json.NewEncoder(w)
				_ = encoder.Encode(&errResp)
				return
			}

			// Fetch current pubkey
			url := fmt.Sprintf("%s/key", authURL)
			ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
			defer cancel()

			req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
			if err != nil {
				// TODO: Should log
				w.WriteHeader(http.StatusInternalServerError)
				errResp := &auth.ErrorResponse{
					Message: fmt.Sprintf("Error getting pubkey for token verification: %s", err),
					Status:  http.StatusUnauthorized,
				}

				encoder := json.NewEncoder(w)
				_ = encoder.Encode(&errResp)
				return
			}

			client := &http.Client{
				Transport: otelhttp.NewTransport(http.DefaultTransport),
			}

			resp, err := client.Do(req)
			if err != nil {
				// TODO: Should log
				w.WriteHeader(http.StatusInternalServerError)
				errResp := &auth.ErrorResponse{
					Message: fmt.Sprintf("Error getting pubkey for token verification: %s", err),
					Status:  http.StatusUnauthorized,
				}

				encoder := json.NewEncoder(w)
				_ = encoder.Encode(&errResp)
				return
			}
			defer resp.Body.Close()

			var authResponse auth.PubkeyResponse
			decoder := json.NewDecoder(resp.Body)
			if err := decoder.Decode(&authResponse); err != nil {
				// TODO: Should log
				w.WriteHeader(http.StatusInternalServerError)
				errResp := &auth.ErrorResponse{
					Message: fmt.Sprintf("Error getting pubkey for token verification: %s", err),
					Status:  http.StatusUnauthorized,
				}

				encoder := json.NewEncoder(w)
				_ = encoder.Encode(&errResp)
				return
			}

			// Parse pubkey
			pub, err := x509.ParsePKIXPublicKey(authResponse.PubKey)
			if err != nil {
				// TODO: Should log
				w.WriteHeader(http.StatusInternalServerError)
				errResp := &auth.ErrorResponse{
					Message: fmt.Sprintf("Error getting pubkey for token verification: %s", err),
					Status:  http.StatusUnauthorized,
				}

				encoder := json.NewEncoder(w)
				_ = encoder.Encode(&errResp)
				return
			}

			// Validate token
			tokenString := strings.Split(authHeader, " ")[1]
			token, err := jwt.ParseWithClaims(tokenString, &auth.MicrofilmClaims{}, func(t *jwt.Token) (interface{}, error) { return pub, nil })
			if err != nil {
				// Reject and write error response
				w.WriteHeader(http.StatusUnauthorized)
				var errResp auth.ErrorResponse
				errResp.Message = fmt.Sprintf("Token verification failed: %s", err)
				errResp.Status = http.StatusUnauthorized

				encoder := json.NewEncoder(w)
				_ = encoder.Encode(&errResp)
				return
			}
			// TODO: Check that it is in permitted

			// Add claims to request context
			if claims, ok := token.Claims.(*auth.MicrofilmClaims); ok && token.Valid {
				ctx := context.WithValue(r.Context(), ctxKeyClaims, claims)
				next.ServeHTTP(w, r.WithContext(ctx))
				return
			}

			next.ServeHTTP(w, r)
		}
		return http.HandlerFunc(fn)
	}

	return fn
}

func ClaimsFromCtx(ctx context.Context) (*auth.MicrofilmClaims, error) {
	rawValue := ctx.Value(ctxKeyClaims)
	value, ok := rawValue.(*auth.MicrofilmClaims)
	if ok {
		return value, nil
	}

	return nil, ErrNoClaimsInRequest
}