From 9b7cd07e542baba12c4b4535e42c4226aebab3e8 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Sat, 17 Oct 2026 13:34:21 +0000 Subject: [PATCH] db: keep only the current policy Every write inserted a row and nothing deleted one, while every reader takes only the newest. Writers store and install under one lock, so an uncommitted insert cannot outlive the newer row that pruned the rest. --- hscontrol/db/policy.go | 29 ++++++---- hscontrol/db/policy_test.go | 31 +++++++++++ hscontrol/state/persist_test.go | 95 +++++++++++++++++++++++++++++++++ hscontrol/state/state.go | 16 ++++-- 4 files changed, 158 insertions(+), 13 deletions(-) create mode 100644 hscontrol/db/policy_test.go diff --git a/hscontrol/db/policy.go b/hscontrol/db/policy.go index 3c60534d6..da17c140b 100644 --- a/hscontrol/db/policy.go +++ b/hscontrol/db/policy.go @@ -10,19 +10,28 @@ import ( "gorm.io/gorm/clause" ) -// SetPolicy sets the policy in the database. +// SetPolicy replaces the policy in the database. Only the newest row is +// ever read, so older ones are removed rather than kept as history. func (hsdb *HSDatabase) SetPolicy(policy string) (*types.Policy, error) { - // Create a new policy. - p := types.Policy{ - Data: policy, - } + return Write(hsdb.DB, func(tx *gorm.DB) (*types.Policy, error) { + p := types.Policy{ + Data: policy, + } - err := hsdb.DB.Clauses(clause.Returning{}).Create(&p).Error - if err != nil { - return nil, err - } + err := tx.Clauses(clause.Returning{}).Create(&p).Error + if err != nil { + return nil, err + } - return &p, nil + // Unscoped: a soft delete would keep the data. Deleting by id + // never removes a newer row from a concurrent writer. + err = tx.Unscoped().Where("id < ?", p.ID).Delete(&types.Policy{}).Error + if err != nil { + return nil, err + } + + return &p, nil + }) } // GetPolicy returns the latest policy in the database. diff --git a/hscontrol/db/policy_test.go b/hscontrol/db/policy_test.go new file mode 100644 index 000000000..f62802db0 --- /dev/null +++ b/hscontrol/db/policy_test.go @@ -0,0 +1,31 @@ +package db + +import ( + "strconv" + "testing" + + "github.com/juanfont/headscale/hscontrol/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestSetPolicyKeepsOnlyCurrent checks repeated writes replace the stored +// policy instead of piling up rows. +func TestSetPolicyKeepsOnlyCurrent(t *testing.T) { + db := dbForTest(t) + + for i := range 5 { + _, err := db.SetPolicy( + `{"hosts": {"h": "100.64.0.` + strconv.Itoa(i) + `/32"}}`, + ) + require.NoError(t, err) + } + + var rows int64 + require.NoError(t, db.DB.Unscoped().Model(&types.Policy{}).Count(&rows).Error) + assert.Equal(t, int64(1), rows) + + p, err := db.GetPolicy() + require.NoError(t, err) + assert.JSONEq(t, `{"hosts": {"h": "100.64.0.4/32"}}`, p.Data) +} diff --git a/hscontrol/state/persist_test.go b/hscontrol/state/persist_test.go index 042db6381..77e853d33 100644 --- a/hscontrol/state/persist_test.go +++ b/hscontrol/state/persist_test.go @@ -6,6 +6,7 @@ import ( "maps" "net/http" "net/netip" + "net/url" "os" "path/filepath" "runtime" @@ -30,6 +31,7 @@ import ( "tailscale.com/tailcfg" "tailscale.com/types/key" "tailscale.com/types/views" + "zombiezen.com/go/postgrestest" ) // persistTestSetup pre-creates a sqlite database on disk with a single @@ -1341,6 +1343,99 @@ func TestUpdatePolicyStoreFailure(t *testing.T) { assert.Equal(t, before, after, "an unstored policy must not be installed") } +// TestPostgresUpdatePolicyConcurrentWrites runs a second write while the +// first one's insert is not yet committed. Pruning cannot see that row, so +// unless writers are serialised up to the install, it outlives the newer one. +func TestPostgresUpdatePolicyConcurrentWrites(t *testing.T) { + srv, err := postgrestest.Start(t.Context()) + if err != nil { + t.Skipf("start postgres: %s", err) + } + + t.Cleanup(srv.Cleanup) + + dsn, err := srv.CreateDatabase(t.Context()) + require.NoError(t, err) + + u, err := url.Parse(dsn) + require.NoError(t, err) + + pass, _ := u.User.Password() + port, err := strconv.Atoi(u.Port()) + require.NoError(t, err) + + cfg := persistTestConfig("") + cfg.Database = types.DatabaseConfig{ + Type: types.DatabasePostgres, + Postgres: types.PostgresConfig{ + Host: u.Hostname(), + Port: port, + Name: strings.TrimLeft(u.Path, "/"), + User: u.User.Username(), + Pass: pass, + Ssl: "disable", + }, + } + + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + + const ( + first = `{"acls": [{"action": "accept", "src": ["*"], "dst": ["*:*"]}]}` + second = `{"acls": [{"action": "accept", "src": ["*"], "dst": ["*:22"]}]}` + ) + + _, _, err = s.UpdatePolicy([]byte(second)) + require.NoError(t, err) + + want, _ := s.Filter() + + var creates atomic.Int32 + + written := make(chan struct{}) + + require.NoError(t, s.db.DB.Callback().Create().After("gorm:create"). + Register("write_while_storing", func(tx *gorm.DB) { + if tx.Statement.Table != "policies" || creates.Add(1) != 1 { + return + } + + go func() { + defer close(written) + + _, _, err := s.UpdatePolicy([]byte(second)) + assert.NoError(t, err) + }() + + // Give the second write time to finish, unless the first + // write blocks it. + select { + case <-written: + case <-time.After(200 * time.Millisecond): + } + })) + t.Cleanup(func() { + _ = s.db.DB.Callback().Create().Remove("write_while_storing") + }) + + _, _, err = s.UpdatePolicy([]byte(first)) + require.NoError(t, err) + <-written + + var rows int64 + require.NoError(t, + s.db.DB.Unscoped().Model(&types.Policy{}).Count(&rows).Error) + assert.Equal(t, int64(1), rows, "only the current policy may be stored") + + stored, err := s.GetPolicy() + require.NoError(t, err) + assert.JSONEq(t, second, stored.Data, "the write that ran last is stored") + + got, _ := s.Filter() + assert.Equal(t, want, got, "the policy in force must be the one stored") +} + // TestReloadPolicyEmpty: with no policy configured a reload has nothing to // install; an emptied policy file is refused. Both keep the live policy. func TestReloadPolicyEmpty(t *testing.T) { diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index 01d9ef3c1..9b8a1d5ea 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -167,8 +167,8 @@ type State struct { // polMan handles policy evaluation and management polMan policy.PolicyManager - // policyMu makes loading the stored policy and installing it one step, - // so of two overlapping reloads the older document cannot land last. + // policyMu makes storing, loading and installing a policy one step, so + // policies land in the database in the order they are installed. policyMu sync.Mutex // authCache holds any pending authentication requests from either auth @@ -417,6 +417,11 @@ func (s *State) ReloadPolicy() ([]change.Change, error) { s.policyMu.Lock() defer s.policyMu.Unlock() + return s.reloadPolicyLocked() +} + +// reloadPolicyLocked is [State.ReloadPolicy] for callers holding policyMu. +func (s *State) reloadPolicyLocked() ([]change.Change, error) { pol, err := hsdb.PolicyBytes(s.db.DB, s.cfg) if err != nil { return nil, fmt.Errorf("loading policy: %w", err) @@ -1521,6 +1526,11 @@ func (s *State) UpdatePolicy(pol []byte) (*types.Policy, []change.Change, error) return nil, nil, err } + // Held from store to install: on PostgreSQL a concurrent writer could + // otherwise commit an older row after a newer one pruned the rest. + s.policyMu.Lock() + defer s.policyMu.Unlock() + stored, err := s.db.SetPolicy(string(pol)) if err != nil { return nil, nil, fmt.Errorf("storing policy: %w", err) @@ -1528,7 +1538,7 @@ func (s *State) UpdatePolicy(pol []byte) (*types.Policy, []change.Change, error) // Reload even when content is unchanged: routes manually disabled before // may now qualify for auto-approval, so they must be re-evaluated. - cs, err := s.ReloadPolicy() + cs, err := s.reloadPolicyLocked() return stored, cs, err }