@x47base/pocketbase-addon 0.1.0

package/security/security.go

@@ -0,0 +1,503 @@
+package security
+
+import (
+	"crypto/hmac"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"net/http"
+	"net/netip"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/pocketbase/pocketbase/apis"
+	"github.com/pocketbase/pocketbase/core"
+	"github.com/pocketbase/pocketbase/tools/hook"
+	"github.com/pocketbase/pocketbase/tools/router"
+)
+
+type bucket struct {
+	tokens  float64
+	updated time.Time
+}
+type Observation struct {
+	Family   string `json:"family"`
+	Count    int    `json:"count"`
+	Rejected int    `json:"rejected"`
+	Failures int    `json:"failures"`
+	Slow     int    `json:"slow"`
+}
+type Incident struct {
+	ID           uint64      `json:"id"`
+	Family       string      `json:"family"`
+	Phase        string      `json:"phase"`
+	FirstSeen    time.Time   `json:"firstSeen"`
+	LastSeen     time.Time   `json:"lastSeen"`
+	Evidence     Observation `json:"evidence"`
+	Acknowledged bool        `json:"acknowledged"`
+}
+type Action struct {
+	Family    string    `json:"family"`
+	ExpiresAt time.Time `json:"expiresAt"`
+	Reason    string    `json:"reason"`
+}
+type subscriptionState struct {
+	count    int
+	updating bool
+}
+
+type TrafficWindow struct {
+	At       time.Time `json:"at"`
+	Count    int       `json:"count"`
+	Rejected int       `json:"rejected"`
+	Failures int       `json:"failures"`
+	Slow     int       `json:"slow"`
+}
+
+type Service struct {
+	startedAt          time.Time
+	lastEvaluated      time.Time
+	history            []TrafficWindow
+	journalMu          sync.Mutex
+	storageError       bool
+	clients            map[string]*subscriptionState
+	totalSubscriptions int
+
+	mu           sync.Mutex
+	config       Config
+	revision     uint64
+	key          [32]byte
+	global       bucket
+	identities   map[string]*bucket
+	overflow     bucket
+	active       map[string]int
+	observations map[string]*Observation
+	streak       map[string]int
+	clear        map[string]int
+	incidents    []Incident
+	nextID       uint64
+	actions      map[string]Action
+	stop         chan struct{}
+	done         chan struct{}
+	once         sync.Once
+	total        uint64
+	rejected     uint64
+	lastSweep    time.Time
+}
+
+func New(c Config) (*Service, error) {
+	if err := c.Validate(); err != nil {
+		return nil, err
+	}
+	s := &Service{startedAt: time.Now(), clients: map[string]*subscriptionState{}, config: cloneConfig(c), revision: 1, identities: map[string]*bucket{}, active: map[string]int{}, observations: map[string]*Observation{}, streak: map[string]int{}, clear: map[string]int{}, actions: map[string]Action{}, stop: make(chan struct{}), done: make(chan struct{})}
+	if _, err := rand.Read(s.key[:]); err != nil {
+		return nil, err
+	}
+	s.loadState()
+	go func() {
+		defer close(s.done)
+		t := time.NewTicker(10 * time.Second)
+		defer t.Stop()
+		for {
+			select {
+			case now := <-t.C:
+				s.evaluate(now)
+				s.persist()
+			case <-s.stop:
+				return
+			}
+		}
+	}()
+	return s, nil
+}
+func (s *Service) Close() { s.once.Do(func() { close(s.stop) }); <-s.done; s.persist() }
+func Register(app core.App, c Config) error {
+	if app.Store().Has("pb.security.service") {
+		return errors.New("security already registered")
+	}
+	s, err := New(c)
+	if err != nil {
+		return err
+	}
+	app.Store().Set("pb.security.service", s)
+	app.OnServe().Bind(&hook.Handler[*core.ServeEvent]{Id: "pb.security.serve", Func: func(e *core.ServeEvent) error {
+		s.Bind(e.Router)
+		e.UIExtensions = append(e.UIExtensions, core.UIExtension{Name: "security", FS: dashboardFiles})
+		return e.Next()
+	}})
+	app.OnTerminate().Bind(&hook.Handler[*core.TerminateEvent]{Id: "pb.security.stop", Func: func(e *core.TerminateEvent) error { s.Close(); return e.Next() }})
+	app.OnBatchRequest().Bind(&hook.Handler[*core.BatchRequestEvent]{Id: "pb.security.batch", Func: func(e *core.BatchRequestEvent) error {
+		extra := max(0, len(e.Batch)-1)
+		s.mu.Lock()
+		mode := s.config.Mode
+		allowed := s.active["writes"]+extra <= s.config.MaxWrites
+		if mode == "disabled" {
+			s.mu.Unlock()
+			return e.Next()
+		}
+		if !allowed {
+			e.Set("pb.security.rejected", true)
+		}
+		if !allowed && mode == "enforce" {
+			s.mu.Unlock()
+			e.Response.Header().Set("Retry-After", "1")
+			return e.TooManyRequestsError("Batch operation budget exhausted", nil)
+		}
+		s.active["writes"] += extra
+		s.mu.Unlock()
+		defer func() { s.mu.Lock(); s.active["writes"] -= extra; s.mu.Unlock() }()
+		return e.Next()
+	}})
+	app.OnRealtimeConnectRequest().Bind(&hook.Handler[*core.RealtimeConnectRequestEvent]{Id: "pb.security.connect", Func: func(e *core.RealtimeConnectRequestEvent) error {
+		s.mu.Lock()
+		mode := s.config.Mode
+		if mode == "disabled" {
+			s.mu.Unlock()
+			return e.Next()
+		}
+		if len(s.clients) >= s.config.MaxRealtime {
+			s.mu.Unlock()
+			if mode == "enforce" {
+				return e.TooManyRequestsError("Realtime client budget exhausted", nil)
+			}
+			return e.Next()
+		}
+		s.clients[e.Client.Id()] = &subscriptionState{}
+		s.mu.Unlock()
+		defer func() {
+			s.mu.Lock()
+			if state := s.clients[e.Client.Id()]; state != nil {
+				s.totalSubscriptions -= state.count
+				delete(s.clients, e.Client.Id())
+			}
+			s.mu.Unlock()
+		}()
+		return e.Next()
+	}})
+	app.OnRealtimeSubscribeRequest().Bind(&hook.Handler[*core.RealtimeSubscribeRequestEvent]{Id: "pb.security.subscribe", Func: func(e *core.RealtimeSubscribeRequestEvent) error {
+		s.mu.Lock()
+		mode := s.config.Mode
+		state := s.clients[e.Client.Id()]
+		if mode == "disabled" || mode == "observe" && state == nil {
+			s.mu.Unlock()
+			return e.Next()
+		}
+		if mode == "observe" && state.updating {
+			s.mu.Unlock()
+			return e.Next()
+		}
+		if state == nil || state.updating {
+			s.mu.Unlock()
+			e.Response.Header().Set("Retry-After", "1")
+			return e.TooManyRequestsError("Realtime client unavailable or busy", nil)
+		}
+		old := state.count
+		count := len(e.Subscriptions)
+		allowed := count <= s.config.MaxSubscriptions && s.totalSubscriptions-old+count <= s.config.MaxSubscriptionsTotal
+		if !allowed {
+			e.Set("pb.security.rejected", true)
+		}
+		if !allowed && mode == "enforce" {
+			s.mu.Unlock()
+			e.Response.Header().Set("Retry-After", "1")
+			return e.TooManyRequestsError("Subscription budget exhausted", nil)
+		}
+		state.updating = true
+		state.count = count
+		s.totalSubscriptions += count - old
+		s.mu.Unlock()
+		defer func() {
+			s.mu.Lock()
+			if s.clients[e.Client.Id()] == state {
+				actual := len(e.Client.Subscriptions())
+				s.totalSubscriptions += actual - state.count
+				state.count = actual
+				state.updating = false
+			}
+			s.mu.Unlock()
+		}()
+		return e.Next()
+	}})
+	app.OnRealtimeMessageSend().Bind(&hook.Handler[*core.RealtimeMessageEvent]{Id: "pb.security.sse-write", Func: func(e *core.RealtimeMessageEvent) error {
+		s.mu.Lock()
+		enabled := s.config.Mode == "enforce"
+		s.mu.Unlock()
+		if enabled {
+			rc := http.NewResponseController(e.Response)
+			if err := rc.SetWriteDeadline(time.Now().Add(10 * time.Second)); err != nil && !errors.Is(err, http.ErrNotSupported) {
+				return err
+			}
+		}
+		return e.Next()
+	}})
+	return nil
+}
+func family(r *http.Request) string {
+	p := r.URL.Path
+	switch {
+	case strings.HasPrefix(p, "/api/security/"):
+		return "management"
+	case p == "/api/realtime":
+		return "realtime"
+	case strings.HasPrefix(p, "/api/files/"):
+		return "files"
+	case p == "/api/batch":
+		return "writes"
+	case strings.Contains(p, "/auth-") || strings.Contains(p, "/request-") || strings.Contains(p, "/confirm-"):
+		return "auth"
+	case r.Method == "POST" || r.Method == "PUT" || r.Method == "PATCH" || r.Method == "DELETE":
+		return "writes"
+	case strings.HasPrefix(p, "/api/collections/"):
+		return "reads"
+	default:
+		return "other"
+	}
+}
+func (s *Service) client(e *core.RequestEvent, c Config) string {
+	ip := e.RemoteIP()
+	if !inNetworks(ip, c.TrustedPeers) {
+		return ip
+	}
+	parts := strings.Split(e.Request.Header.Get("X-Forwarded-For"), ",")
+	if len(parts) > 32 {
+		return ip
+	}
+	for i := len(parts) - 1; i >= 0; i-- {
+		a, err := netip.ParseAddr(strings.TrimSpace(parts[i]))
+		if err != nil {
+			return ip
+		}
+		ip = a.Unmap().String()
+		if !inNetworks(ip, c.TrustedPeers) {
+			return ip
+		}
+	}
+	return ip
+}
+func (s *Service) fingerprint(value string) string {
+	m := hmac.New(sha256.New, s.key[:])
+	m.Write([]byte(value))
+	return hex.EncodeToString(m.Sum(nil)[:16])
+}
+func consume(b *bucket, rate float64, burst int, now time.Time) bool {
+	if b.updated.IsZero() {
+		b.tokens = float64(burst)
+	} else {
+		b.tokens = min(float64(burst), b.tokens+max(0, now.Sub(b.updated).Seconds())*rate)
+	}
+	b.updated = now
+	if b.tokens < 1 {
+		return false
+	}
+	b.tokens--
+	return true
+}
+func (s *Service) identityAllowed(key string, c Config, now time.Time) bool {
+	b := s.identities[key]
+	if b == nil {
+		if len(s.identities) >= c.MaxIdentities && now.Sub(s.lastSweep) > time.Minute {
+			for k, v := range s.identities {
+				if now.Sub(v.updated) > 2*time.Minute {
+					delete(s.identities, k)
+				}
+			}
+			s.lastSweep = now
+		}
+		if len(s.identities) >= c.MaxIdentities {
+			return consume(&s.overflow, c.IdentityPerSecond, c.IdentityBurst, now)
+		}
+		b = &bucket{}
+		s.identities[key] = b
+	}
+	return consume(b, c.IdentityPerSecond, c.IdentityBurst, now)
+}
+func (s *Service) observe(f string, rejected bool, status int, elapsed time.Duration) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	o := s.observations[f]
+	if o == nil {
+		o = &Observation{Family: f}
+		s.observations[f] = o
+	}
+	o.Count++
+	s.total++
+	if rejected {
+		o.Rejected++
+		s.rejected++
+	}
+	if status >= 400 {
+		o.Failures++
+	}
+	if elapsed > time.Second {
+		o.Slow++
+	}
+}
+func (s *Service) Bind(r *router.Router[*core.RequestEvent]) {
+	r.Bind(&hook.Handler[*core.RequestEvent]{Id: "pb.security.admission", Priority: apis.DefaultLoadAuthTokenMiddlewarePriority - 5, Func: func(e *core.RequestEvent) (err error) {
+		start := time.Now()
+		f := family(e.Request)
+		s.mu.Lock()
+		c := cloneConfig(s.config)
+		s.mu.Unlock()
+		if c.Mode == "disabled" {
+			return e.Next()
+		}
+		s.mu.Lock()
+		operator := e.Request.Header.Get("X-Spink-Operator") == "1" && inNetworks(e.RemoteIP(), c.OperatorPeers)
+		if operator {
+			f = "management"
+			e.Set("pb.security.operator", true)
+		}
+		management := operator || f == "management" && inNetworks(e.RemoteIP(), c.ManagementPeers)
+		allowed := true
+		status := 429
+		if management {
+			allowed = s.active["management"] < 4
+			status = 503
+		} else {
+			allowed = consume(&s.global, c.RequestsPerSecond, c.Burst, start)
+			count := s.active["ordinary"]
+			laneMax := c.MaxConcurrent
+			switch f {
+			case "writes":
+				laneMax = c.MaxWrites
+			case "files":
+				laneMax = c.MaxFiles
+			case "realtime":
+				laneMax = c.MaxRealtime
+				count = s.active["realtime"]
+			}
+			if count >= c.MaxConcurrent && f != "realtime" || s.active[f] >= laneMax {
+				allowed = false
+				status = 503
+			}
+			if a, ok := s.actions[f]; ok && start.Before(a.ExpiresAt) {
+				allowed = false
+			}
+		}
+		block := !allowed && c.Mode == "enforce"
+		if !block {
+			s.active[f]++
+			if f != "realtime" && f != "management" {
+				s.active["ordinary"]++
+			}
+		}
+		s.mu.Unlock()
+		if block {
+			s.observe(f, true, status, 0)
+			e.Response.Header().Set("Retry-After", "1")
+			e.Response.Header().Set("X-Content-Type-Options", "nosniff")
+			return router.NewApiError(status, "Security capacity limit reached", nil)
+		}
+		defer func() {
+			s.mu.Lock()
+			s.active[f]--
+			if f != "realtime" && f != "management" {
+				s.active["ordinary"]--
+			}
+			s.mu.Unlock()
+			status := e.Status()
+			if err != nil && status < 400 {
+				status = 400
+			}
+			s.observe(f, !allowed || e.Get("pb.security.rejected") == true, status, time.Since(start))
+		}()
+		return e.Next()
+	}})
+	r.Bind(&hook.Handler[*core.RequestEvent]{Id: "pb.security.identity", Priority: apis.DefaultRateLimitMiddlewarePriority - 5, Func: func(e *core.RequestEvent) error {
+		if family(e.Request) == "management" || e.Get("pb.security.operator") == true {
+			return e.Next()
+		}
+		s.mu.Lock()
+		c := s.config
+		allowed := true
+		if c.Mode != "disabled" {
+			key := "source:" + s.fingerprint(s.client(e, c))
+			if e.Auth != nil {
+				key = "user:" + s.fingerprint(e.Auth.Collection().Id+":"+e.Auth.Id)
+			}
+			allowed = s.identityAllowed(key, c, time.Now())
+		}
+		s.mu.Unlock()
+		if !allowed {
+			e.Set("pb.security.rejected", true)
+		}
+		if !allowed && c.Mode == "enforce" {
+			e.Response.Header().Set("Retry-After", "1")
+			return e.TooManyRequestsError("Identity budget exhausted", nil)
+		}
+		return e.Next()
+	}})
+	s.bindManagement(r)
+	s.bindReview(r)
+}
+func (s *Service) evaluate(now time.Time) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.config.Mode == "disabled" {
+		s.observations = map[string]*Observation{}
+		return
+	}
+	sample := TrafficWindow{At: now}
+	for _, o := range s.observations {
+		sample.Count += o.Count
+		sample.Rejected += o.Rejected
+		sample.Failures += o.Failures
+		sample.Slow += o.Slow
+	}
+	s.history = append(s.history, sample)
+	if len(s.history) > 60 {
+		s.history = append([]TrafficWindow(nil), s.history[len(s.history)-60:]...)
+	}
+	s.lastEvaluated = now
+	for _, f := range []string{"auth", "reads", "writes", "files", "realtime", "other"} {
+		o := s.observations[f]
+		if o == nil {
+			o = &Observation{Family: f}
+		}
+		suspicious := o.Count >= 20 && ((o.Rejected*2 > o.Count) || (f == "auth" && o.Failures*5 > o.Count*4) || (o.Slow*2 > o.Count && o.Failures*2 > o.Count))
+		if suspicious {
+			s.streak[f]++
+			s.clear[f] = 0
+		} else {
+			s.streak[f] = 0
+			s.clear[f]++
+		}
+		idx := -1
+		for i := len(s.incidents) - 1; i >= 0; i-- {
+			if s.incidents[i].Family == f && s.incidents[i].Phase != "resolved" {
+				idx = i
+				break
+			}
+		}
+		if s.streak[f] >= 3 {
+			if idx < 0 {
+				s.nextID++
+				s.incidents = append(s.incidents, Incident{ID: s.nextID, Family: f, FirstSeen: now.Add(-30 * time.Second), Phase: "active"})
+				if len(s.incidents) > 256 {
+					s.incidents = s.incidents[1:]
+				}
+				idx = len(s.incidents) - 1
+			}
+			s.incidents[idx].LastSeen = now
+			s.incidents[idx].Evidence = *o
+			s.incidents[idx].Phase = "active"
+		} else if idx >= 0 && s.clear[f] > 0 {
+			s.incidents[idx].Phase = "recovering"
+			if s.clear[f] >= 6 {
+				s.incidents[idx].Phase = "resolved"
+			}
+		}
+	}
+	s.observations = map[string]*Observation{}
+	for k, a := range s.actions {
+		if !now.Before(a.ExpiresAt) {
+			delete(s.actions, k)
+		}
+	}
+}
+
+// RequestFamily exposes the classifier to the pre-database gateway.
+func RequestFamily(r *http.Request) string { return family(r) }

package/security/security_test.go

@@ -0,0 +1,146 @@
+package security
+
+import (
+	"math"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/pocketbase/pocketbase/apis"
+	"github.com/pocketbase/pocketbase/core"
+	"github.com/pocketbase/pocketbase/tests"
+)
+
+func TestConfig(t *testing.T) {
+	for _, rate := range []float64{0, -1, math.NaN(), math.Inf(1)} {
+		c := DefaultConfig()
+		c.RequestsPerSecond = rate
+		if c.Validate() == nil {
+			t.Fatal("accepted invalid rate")
+		}
+	}
+}
+func TestBoundedIdentityState(t *testing.T) {
+	c := DefaultConfig()
+	c.MaxIdentities = 2
+	s, _ := New(c)
+	defer s.Close()
+	now := time.Now()
+	for _, id := range []string{"a", "b", "c", "d"} {
+		s.identityAllowed(id, c, now)
+	}
+	if len(s.identities) != 2 {
+		t.Fatal("unbounded identities")
+	}
+}
+func TestAdmissionModesAndRelease(t *testing.T) {
+	for _, mode := range []string{"disabled", "observe", "enforce"} {
+		t.Run(mode, func(t *testing.T) {
+			app, _ := tests.NewTestApp()
+			defer app.Cleanup()
+			c := DefaultConfig()
+			c.Mode = mode
+			c.Burst = 1
+			c.IdentityBurst = 1
+			c.RequestsPerSecond = .001
+			c.IdentityPerSecond = .001
+			s, _ := New(c)
+			defer s.Close()
+			r, _ := apis.NewRouter(app)
+			s.Bind(r)
+			r.GET("/probe", func(e *core.RequestEvent) error { return e.String(200, "ok") })
+			h, err := r.BuildMux()
+			if err != nil {
+				t.Fatal(err)
+			}
+			for i := 0; i < 3; i++ {
+				rec := httptest.NewRecorder()
+				h.ServeHTTP(rec, httptest.NewRequest("GET", "/probe", nil))
+				expected := 200
+				if mode == "enforce" && i > 0 {
+					expected = 429
+				}
+				if rec.Code != expected {
+					t.Fatalf("got %d expected %d", rec.Code, expected)
+				}
+			}
+			for _, n := range s.active {
+				if n != 0 {
+					t.Fatal("permit leaked")
+				}
+			}
+		})
+	}
+}
+func TestIdentityTrust(t *testing.T) {
+	s, _ := New(DefaultConfig())
+	defer s.Close()
+	e := &core.RequestEvent{}
+	e.Request = httptest.NewRequest("GET", "/", nil)
+	e.Request.RemoteAddr = "192.0.2.1:1000"
+	e.Request.Header.Set("X-Forwarded-For", "203.0.113.8, 10.0.0.2")
+	c := DefaultConfig()
+	if s.client(e, c) != "192.0.2.1" {
+		t.Fatal("trusted unknown peer")
+	}
+	c.TrustedPeers = []string{"192.0.2.0/24", "10.0.0.0/8"}
+	if s.client(e, c) != "203.0.113.8" {
+		t.Fatal("wrong proxy chain")
+	}
+}
+func TestDetectionRecoveryAndNoAutomaticAction(t *testing.T) {
+	s, _ := New(DefaultConfig())
+	defer s.Close()
+	now := time.Now()
+	for i := 0; i < 3; i++ {
+		for j := 0; j < 30; j++ {
+			s.observe("auth", false, 400, 0)
+		}
+		s.evaluate(now.Add(time.Duration(i) * 10 * time.Second))
+	}
+	if len(s.incidents) != 1 || s.incidents[0].Phase != "active" {
+		t.Fatal("missing incident")
+	}
+	if len(s.actions) != 0 {
+		t.Fatal("detector automatically blocks")
+	}
+	for i := 0; i < 6; i++ {
+		s.evaluate(now.Add(time.Duration(i+3) * 10 * time.Second))
+	}
+	if s.incidents[0].Phase != "resolved" {
+		t.Fatal("no recovery")
+	}
+}
+func TestManagementRequiresAuth(t *testing.T) {
+	app, _ := tests.NewTestApp()
+	defer app.Cleanup()
+	s, _ := New(DefaultConfig())
+	defer s.Close()
+	r, _ := apis.NewRouter(app)
+	s.Bind(r)
+	h, _ := r.BuildMux()
+	w := httptest.NewRecorder()
+	h.ServeHTTP(w, httptest.NewRequest(http.MethodPut, "/api/security/policy", strings.NewReader(`{}`)))
+	if w.Code != 401 {
+		t.Fatalf("unauthorized policy status %d", w.Code)
+	}
+}
+func TestConcurrentObservations(t *testing.T) {
+	s, _ := New(DefaultConfig())
+	defer s.Close()
+	var wg sync.WaitGroup
+	for i := 0; i < 20; i++ {
+		wg.Go(func() {
+			for j := 0; j < 100; j++ {
+				s.observe("reads", false, 200, 0)
+			}
+		})
+	}
+	wg.Wait()
+	if s.total != 2000 {
+		t.Fatal(s.total)
+	}
+}