Janady13/missionmatch / internal/platform/tenant/tenant.go

Source identity

Field	Recorded value
Source type	GitHub source
Original record	Janady13/missionmatch / internal/platform/tenant/tenant.go
Date	Not specified in source
Repository	Janady13/missionmatch
Revision	39df94f92b2a33f852481fe1954dabed841524bb
Branch	w0-aedt-compliance
Path or archive member	internal/platform/tenant/tenant.go
Original SHA-256	904b45631bcecff995bc4bf455f93825adc77efa605dd4af73c36b1e6344954f
Git blob SHA	73d26789d2875d77efaf964d625711a39db89b75
Reading-copy SHA-256	904b45631bcecff995bc4bf455f93825adc77efa605dd4af73c36b1e6344954f

Original source ↗

Recovered source transcript

Download complete source text ↓138 lines · 5.2 KB

// Package tenant is the multi-tenancy correctness boundary.
//
// Tenant isolation is enforced by Postgres Row-Level Security, NOT by Go WHERE clauses. Every
// tenant-scoped request runs inside a transaction whose `app.current_org_id` GUC is set
// transaction-locally, so Postgres discards it at COMMIT/ROLLBACK and it can never leak to the next
// request that reuses a pooled connection (the pgx-pool gotcha — see docs/STACK_DECISIONS.md §2).
//
// Two fail-safes stack here:
//   - Go side: no validated org in context => RunInTenantTx refuses to run (fail-closed).
//   - DB side: an unset GUC makes current_setting(...,true) NULL, so the RLS policy matches zero
//     rows (default-deny).
//
// The org id MUST originate from a verified session/JWT claim resolved by auth middleware — never
// from a client-supplied body, header, or query parameter.
package tenant

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

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// ErrNoTenant is returned when tenant-scoped DB work is attempted without a validated organization.
var ErrNoTenant = errors.New("tenant: no organization_id in context (refusing to run unscoped query)")

// ctxKey is an unexported context key so a client cannot spoof the org via a context-key collision.
type ctxKey struct{}

var orgKey = ctxKey{}

// WithOrg stores the validated org id in ctx. Call this ONLY from auth middleware, with a value
// taken from verified claims.
func WithOrg(ctx context.Context, orgID string) context.Context {
	return context.WithValue(ctx, orgKey, orgID)
}

// OrgFromContext returns the validated org id and whether a non-empty one is present.
func OrgFromContext(ctx context.Context) (string, bool) {
	v, ok := ctx.Value(orgKey).(string)
	return v, ok && v != ""
}

// principalCtxKey is a SEPARATE key type from ctxKey. Reusing ctxKey for the user would collide with
// orgKey, because every empty struct value of the same type is equal, so the two context lookups
// would return each other's value. A distinct type keeps the org and the principal independent.
type principalCtxKey struct{}

var principalKey = principalCtxKey{}

// Principal is the authenticated actor resolved from a validated session: the user id and the role
// used for authorization. Like the org, it MUST originate from auth middleware reading verified
// session state, never a client supplied value.
type Principal struct {
	UserID string
	Role   string
}

// WithUser stores the validated principal in ctx. Call this ONLY from auth middleware.
func WithUser(ctx context.Context, p Principal) context.Context {
	return context.WithValue(ctx, principalKey, p)
}

// UserFromContext returns the validated principal and whether one with a non-empty user id is
// present.
func UserFromContext(ctx context.Context) (Principal, bool) {
	p, ok := ctx.Value(principalKey).(Principal)
	return p, ok && p.UserID != ""
}

// DB wraps a pgxpool.Pool whose DSN authenticates as the NOBYPASSRLS application role.
type DB struct{ Pool *pgxpool.Pool }

// New wraps a pool as a tenant-scoped DB.
func New(pool *pgxpool.Pool) *DB { return &DB{Pool: pool} }

// Close releases the underlying pool.
func (db *DB) Close() {
	if db.Pool != nil {
		db.Pool.Close()
	}
}

// RunInTenantTx is THE correctness boundary. It:
//  1. requires a validated org id (fail-closed: no org => no query, before the pool is touched),
//  2. opens a transaction on a single pooled connection,
//  3. sets app.current_org_id transaction-locally via set_config(..., is_local=true) so it
//     auto-resets at COMMIT/ROLLBACK and cannot leak across the pool,
//  4. runs fn against that same tx, then commits (rolling back on any error).
//
// All tenant-scoped repository work MUST go through this method.
func (db *DB) RunInTenantTx(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) (err error) {
	orgID, ok := OrgFromContext(ctx)
	if !ok {
		return ErrNoTenant
	}

	tx, err := db.Pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("begin tenant tx: %w", err)
	}
	// Rollback is a no-op once Commit succeeds; this guarantees the connection is released on every path.
	defer func() { _ = tx.Rollback(ctx) }()

	// Transaction-local GUC, parameterized (no SQL-injection surface). is_local=true is what makes
	// this safe on a pooled connection.
	if _, err = tx.Exec(ctx, "SELECT set_config('app.current_org_id', $1, true)", orgID); err != nil {
		return fmt.Errorf("set tenant guc: %w", err)
	}

	if err = fn(ctx, tx); err != nil {
		return err // defer rolls back
	}
	if err = tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit tenant tx: %w", err)
	}
	return nil
}

// Middleware resolves the org id from an already-validated request and stores it in context.
// resolveOrg MUST read verified session/JWT claims, never a raw client value. A request whose org
// cannot be resolved is rejected with 403 before any downstream handler runs.
func Middleware(resolveOrg func(*http.Request) (string, bool)) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			orgID, ok := resolveOrg(r)
			if !ok {
				http.Error(w, "forbidden", http.StatusForbidden)
				return
			}
			next.ServeHTTP(w, r.WithContext(WithOrg(r.Context(), orgID)))
		})
	}
}