go-ohif-proxy/internal/auth/jwt.go

package auth

import (
	"errors"
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

var (
	ErrInvalidToken = errors.New("invalid token")
	ErrExpiredToken = errors.New("token has expired")
)

// JWTManager handles JWT token operations
type JWTManager struct {
	secretKey     string
	accessExpiry  time.Duration
	refreshExpiry time.Duration
}

// NewJWTManager creates a new JWT manager
func NewJWTManager(secretKey string, accessExpiry, refreshExpiry time.Duration) *JWTManager {
	return &JWTManager{
		secretKey:     secretKey,
		accessExpiry:  accessExpiry,
		refreshExpiry: refreshExpiry,
	}
}

// CustomClaims contains the claims we want in our tokens
type CustomClaims struct {
	UserID    string `json:"user_id"`
	Email     string `json:"email"`
	Role      string `json:"role"`
	UserName  string `json:"user_name"`
	TokenType string `json:"token_type"` // access or refresh

	// Patient-specific fields
	PatientID        string   `json:"patient_id,omitempty"`
	PatientName      string   `json:"patient_name,omitempty"`
	AccessionNumber  string   `json:"accession_number,omitempty"`  // For backward compatibility
	AccessionNumbers []string `json:"accession_numbers,omitempty"` // Multiple accession numbers
	StudyIUID        string   `json:"study_iuid,omitempty"`        // For backward compatibility
	StudyIUIDs       []string `json:"study_iuids,omitempty"`       // Multiple study IUIDs

	// Navigation and permissions
	HomeURL   string `json:"home_url,omitempty"`
	StudyList string `json:"study_list,omitempty"` // enabled or disabled
	FilterURL string `json:"filter_url,omitempty"` // for ref_doctor filtering

	jwt.RegisteredClaims
}

// GenerateAccessToken creates a new access token with role-specific claims
func (m *JWTManager) GenerateAccessToken(userID, email, role, userName string, additionalClaims map[string]interface{}) (string, error) {
	claims := CustomClaims{
		UserID:    userID,
		Email:     email,
		Role:      role,
		UserName:  userName,
		TokenType: "access",
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(m.accessExpiry)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
		},
	}

	// Add role-specific additional claims
	if additionalClaims != nil {
		// Handle string claims
		if val, ok := additionalClaims["patient_id"].(string); ok {
			claims.PatientID = val
		}
		if val, ok := additionalClaims["patient_name"].(string); ok {
			claims.PatientName = val
		}
		if val, ok := additionalClaims["accession_number"].(string); ok {
			claims.AccessionNumber = val
		}
		if val, ok := additionalClaims["study_iuid"].(string); ok {
			claims.StudyIUID = val
		}
		if val, ok := additionalClaims["home_url"].(string); ok {
			claims.HomeURL = val
		}
		if val, ok := additionalClaims["study_list"].(string); ok {
			claims.StudyList = val
		}
		if val, ok := additionalClaims["filter_url"].(string); ok {
			claims.FilterURL = val
		}

		// Handle array claims
		if val, ok := additionalClaims["study_iuids"].([]string); ok {
			claims.StudyIUIDs = val
		}
		if val, ok := additionalClaims["accession_numbers"].([]string); ok {
			claims.AccessionNumbers = val
		}
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(m.secretKey))
}

// GenerateRefreshToken creates a new refresh token with the same claims as access token
func (m *JWTManager) GenerateRefreshToken(userID, email, role, userName string, additionalClaims map[string]interface{}) (string, error) {
	claims := CustomClaims{
		UserID:    userID,
		Email:     email,
		Role:      role,
		UserName:  userName,
		TokenType: "refresh",
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(m.refreshExpiry)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
		},
	}

	// Add role-specific additional claims
	if additionalClaims != nil {
		// Handle string claims
		if val, ok := additionalClaims["patient_id"].(string); ok {
			claims.PatientID = val
		}
		if val, ok := additionalClaims["patient_name"].(string); ok {
			claims.PatientName = val
		}
		if val, ok := additionalClaims["accession_number"].(string); ok {
			claims.AccessionNumber = val
		}
		if val, ok := additionalClaims["study_iuid"].(string); ok {
			claims.StudyIUID = val
		}
		if val, ok := additionalClaims["home_url"].(string); ok {
			claims.HomeURL = val
		}
		if val, ok := additionalClaims["study_list"].(string); ok {
			claims.StudyList = val
		}
		if val, ok := additionalClaims["filter_url"].(string); ok {
			claims.FilterURL = val
		}

		// Handle array claims
		if val, ok := additionalClaims["study_iuids"].([]string); ok {
			claims.StudyIUIDs = val
		}
		if val, ok := additionalClaims["accession_numbers"].([]string); ok {
			claims.AccessionNumbers = val
		}
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString([]byte(m.secretKey))
}

// ValidateToken validates a token and returns the claims
func (m *JWTManager) ValidateToken(tokenString string) (*CustomClaims, error) {
	token, err := jwt.ParseWithClaims(
		tokenString,
		&CustomClaims{},
		func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
			}
			return []byte(m.secretKey), nil
		},
	)

	if err != nil {
		if errors.Is(err, jwt.ErrTokenExpired) {
			return nil, ErrExpiredToken
		}
		return nil, ErrInvalidToken
	}

	claims, ok := token.Claims.(*CustomClaims)
	if !ok || !token.Valid {
		return nil, ErrInvalidToken
	}

	return claims, nil
}