go-guess/backend/internal/web/middleware/TenantSchemaMiddleware.go

package middleware

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

	"github.com/danyel/go-guess/backend/internal/golooseauth"
	"github.com/danyel/go-loose/client"
	"gorm.io/gorm"
)

type contextKey string

const TxKey contextKey = "tenant_db_tx"

// Alphanumeric or underscores only to prevent SQL Injection in schema names
var safeSchemaRegex = regexp.MustCompile(`^[a-zA-Z0-9_:-]+$`)

// TenantSchemaMiddleware isolates database connections per request using PostgreSQL schemas
func TenantSchemaMiddleware(db *gorm.DB, appDomain string) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path == "/api/health" || r.URL.Path == "/api/openapi.json" {
				next.ServeHTTP(w, r)
				return
			}
			tenantID := r.Header.Get(client.XTenantId)
			if tenantID == "" {
				tenantID = golooseauth.TenantFromHost(r.Host, appDomain)
			}
			if tenantID == "" || db == nil {
				next.ServeHTTP(w, r)
				return
			}

			// Validate schema name to guarantee safety against SQL injection
			if !safeSchemaRegex.MatchString(tenantID) {
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusBadRequest)
				_, _ = w.Write([]byte(`{"error": "Invalid characters in tenant ID"}`))
				return
			}

			// Start a fresh transaction scoped exclusively to this HTTP request
			tx := db.Begin()

			// Set the local search_path for the duration of this transaction block
			if err := tx.Exec(fmt.Sprintf("SET LOCAL search_path TO %s", tenantID)).Error; err != nil {
				tx.Rollback()
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusInternalServerError)
				_, _ = w.Write([]byte(`{"error": "Failed to resolve tenant storage isolation"}`))
				return
			}

			// Embed the transaction database instance into the request context
			ctx := context.WithValue(r.Context(), TxKey, tx)

			// Execute the subsequent request handlers
			next.ServeHTTP(w, r.WithContext(ctx))

			// Commit or Rollback based on the transaction error status
			if tx.Error != nil {
				tx.Rollback()
			} else {
				tx.Commit()
			}
		})
	}
}

// GetDB retrieves the transaction-bound tenant GORM DB instance from context
func GetDB(ctx context.Context) *gorm.DB {
	if tx, ok := ctx.Value(TxKey).(*gorm.DB); ok {
		return tx.WithContext(ctx)
	}
	panic("GetDB was invoked outside of TenantSchemaMiddleware")
}