go-ohif-proxy/internal/api/middleware/auth.go

package middleware

import (
	"context"
	"encoding/json"
	"net/http"
	"regexp"
	"strings"

	"devone.aplikasi.web.id/gitea/mario/go-ohif-proxy/internal/api/service"

	"go.uber.org/zap"
)

type contextKey string

const (
	UserIDKey    contextKey = "user_id"
	UserRoleKey  contextKey = "user_role"
	UserEmailKey contextKey = "user_email"
)

// WhitelistedEndpoints contains paths that can be accessed without authentication
var WhitelistedEndpoints = []*regexp.Regexp{
	// Study by UID
	regexp.MustCompile(`^/dicomWeb/studies\?.*StudyInstanceUID=.+`),

	// Frame endpoint
	regexp.MustCompile(`^/dicomWeb/studies/[^/]+/series/[^/]+/instances/[^/]+/frames/\d+$`),
}

// Auth middleware authenticates requests using JWT tokens
func Auth(authService *service.AuthService, logger *zap.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Check if the request path is whitelisted
			path := r.URL.Path
			if r.URL.RawQuery != "" {
				path = path + "?" + r.URL.RawQuery
			}

			for _, pattern := range WhitelistedEndpoints {
				if pattern.MatchString(path) {
					// Path is whitelisted, skip authentication
					logger.Debug("Skipping authentication for whitelisted path", zap.String("path", path))
					next.ServeHTTP(w, r)
					return
				}
			}

			// Get authorization header
			authHeader := r.Header.Get("Authorization")
			if authHeader == "" {
				logger.Warn("Missing Authorization header")
				respondWithError(w, http.StatusUnauthorized, "missing authorization header")
				return
			}

			// Extract token from Bearer token
			bearerToken := strings.Split(authHeader, " ")
			if len(bearerToken) != 2 || strings.ToLower(bearerToken[0]) != "bearer" {
				logger.Warn("Invalid Authorization header format")
				respondWithError(w, http.StatusUnauthorized, "invalid authorization format")
				return
			}

			token := bearerToken[1]

			// Validate token
			claims, err := authService.ValidateToken(token)
			if err != nil {
				logger.Warn("Invalid or expired token", zap.Error(err))
				respondWithError(w, http.StatusUnauthorized, "invalid or expired token")
				return
			}

			// Check token type
			if claims.TokenType != "access" {
				logger.Warn("Invalid token type", zap.String("tokenType", claims.TokenType))
				respondWithError(w, http.StatusUnauthorized, "invalid token type")
				return
			}

			// Add user info to request context
			ctx := context.WithValue(r.Context(), UserIDKey, claims.UserID)
			ctx = context.WithValue(ctx, UserRoleKey, claims.Role)
			ctx = context.WithValue(ctx, UserEmailKey, claims.Email)

			// Log user info
			logger.Info("Authenticated user", zap.String("userID", claims.UserID), zap.String("role", claims.Role), zap.String("email", claims.Email))

			// Continue with the request
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// RoleRequired middleware checks if user has the required role
func RoleRequired(roles ...string) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Check if the request path is whitelisted first
			path := r.URL.Path
			if r.URL.RawQuery != "" {
				path = path + "?" + r.URL.RawQuery
			}

			for _, pattern := range WhitelistedEndpoints {
				if pattern.MatchString(path) {
					// Path is whitelisted, skip role check
					next.ServeHTTP(w, r)
					return
				}
			}

			// Get user role from context
			userRole, ok := r.Context().Value(UserRoleKey).(string)
			if !ok {
				respondWithError(w, http.StatusUnauthorized, "user context not found")
				return
			}

			// Check if user has one of the required roles
			hasRole := false
			for _, role := range roles {
				if userRole == role {
					hasRole = true
					break
				}
			}

			if !hasRole {
				respondWithError(w, http.StatusForbidden, "insufficient permissions")
				return
			}

			// Continue with the request
			next.ServeHTTP(w, r)
		})
	}
}

// PatientViewRestriction middleware restricts patients to view only their studies
func PatientViewRestriction(logger *zap.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Check if the request path is whitelisted first
			path := r.URL.Path
			if r.URL.RawQuery != "" {
				path = path + "?" + r.URL.RawQuery
			}

			for _, pattern := range WhitelistedEndpoints {
				if pattern.MatchString(path) {
					// Path is whitelisted, skip restrictions
					logger.Debug("Skipping patient restrictions for whitelisted path", zap.String("path", path))
					next.ServeHTTP(w, r)
					return
				}
			}

			// Get user role from context
			userRole, ok := r.Context().Value(UserRoleKey).(string)
			if !ok {
				respondWithError(w, http.StatusUnauthorized, "user context not found")
				return
			}

			// Only apply restrictions to patients
			if userRole != "patient" {
				next.ServeHTTP(w, r)
				return
			}

			// TODO: Logic to restrict patients to only access their assigned studies
			// For now, we're just letting the request through, but in a real
			// implementation, you would check the study ID against the patient's
			// assigned studies.

			// TODO: Check if the requested study is assigned to the patient
			// This would likely involve parsing the URL path to extract study ID
			// and checking it against a database of patient assignments

			next.ServeHTTP(w, r)
		})
	}
}

// Helper function to respond with JSON error
func respondWithError(w http.ResponseWriter, code int, message string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(map[string]string{"error": message})
}