74 lines
2.2 KiB
Go
74 lines
2.2 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"net/http"
|
||
|
|
"regexp"
|
||
|
|
|
||
|
|
"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) func(http.Handler) http.Handler {
|
||
|
|
return func(next http.Handler) http.Handler {
|
||
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
tenantID := r.Header.Get("X-Tenant-Id")
|
||
|
|
if tenantID == "" {
|
||
|
|
w.Header().Set("Content-Type", "application/json")
|
||
|
|
w.WriteHeader(http.StatusBadRequest)
|
||
|
|
_, _ = w.Write([]byte(`{"error": "Missing X-Tenant-Id header"}`))
|
||
|
|
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")
|
||
|
|
}
|