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 534763245..976917723 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 @@ -1339,6 +1341,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 f0eaf5c27..aab3eadf0 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -182,8 +182,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 pending Web, OIDC and SSH authentication requests. @@ -472,6 +472,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) @@ -1576,6 +1581,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) @@ -1583,7 +1593,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 }