tree-test-1770068050/pkg/httpcontext/keys.go

// Package httpcontext provides type-safe context keys and helpers for HTTP request contexts.
//
// This package standardizes how context values are stored and retrieved across all services.
// Using unexported types for context keys prevents collisions with other packages.
//
// Usage in middleware:
//
//	func AuthMiddleware() func(http.Handler) http.Handler {
//	    return func(next http.Handler) http.Handler {
//	        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
//	            user := extractUserFromAuth(r)
//	            ctx := httpcontext.SetUser(r.Context(), user)
//	            ctx = httpcontext.SetOrgID(ctx, user.OrganizationID)
//	            next.ServeHTTP(w, r.WithContext(ctx))
//	        })
//	    }
//	}
//
// Usage in handlers:
//
//	func (h *Handler) GetProfile(w http.ResponseWriter, r *http.Request) {
//	    user, ok := httpcontext.GetUser(r.Context())
//	    if !ok {
//	        http.Error(w, "unauthorized", http.StatusUnauthorized)
//	        return
//	    }
//	    // ... use user
//	}
package httpcontext

import "context"

// contextKey is an unexported type used for context keys to prevent collisions.
// Other packages cannot create values of this type, ensuring our keys are unique.
type contextKey string

// Standard context keys used across services.
const (
	keyUser      contextKey = "user"
	keyOrgID     contextKey = "organization_id"
	keyRequestID contextKey = "request_id"
	keyTraceID   contextKey = "trace_id"
	keyJWTClaims contextKey = "jwt_claims"
)

// SetUser adds a user to the context.
// The user can be any type - typically a domain user struct.
// Returns a new context with the user attached.
func SetUser(ctx context.Context, user any) context.Context {
	if user == nil {
		return ctx
	}
	return context.WithValue(ctx, keyUser, user)
}

// GetUser retrieves the user from context.
// Returns (user, true) if present, (nil, false) if not found.
// Caller should type-assert the returned value to their user type.
//
// Example:
//
//	if val, ok := httpcontext.GetUser(ctx); ok {
//	    user := val.(*domain.User)
//	    // ... use user
//	}
func GetUser(ctx context.Context) (any, bool) {
	val := ctx.Value(keyUser)
	if val == nil {
		return nil, false
	}
	return val, true
}

// SetOrgID adds an organization ID to the context.
// Returns a new context with the organization ID attached.
func SetOrgID(ctx context.Context, orgID string) context.Context {
	if orgID == "" {
		return ctx
	}
	return context.WithValue(ctx, keyOrgID, orgID)
}

// GetOrgID retrieves the organization ID from context.
// Returns (orgID, true) if present, ("", false) if not found.
func GetOrgID(ctx context.Context) (string, bool) {
	val := ctx.Value(keyOrgID)
	if val == nil {
		return "", false
	}
	if orgID, ok := val.(string); ok {
		return orgID, true
	}
	return "", false
}

// SetRequestID adds a request ID to the context.
// Returns a new context with the request ID attached.
func SetRequestID(ctx context.Context, requestID string) context.Context {
	if requestID == "" {
		return ctx
	}
	return context.WithValue(ctx, keyRequestID, requestID)
}

// GetRequestID retrieves the request ID from context.
// Returns (requestID, true) if present, ("", false) if not found.
func GetRequestID(ctx context.Context) (string, bool) {
	val := ctx.Value(keyRequestID)
	if val == nil {
		return "", false
	}
	if requestID, ok := val.(string); ok {
		return requestID, true
	}
	return "", false
}

// SetTraceID adds a trace ID to the context.
// Returns a new context with the trace ID attached.
func SetTraceID(ctx context.Context, traceID string) context.Context {
	if traceID == "" {
		return ctx
	}
	return context.WithValue(ctx, keyTraceID, traceID)
}

// GetTraceID retrieves the trace ID from context.
// Returns (traceID, true) if present, ("", false) if not found.
func GetTraceID(ctx context.Context) (string, bool) {
	val := ctx.Value(keyTraceID)
	if val == nil {
		return "", false
	}
	if traceID, ok := val.(string); ok {
		return traceID, true
	}
	return "", false
}

// SetJWTClaims adds JWT claims to the context.
// The claims can be any type - typically a custom claims struct.
// Returns a new context with the claims attached.
func SetJWTClaims(ctx context.Context, claims any) context.Context {
	if claims == nil {
		return ctx
	}
	return context.WithValue(ctx, keyJWTClaims, claims)
}

// GetJWTClaims retrieves JWT claims from context.
// Returns (claims, true) if present, (nil, false) if not found.
// Caller should type-assert the returned value to their claims type.
//
// Example:
//
//	if val, ok := httpcontext.GetJWTClaims(ctx); ok {
//	    claims := val.(*auth.CustomClaims)
//	    // ... use claims
//	}
func GetJWTClaims(ctx context.Context) (any, bool) {
	val := ctx.Value(keyJWTClaims)
	if val == nil {
		return nil, false
	}
	return val, true
}

// MustGetUser retrieves the user from context and panics if not found.
// Use only when authentication middleware guarantees user presence.
func MustGetUser(ctx context.Context) any {
	user, ok := GetUser(ctx)
	if !ok {
		panic("httpcontext: user not found in context")
	}
	return user
}

// MustGetRequestID retrieves the request ID from context and panics if not found.
// Use only when middleware guarantees request ID presence.
func MustGetRequestID(ctx context.Context) string {
	requestID, ok := GetRequestID(ctx)
	if !ok {
		panic("httpcontext: request_id not found in context")
	}
	return requestID
}