YT-21684: Fix TIP6Network serialization.

c23d66825b6b85d3f287e56fc70e867f88f99d47

3 files changed, 101 insertions, 20 deletions

diff --git a/yt/yt/core/net/address.cpp b/yt/yt/core/net/address.cpp
index 97a9112816d..8787ec80ffd 100644
--- a/yt/yt/core/net/address.cpp
+++ b/yt/yt/core/net/address.cpp
@@ -803,10 +803,35 @@ void Serialize(const TIP6Address& value, IYsonConsumer* consumer)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+int GetMaskSize(const TIP6Address& mask)
+{
+    int size = 0;
+    const auto* parts = mask.GetRawDWords();
+    for (size_t partIndex = 0; partIndex < 4; ++partIndex) {
+        size += __builtin_popcount(parts[partIndex]);
+    }
+    return size;
+}
+
 TIP6Network::TIP6Network(const TIP6Address& network, const TIP6Address& mask)
     : Network_(network)
     , Mask_(mask)
-{ }
+{
+    // We allow only masks that can be written by using prefix notation,
+    // e.g. /64. In prefix notation no ones can go after zeros.
+    bool seenOne = false;
+    auto bytes = mask.GetRawBytes();
+    for (int i = 0; i < static_cast<int>(TIP6Address::ByteSize * 8); ++i) {
+        auto val = *(bytes + i / 8) & (1 << (i % 8));
+        if (val) {
+            seenOne = true;
+        } else {
+            if (seenOne) {
+                THROW_ERROR_EXCEPTION("Invalid network mask %Qv", ToString(mask));
+            }
+        }
+    }
+}
 
 const TIP6Address& TIP6Network::GetAddress() const
 {
@@ -820,12 +845,12 @@ const TIP6Address& TIP6Network::GetMask() const
 
 int TIP6Network::GetMaskSize() const
 {
-    int size = 0;
-    const auto* parts = Mask_.GetRawDWords();
-    for (size_t partIndex = 0; partIndex < 4; ++partIndex) {
-        size += __builtin_popcount(parts[partIndex]);
-    }
-    return size;
+    return ::NYT::NNet::GetMaskSize(Mask_);
+}
+
+std::optional<ui32> TIP6Network::GetProjectId() const
+{
+    return ProjectId_;
 }
 
 bool TIP6Network::Contains(const TIP6Address& address) const
@@ -847,8 +872,7 @@ TIP6Network TIP6Network::FromString(TStringBuf str)
 bool TIP6Network::FromString(TStringBuf str, TIP6Network* network)
 {
     auto buf = str;
-    std::optional<ui32> projectId = std::nullopt;
-    if (!ParseProjectId(&buf, &projectId)) {
+    if (!ParseProjectId(&buf, &network->ProjectId_)) {
         return false;
     }
 
@@ -861,11 +885,7 @@ bool TIP6Network::FromString(TStringBuf str, TIP6Network* network)
         return false;
     }
 
-    if (projectId) {
-        network->Network_.GetRawWords()[2] = *projectId;
-        network->Network_.GetRawWords()[3] = *projectId >> 16;
-    }
-
+    // Set mask based on mask size.
     network->Mask_ = TIP6Address();
     auto bytes = network->Mask_.GetRawBytes();
     for (int i = 0; i < static_cast<int>(TIP6Address::ByteSize * 8); ++i) {
@@ -876,9 +896,16 @@ bool TIP6Network::FromString(TStringBuf str, TIP6Network* network)
         }
     }
 
-    if (projectId) {
+    if (network->ProjectId_) {
         static_assert(TIP6Address::ByteSize == 16);
-        network->Mask_.GetRawDWords()[1] = 0xffffffff;
+
+        // Set 64-95 bits of address to ::::[project_id_low:project_id_high]::
+        network->Network_.GetRawWords()[2] = *network->ProjectId_;
+        network->Network_.GetRawWords()[3] = *network->ProjectId_ >> 16;
+
+        // Set 64-95 bits of mask to ::::[ffff:ffff]::
+        network->Mask_.GetRawWords()[2] = 0xffff;
+        network->Mask_.GetRawWords()[3] = 0xffff;
     }
 
     return true;
@@ -886,9 +913,30 @@ bool TIP6Network::FromString(TStringBuf str, TIP6Network* network)
 
 void FormatValue(TStringBuilderBase* builder, const TIP6Network& network, TStringBuf /*spec*/)
 {
-    builder->AppendFormat("%v/%v",
-        network.GetAddress(),
-        network.GetMaskSize());
+    auto projectId = network.GetProjectId();
+    if (projectId) {
+        // The network has been created from string in
+        // project id notation. Save it just the way it came.
+
+        // Recover original network address.
+        auto address = network.GetAddress();
+        address.GetRawWords()[2] = 0;
+        address.GetRawWords()[3] = 0;
+
+        // Recover original mask.
+        auto mask = network.GetMask();
+        mask.GetRawWords()[2] = 0;
+        mask.GetRawWords()[3] = 0;
+
+        builder->AppendFormat("%x@%v/%v",
+            *projectId,
+            address,
+            GetMaskSize(mask));
+    } else {
+        builder->AppendFormat("%v/%v",
+            network.GetAddress(),
+            network.GetMaskSize());
+    }
 }
 
 TString ToString(const TIP6Network& network)
diff --git a/yt/yt/core/net/address.h b/yt/yt/core/net/address.h
index b9352f8c3f8..37626d5d234 100644
--- a/yt/yt/core/net/address.h
+++ b/yt/yt/core/net/address.h
@@ -157,9 +157,13 @@ public:
     const TIP6Address& GetMask() const;
     int GetMaskSize() const;
 
+    //! Get project id as extracted by FromString().
+    std::optional<ui32> GetProjectId() const;
+
 private:
     TIP6Address Network_;
     TIP6Address Mask_;
+    std::optional<ui32> ProjectId_;
 };
 
 void FormatValue(TStringBuilderBase* builder, const TIP6Network& network, TStringBuf spec);
diff --git a/yt/yt/core/net/unittests/network_address_ut.cpp b/yt/yt/core/net/unittests/network_address_ut.cpp
index d93fba8d88e..0c2e4635162 100644
--- a/yt/yt/core/net/unittests/network_address_ut.cpp
+++ b/yt/yt/core/net/unittests/network_address_ut.cpp
@@ -207,6 +207,23 @@ TEST(TIP6AddressTest, CanonicalText)
     }
 }
 
+TEST(TIP6AddressTest, Constructors)
+{
+    {
+        auto address = TIP6Address::FromString("0:0:0:0:0:0:0:3");
+        auto mask = TIP6Address::FromString("ffff:ffff:ffff:ffff:ffff:ffff:ffff:0000");
+        TIP6Network network(address, mask);
+        EXPECT_EQ(network.GetAddress(), address);
+        EXPECT_EQ(network.GetMask(), mask);
+    }
+
+    {
+        auto address = TIP6Address::FromString("0:0:0:0:0:0:0:3");
+        auto mask = TIP6Address::FromString("ffff:ffff:ffff:ffff:ffff:ffff:0000:ffff");
+        EXPECT_THROW(TIP6Network(address, mask), TErrorException);
+    }
+}
+
 TEST(TIP6AddressTest, NetworkMask)
 {
     using TTestCase = std::tuple<const char*, std::array<ui16, 8>, int>;
@@ -222,6 +239,12 @@ TEST(TIP6AddressTest, NetworkMask)
         EXPECT_EQ(AddressToWords(network.GetMask()), std::get<1>(testCase));
         EXPECT_EQ(network.GetMaskSize(), std::get<2>(testCase));
         EXPECT_EQ(ToString(network), std::get<0>(testCase));
+
+        auto networkAsString = ToString(network);
+        auto networkFromSerializedString = TIP6Network::FromString(networkAsString);
+        EXPECT_EQ(AddressToWords(networkFromSerializedString.GetMask()), std::get<1>(testCase));
+        EXPECT_EQ(networkFromSerializedString.GetMaskSize(), std::get<2>(testCase));
+        EXPECT_EQ(ToString(networkFromSerializedString), std::get<0>(testCase));
     }
 
     EXPECT_THROW(TIP6Network::FromString("::/129"), TErrorException);
@@ -249,7 +272,13 @@ TEST(TIP6AddressTest, ProjectId)
         auto testAddress = TIP6Address::FromString(std::get<1>(testCase));
         EXPECT_EQ(AddressToWords(network.GetAddress()), std::get<2>(testCase));
         EXPECT_EQ(AddressToWords(network.GetMask()), std::get<3>(testCase));
-        EXPECT_TRUE(network.Contains(testAddress) );
+        EXPECT_TRUE(network.Contains(testAddress));
+
+        auto networkAsString = ToString(network);
+        auto networkFromSerializedString = TIP6Network::FromString(networkAsString);
+        EXPECT_EQ(AddressToWords(networkFromSerializedString.GetAddress()), std::get<2>(testCase));
+        EXPECT_EQ(AddressToWords(networkFromSerializedString.GetMask()), std::get<3>(testCase));
+        EXPECT_TRUE(networkFromSerializedString.Contains(testAddress));
     }
 
     EXPECT_THROW(TIP6Network::FromString("@1::1"), TErrorException);