From 3b81f3e948ffc42ff53275b52e028ae4410f4494 Mon Sep 17 00:00:00 2001 From: Anthony Roy Date: Sun, 14 Jun 2026 19:30:05 -0700 Subject: [PATCH] [NetKAT] Make the PacketSetManager constructor private. PiperOrigin-RevId: 932174966 --- netkat/BUILD.bazel | 35 ++++++++++--------------- netkat/manager_handle_pattern.md | 29 ++++++++++++++------- netkat/packet_set.cc | 3 +++ netkat/packet_set.h | 19 ++++++++++---- netkat/packet_set_benchmark.cc | 13 +++++++--- netkat/packet_set_test.cc | 5 ++-- netkat/packet_set_test_runner.cc | 4 ++- netkat/packet_transformer.cc | 34 ++++++++++++++++++++---- netkat/packet_transformer.h | 12 +++++---- netkat/packet_transformer_test.cc | 43 ++++++++++++++++++++++++------- 10 files changed, 133 insertions(+), 64 deletions(-) diff --git a/netkat/BUILD.bazel b/netkat/BUILD.bazel index 9121fef..2d634f6 100644 --- a/netkat/BUILD.bazel +++ b/netkat/BUILD.bazel @@ -119,26 +119,9 @@ cc_test( ], ) -cc_library( +alias( name = "packet_set", - srcs = ["packet_set.cc"], - hdrs = ["packet_set.h"], - deps = [ - ":netkat_cc_proto", - ":packet", - ":packet_field", - ":paged_stable_vector", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:fixed_array", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:str_format", - "@com_google_gutil//gutil:status", - ], + actual = ":packet_transformer", ) cc_test( @@ -147,6 +130,7 @@ cc_test( shard_count = 8, deps = [ ":evaluator", + ":gtest_utils", ":netkat_proto_constructors", ":packet", ":packet_set", @@ -229,6 +213,7 @@ cc_test( shard_count = 8, deps = [ ":evaluator", + ":gtest_utils", ":netkat_cc_proto", ":netkat_proto_constructors", ":packet", @@ -362,13 +347,18 @@ cc_test( cc_library( name = "packet_transformer", - srcs = ["packet_transformer.cc"], - hdrs = ["packet_transformer.h"], + srcs = [ + "packet_set.cc", + "packet_transformer.cc", + ], + hdrs = [ + "packet_set.h", + "packet_transformer.h", + ], deps = [ ":netkat_cc_proto", ":packet", ":packet_field", - ":packet_set", ":paged_stable_vector", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:btree", @@ -403,6 +393,7 @@ cc_test( shard_count = 5, deps = [ ":evaluator", + ":gtest_utils", ":netkat_cc_proto", ":netkat_proto_constructors", ":packet", diff --git a/netkat/manager_handle_pattern.md b/netkat/manager_handle_pattern.md index e2c81b4..98b8437 100644 --- a/netkat/manager_handle_pattern.md +++ b/netkat/manager_handle_pattern.md @@ -26,19 +26,28 @@ combining them, or inspecting the underlying sets, one must call methods on the manager class, which acts as an arena allocator that owns all memory associated with the handles. -For example: ``` // We need a manager to construct handles. PacketSetManager -manager; PacketSetHandle a = manager.EmptySet() PacketSetHandle b = -manager.Match("src_mac", 0xFF'FF'FF'FF); +For example: -// Handles can be compared and hashed without the help of the manager, // but -that's about it. CHECK(a != b); absl::flat_hash_map ab_set{a, -b}; +``` +// We need a manager to construct handles. +PacketSetManager manager; +PacketSetHandle a = manager.EmptySet() +PacketSetHandle b = manager.Match("src_mac", 0xFF'FF'FF'FF); + +// Handles can be compared and hashed without the help of the manager, +// but that's about it. +CHECK(a != b); +absl::flat_hash_map ab_set{a, b}; // To do interesting things with the handles, we need the manager. -PacketSetHandle c = manager.And(a, b); // The set union of `a` and `b`. -PacketSetHandle not_c = manager.Not(c); // The set complement of `c`. if -(manager.Contains(c, packet)) { CHECK(!manager.Contains(not_c, packet)); } else -{ CHECK(manager.Contains(not_c, packet)); } ``` +PacketSetHandle c = manager.And(a, b); // The set union of `a` and `b`. +PacketSetHandle not_c = manager.Not(c); // The set complement of `c`. +if (manager.Contains(c, packet)) { + CHECK(!manager.Contains(not_c, packet)); +} else { + CHECK(manager.Contains(not_c, packet)); +} +``` ## Motivation for Using the Pattern diff --git a/netkat/packet_set.cc b/netkat/packet_set.cc index 2c7bd1c..d994f17 100644 --- a/netkat/packet_set.cc +++ b/netkat/packet_set.cc @@ -63,6 +63,9 @@ std::string PacketSetHandle::ToString() const { } } +PacketSetManager::PacketSetManager(PacketTransformerManager& transformer) + : transformer_(&transformer) {} + PacketSetHandle PacketSetManager::EmptySet() const { return PacketSetHandle(SentinelNodeIndex::kEmptySet); } diff --git a/netkat/packet_set.h b/netkat/packet_set.h index f11088f..cae37da 100644 --- a/netkat/packet_set.h +++ b/netkat/packet_set.h @@ -131,18 +131,20 @@ static_assert(sizeof(PacketSetHandle) <= 4); // CAUTION: Using a `PacketSetHandle` returned by one `PacketSetManager` // object with a different manager is undefined behavior. // +// This class is not constructible or movable publicly. It must be managed +// by a `PacketTransformerManager` (which owns a `PacketSetManager` by value). +// This restriction ensures that the `PacketSetManager` always has a valid +// `PacketTransformerManager` context, which is required to compile `Pull` +// operations (as they compile down to policies). +// // TODO(b/398303840): Persistent use of an `PacketSetManager` object can // incur unbounded memory growth. Consider adding some garbage collection // mechanism. class PacketSetManager { public: - PacketSetManager() = default; - - // The class is move-only: not copyable, but movable. + // The class is move-only: not copyable. PacketSetManager(const PacketSetManager&) = delete; PacketSetManager& operator=(const PacketSetManager&) = delete; - PacketSetManager(PacketSetManager&&) = default; - PacketSetManager& operator=(PacketSetManager&&) = default; // Returns true iff this packet set represents the empty set of packets. bool IsEmptySet(PacketSetHandle packet_set) const; @@ -159,6 +161,7 @@ class PacketSetManager { // Compiles the given `PredicateProto` into a `PacketSetHandle` that // represents the set of packets satisfying the predicate. + PacketSetHandle Compile(const PredicateProto& pred); // The packet set representing the empty set of packets. @@ -372,6 +375,12 @@ class PacketSetManager { // INVARIANT: All `DecisionNode` fields are interned by this manager. PacketFieldManager field_manager_; + explicit PacketSetManager(class PacketTransformerManager& transformer); + PacketSetManager(PacketSetManager&&) = default; + PacketSetManager& operator=(PacketSetManager&&) = default; + + class PacketTransformerManager* transformer_ = nullptr; + // Allow `PacketTransformerManager` to access private methods. friend class PacketTransformerManager; friend class PacketTransformerManagerTestPeer; diff --git a/netkat/packet_set_benchmark.cc b/netkat/packet_set_benchmark.cc index 7de2fd0..5c9a6ad 100644 --- a/netkat/packet_set_benchmark.cc +++ b/netkat/packet_set_benchmark.cc @@ -19,6 +19,7 @@ #include "netkat/netkat.pb.h" #include "netkat/netkat_proto_constructors.h" #include "netkat/packet_set.h" +#include "netkat/packet_transformer.h" namespace netkat { // Create an arbitrary fixed policy with some relative complexity. In this @@ -56,7 +57,8 @@ void BM_FirstTimeCompileNonOverlappingPredicate(benchmark::State& state) { PredicateProto policy = OrProto(sub_policy1, sub_policy2); for (auto s : state) { - PacketSetManager manager; + PacketTransformerManager transformer; + PacketSetManager& manager = transformer.GetPacketSetManager(); PacketSetHandle handle = manager.Compile(policy); benchmark::DoNotOptimize(handle); } @@ -75,7 +77,8 @@ void BM_ReCompileNonOverlappingPredicate(benchmark::State& state) { CreateFixedArbitraryPredicateProto(/*id_suffix=*/4)); PredicateProto policy = OrProto(sub_policy1, sub_policy2); - PacketSetManager manager; + PacketTransformerManager transformer; + PacketSetManager& manager = transformer.GetPacketSetManager(); PacketSetHandle handle = manager.Compile(policy); for (auto s : state) { handle = manager.Compile(policy); @@ -94,7 +97,8 @@ void BM_FirstTimeCompileOverlappingPredicate(benchmark::State& state) { PredicateProto policy = OrProto(sub_policy1, sub_policy2); for (auto s : state) { - PacketSetManager manager; + PacketTransformerManager transformer; + PacketSetManager& manager = transformer.GetPacketSetManager(); PacketSetHandle handle = manager.Compile(policy); benchmark::DoNotOptimize(handle); } @@ -111,7 +115,8 @@ void BM_ReCompileOverlappingPredicate(benchmark::State& state) { CreateFixedArbitraryPredicateProto()); PredicateProto policy = OrProto(sub_policy1, sub_policy2); - PacketSetManager manager; + PacketTransformerManager transformer; + PacketSetManager& manager = transformer.GetPacketSetManager(); PacketSetHandle handle = manager.Compile(policy); for (auto s : state) { handle = manager.Compile(policy); diff --git a/netkat/packet_set_test.cc b/netkat/packet_set_test.cc index 40731f6..d85463a 100644 --- a/netkat/packet_set_test.cc +++ b/netkat/packet_set_test.cc @@ -32,6 +32,7 @@ #include "netkat/evaluator.h" #include "netkat/netkat_proto_constructors.h" #include "netkat/packet.h" +#include "netkat/packet_transformer.h" #include "re2/re2.h" namespace netkat { @@ -40,8 +41,8 @@ namespace netkat { // test cases. This also enables better pretty printing for debugging, see // `PrintTo`. PacketSetManager& Manager() { - static absl::NoDestructor manager; - return *manager; + static absl::NoDestructor transformer; + return transformer->GetPacketSetManager(); } // The default `PacketSetHandle` pretty printer sucks! It does not have access diff --git a/netkat/packet_set_test_runner.cc b/netkat/packet_set_test_runner.cc index 30fd670..c7dfeb9 100644 --- a/netkat/packet_set_test_runner.cc +++ b/netkat/packet_set_test_runner.cc @@ -22,6 +22,7 @@ #include "netkat/netkat_proto_constructors.h" #include "netkat/packet_set.h" +#include "netkat/packet_transformer.h" namespace netkat { namespace { @@ -84,7 +85,8 @@ std::vector TestCases() { void main() { // This test needs a deterministic field interning order, and thus must start // from a fresh manager. - PacketSetManager manager; + PacketTransformerManager transformer; + PacketSetManager& manager = transformer.GetPacketSetManager(); for (const TestCase& test_case : TestCases()) { netkat::PacketSetHandle packet_set = manager.Compile(test_case.predicate); std::cout << kBanner << "Test case: " << test_case.description << std::endl diff --git a/netkat/packet_transformer.cc b/netkat/packet_transformer.cc index a125c8b..b32dc78 100644 --- a/netkat/packet_transformer.cc +++ b/netkat/packet_transformer.cc @@ -71,6 +71,30 @@ std::string PacketTransformerHandle::ToString() const { } } +PacketTransformerManager::PacketTransformerManager() + : packet_set_manager_(*this) {} + +PacketTransformerManager::PacketTransformerManager( + PacketTransformerManager&& other) + : nodes_(std::move(other.nodes_)), + transformer_by_node_(std::move(other.transformer_by_node_)), + transformer_by_hash_(std::move(other.transformer_by_hash_)), + packet_set_manager_(std::move(other.packet_set_manager_)) { + packet_set_manager_.transformer_ = this; +} + +PacketTransformerManager& PacketTransformerManager::operator=( + PacketTransformerManager&& other) { + if (this != &other) { + nodes_ = std::move(other.nodes_); + transformer_by_node_ = std::move(other.transformer_by_node_); + transformer_by_hash_ = std::move(other.transformer_by_hash_); + packet_set_manager_ = std::move(other.packet_set_manager_); + packet_set_manager_.transformer_ = this; + } + return *this; +} + const PacketTransformerManager::DecisionNode& PacketTransformerManager::GetNodeOrDie( PacketTransformerHandle transformer) const { @@ -771,8 +795,8 @@ PacketTransformerHandle PacketTransformerManager::Iterate( PacketSetHandle PacketTransformerManager::GetAllPossibleOutputPackets( PacketTransformerHandle transformer) { - if (IsAccept(transformer)) return PacketSetManager().FullSet(); - if (IsDeny(transformer)) return PacketSetManager().EmptySet(); + if (IsAccept(transformer)) return packet_set_manager_.FullSet(); + if (IsDeny(transformer)) return packet_set_manager_.EmptySet(); const DecisionNode& node = GetNodeOrDie(transformer); PacketSetHandle default_output = @@ -821,7 +845,7 @@ PacketSetHandle PacketTransformerManager::GetAllPossibleOutputPackets( // C.3 Push and Pull in KATch: A Fast Symbolic Verifier for NetKAT. for (auto& [match_value, unused] : node.modify_branch_by_field_match) { if (!branch_modify_values.contains(match_value)) { - add_to_output_by_field_value(match_value, PacketSetManager().EmptySet()); + add_to_output_by_field_value(match_value, packet_set_manager_.EmptySet()); } } @@ -862,8 +886,8 @@ PacketSetHandle PacketTransformerManager::Push( PacketSetHandle PacketTransformerManager::GetAllInputPacketsThatProduceAnyOutput( PacketTransformerHandle transformer) { - if (IsAccept(transformer)) return PacketSetManager().FullSet(); - if (IsDeny(transformer)) return PacketSetManager().EmptySet(); + if (IsAccept(transformer)) return packet_set_manager_.FullSet(); + if (IsDeny(transformer)) return packet_set_manager_.EmptySet(); const DecisionNode& node = GetNodeOrDie(transformer); diff --git a/netkat/packet_transformer.h b/netkat/packet_transformer.h index 4c9c90d..5a18fe9 100644 --- a/netkat/packet_transformer.h +++ b/netkat/packet_transformer.h @@ -138,23 +138,25 @@ static_assert(sizeof(PacketTransformerHandle) <= 4); // `PacketTransformerManager` object with a different manager is // undefined behavior. `PacketSetHandles` and `PacketTransformerHandles` // returned by this class are not invalidated on move. + class PacketTransformerManager { public: - PacketTransformerManager() = default; - explicit PacketTransformerManager(PacketSetManager&& manager) - : packet_set_manager_(std::move(manager)) {}; + PacketTransformerManager(); // The class is move-only: not copyable, but movable. // `PacketSetHandles` and `PacketTransformerHandles` returned by this class // are not invalidated on move. PacketTransformerManager(const PacketTransformerManager&) = delete; PacketTransformerManager& operator=(const PacketTransformerManager&) = delete; - PacketTransformerManager(PacketTransformerManager&&) = default; - PacketTransformerManager& operator=(PacketTransformerManager&&) = default; + PacketTransformerManager(PacketTransformerManager&& other); + PacketTransformerManager& operator=(PacketTransformerManager&& other); // Returns the `PacketSetManager` used by this object to compile // predicates. PacketSetManager& GetPacketSetManager() { return packet_set_manager_; } + const PacketSetManager& GetPacketSetManager() const { + return packet_set_manager_; + } // Returns true iff this transformer represents the Deny policy. bool IsDeny(PacketTransformerHandle transformer) const; diff --git a/netkat/packet_transformer_test.cc b/netkat/packet_transformer_test.cc index 0e7b0fc..1b3a92e 100644 --- a/netkat/packet_transformer_test.cc +++ b/netkat/packet_transformer_test.cc @@ -114,6 +114,35 @@ TEST(PacketTransformerManagerTest, AbslHashValueWorks) { EXPECT_EQ(set.size(), 2); } +TEST(PacketTransformerManagerTest, MoveConstructorPreservesState) { + PacketTransformerManager manager1; + PacketTransformerHandle h1 = manager1.Compile(ModificationProto("f", 1)); + PacketTransformerHandle h2 = manager1.Compile(ModificationProto("f", 2)); + + PacketTransformerManager manager2 = std::move(manager1); + + EXPECT_THAT(manager2.ToString(h1), StartsWith("PacketTransformerHandle")); + EXPECT_THAT(manager2.ToString(h2), StartsWith("PacketTransformerHandle")); + + EXPECT_EQ(manager2.Compile(ModificationProto("f", 1)), h1); + EXPECT_EQ(manager2.Compile(ModificationProto("f", 2)), h2); +} + +TEST(PacketTransformerManagerTest, MoveAssignmentPreservesState) { + PacketTransformerManager manager1; + PacketTransformerHandle h1 = manager1.Compile(ModificationProto("f", 1)); + PacketTransformerHandle h2 = manager1.Compile(ModificationProto("f", 2)); + + PacketTransformerManager manager2; + manager2 = std::move(manager1); + + EXPECT_THAT(manager2.ToString(h1), StartsWith("PacketTransformerHandle")); + EXPECT_THAT(manager2.ToString(h2), StartsWith("PacketTransformerHandle")); + + EXPECT_EQ(manager2.Compile(ModificationProto("f", 1)), h1); + EXPECT_EQ(manager2.Compile(ModificationProto("f", 2)), h2); +} + TEST(PacketTransformerManagerTest, EmptyPolicyCompilesToDeny) { EXPECT_TRUE(Manager().IsDeny(Manager().Compile(PolicyProto()))); } @@ -128,13 +157,6 @@ void CompileIsSameAsOfCompiledPacketSetHandle(PredicateProto predicate) { PacketSetHandle set_1 = Manager().GetPacketSetManager().Compile(predicate); EXPECT_EQ(Manager().Compile(FilterProto(predicate)), Manager().FromPacketSetHandle(set_1)); - - // Using a newly constructed PacketSetManager. - PacketSetManager packet_set_manager; - PacketSetHandle set_2 = packet_set_manager.Compile(predicate); - PacketTransformerManager manager(std::move(packet_set_manager)); - EXPECT_EQ(manager.Compile(FilterProto(predicate)), - manager.FromPacketSetHandle(set_2)); } FUZZ_TEST(PacketTransformerManagerTest, CompileIsSameAsOfCompiledPacketSetHandle); @@ -942,9 +964,9 @@ class PacketTransformerManagerTestPeer { PacketSetHandle GetAllPossibleOutputPacketsReferenceImplementation( PacketTransformerHandle transformer) { if (packet_transformer_manager_->IsAccept(transformer)) - return PacketSetManager().FullSet(); + return packet_transformer_manager_->GetPacketSetManager().FullSet(); if (packet_transformer_manager_->IsDeny(transformer)) - return PacketSetManager().EmptySet(); + return packet_transformer_manager_->GetPacketSetManager().EmptySet(); const PacketTransformerManager::DecisionNode& node = packet_transformer_manager_->GetNodeOrDie(transformer); const std::string field = packet_transformer_manager_->GetPacketSetManager() @@ -994,7 +1016,8 @@ class PacketTransformerManagerTestPeer { // 1. input.field != match_value for all explicit match branches, thus // output.field != match_value for all explicit match branches by (0) // 2. output.field != modify_value for all default-modify branches - PacketSetHandle fallthrough_output = PacketSetManager().FullSet(); + PacketSetHandle fallthrough_output = + packet_transformer_manager_->GetPacketSetManager().FullSet(); for (const auto& [match_value, unused] : node.modify_branch_by_field_match) { fallthrough_output =