Merge pull request #1002 from ellemouton/sql24

[sql-24] firewalldb: thread contexts through for privacy mapper interfaces

Author: ellemouton

firewall/privacy_mapper.go

@@ -1073,9 +1073,11 @@ func newMockDB(t *testing.T, preloadRealToPseudo map[string]string,
 	db := mockDB{privDB: make(map[string]*mockPrivacyMapDB)}
 	sessDB := db.NewSessionDB(sessID)
 
-	_ = sessDB.Update(func(tx firewalldb.PrivacyMapTx) error {
+	_ = sessDB.Update(context.Background(), func(ctx context.Context,
+		tx firewalldb.PrivacyMapTx) error {
+
 		for r, p := range preloadRealToPseudo {
-			require.NoError(t, tx.NewPair(r, p))
+			require.NoError(t, tx.NewPair(ctx, r, p))
 		}
 		return nil
 	})
@@ -1107,25 +1109,29 @@ type mockPrivacyMapDB struct {
 	p2r map[string]string
 }
 
-func (m *mockPrivacyMapDB) Update(
-	f func(tx firewalldb.PrivacyMapTx) error) error {
+func (m *mockPrivacyMapDB) Update(ctx context.Context,
+	f func(ctx context.Context, tx firewalldb.PrivacyMapTx) error) error {
 
-	return f(m)
+	return f(ctx, m)
 }
 
-func (m *mockPrivacyMapDB) View(
-	f func(tx firewalldb.PrivacyMapTx) error) error {
+func (m *mockPrivacyMapDB) View(ctx context.Context,
+	f func(ctx context.Context, tx firewalldb.PrivacyMapTx) error) error {
 
-	return f(m)
+	return f(ctx, m)
 }
 
-func (m *mockPrivacyMapDB) NewPair(real, pseudo string) error {
+func (m *mockPrivacyMapDB) NewPair(_ context.Context, real,
+	pseudo string) error {
+
 	m.r2p[real] = pseudo
 	m.p2r[pseudo] = real
 	return nil
 }
 
-func (m *mockPrivacyMapDB) PseudoToReal(pseudo string) (string, error) {
+func (m *mockPrivacyMapDB) PseudoToReal(_ context.Context, pseudo string) (
+	string, error) {
+
 	r, ok := m.p2r[pseudo]
 	if !ok {
 		return "", firewalldb.ErrNoSuchKeyFound
@@ -1134,7 +1140,9 @@ func (m *mockPrivacyMapDB) PseudoToReal(pseudo string) (string, error) {
 	return r, nil
 }
 
-func (m *mockPrivacyMapDB) RealToPseudo(real string) (string, error) {
+func (m *mockPrivacyMapDB) RealToPseudo(_ context.Context, real string) (string,
+	error) {
+
 	p, ok := m.r2p[real]
 	if !ok {
 		return "", firewalldb.ErrNoSuchKeyFound
@@ -1143,8 +1151,8 @@ func (m *mockPrivacyMapDB) RealToPseudo(real string) (string, error) {
 	return p, nil
 }
 
-func (m *mockPrivacyMapDB) FetchAllPairs() (*firewalldb.PrivacyMapPairs,
-	error) {
+func (m *mockPrivacyMapDB) FetchAllPairs(_ context.Context) (
+	*firewalldb.PrivacyMapPairs, error) {
 
 	return firewalldb.NewPrivacyMapPairs(m.r2p), nil
 }

firewall/privacy_mapper_test.go

@@ -395,7 +395,7 @@ func (r *RuleEnforcer) initRule(ctx context.Context, reqID uint64, name string,
 		privMap := r.newPrivMap(session.GroupID)
 
 		ruleValues, err = ruleValues.PseudoToReal(
-			privMap, session.PrivacyFlags,
+			ctx, privMap, session.PrivacyFlags,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("could not prepare rule "+

firewall/rule_enforcer.go

@@ -0,0 +1,44 @@
+package firewalldb
+
+import (
+	"context"
+
+	"go.etcd.io/bbolt"
+)
+
+// kvdbExecutor is a concrete implementation of the DBExecutor interface that
+// uses a bbolt database as its backing store.
+type kvdbExecutor[T any] struct {
+	db     *bbolt.DB
+	wrapTx func(tx *bbolt.Tx) T
+}
+
+// Update opens a database read/write transaction and executes the function f
+// with the transaction passed as a parameter. After f exits, if f did not
+// error, the transaction is committed. Otherwise, if f did error, the
+// transaction is rolled back. If the rollback fails, the original error
+// returned by f is still returned. If the commit fails, the commit error is
+// returned.
+//
+// NOTE: this is part of the DBExecutor interface.
+func (e *kvdbExecutor[T]) Update(ctx context.Context,
+	fn func(ctx context.Context, tx T) error) error {
+
+	return e.db.Update(func(tx *bbolt.Tx) error {
+		return fn(ctx, e.wrapTx(tx))
+	})
+}
+
+// View opens a database read transaction and executes the function f with the
+// transaction passed as a parameter. After f exits, the transaction is rolled
+// back. If f errors, its error is returned, not a rollback error (if any
+// occur).
+//
+// NOTE: this is part of the DBExecutor interface.
+func (e *kvdbExecutor[T]) View(ctx context.Context,
+	fn func(ctx context.Context, tx T) error) error {
+
+	return e.db.View(func(tx *bbolt.Tx) error {
+		return fn(ctx, e.wrapTx(tx))
+	})
+}

firewalldb/kvdb_store.go

@@ -107,62 +107,28 @@ type RulesDB interface {
 func (db *DB) GetKVStores(rule string, groupID session.ID,
 	feature string) KVStores {
 
-	return &kvStores{
-		db:          db.DB,
-		ruleName:    rule,
-		groupID:     groupID,
-		featureName: feature,
+	return &kvdbExecutor[KVStoreTx]{
+		db: db.DB,
+		wrapTx: func(tx *bbolt.Tx) KVStoreTx {
+			return &kvStoreTx{
+				boltTx: tx,
+				kvStores: &kvStores{
+					ruleName:    rule,
+					groupID:     groupID,
+					featureName: feature,
+				},
+			}
+		},
 	}
 }
 
 // kvStores implements the rules.KVStores interface.
 type kvStores struct {
-	db          *bbolt.DB
 	ruleName    string
 	groupID     session.ID
 	featureName string
 }
 
-// Update opens a database read/write transaction and executes the function f
-// with the transaction passed as a parameter. After f exits, if f did not
-// error, the transaction is committed. Otherwise, if f did error, the
-// transaction is rolled back. If the rollback fails, the original error
-// returned by f is still returned. If the commit fails, the commit error is
-// returned.
-//
-// NOTE: this is part of the KVStores interface.
-func (s *kvStores) Update(ctx context.Context, fn func(ctx context.Context,
-	tx KVStoreTx) error) error {
-
-	return s.db.Update(func(tx *bbolt.Tx) error {
-		boltTx := &kvStoreTx{
-			boltTx:   tx,
-			kvStores: s,
-		}
-
-		return fn(ctx, boltTx)
-	})
-}
-
-// View opens a database read transaction and executes the function f with the
-// transaction passed as a parameter. After f exits, the transaction is rolled
-// back. If f errors, its error is returned, not a rollback error (if any
-// occur).
-//
-// NOTE: this is part of the KVStores interface.
-func (s *kvStores) View(ctx context.Context, fn func(ctx context.Context,
-	tx KVStoreTx) error) error {
-
-	return s.db.View(func(tx *bbolt.Tx) error {
-		boltTx := &kvStoreTx{
-			boltTx:   tx,
-			kvStores: s,
-		}
-
-		return fn(ctx, boltTx)
-	})
-}
-
 // getBucketFunc defines the signature of the bucket creation/fetching function
 // required by kvStoreTx. If create is true, then all the bucket (and all
 // buckets leading up to the bucket) should be created if they do not already