Import libs 4 (#758)

3 files changed, 59 insertions, 7 deletions

diff --git a/library/cpp/ipmath/ipmath.cpp b/library/cpp/ipmath/ipmath.cpp
index 53f19dbb017..b2dffcfcfdb 100644
--- a/library/cpp/ipmath/ipmath.cpp
+++ b/library/cpp/ipmath/ipmath.cpp
@@ -139,7 +139,7 @@ TIpAddressRange::TIpAddressRangeBuilder& TIpAddressRange::TIpAddressRangeBuilder
     return *this;
 }
 
-TIpAddressRange::TIpAddressRangeBuilder& TIpAddressRange::TIpAddressRangeBuilder::WithPrefix(ui8 len) {
+TIpAddressRange::TIpAddressRangeBuilder& TIpAddressRange::TIpAddressRangeBuilder::WithPrefixImpl(ui8 len, bool checkLowerBound) {
     Y_ENSURE_EX(IsValid(Start_), TInvalidIpRangeException() << "Start value must be set before prefix");
     const auto type = Start_.Type();
     const auto maxLen = MaxPrefixLenForType(type);
@@ -147,14 +147,25 @@ TIpAddressRange::TIpAddressRangeBuilder& TIpAddressRange::TIpAddressRangeBuilder
         << maxLen << ", but requested " << (ui32)len);
 
     const auto lowerBound = LowerBoundForPrefix(Start_, len);
-    Y_ENSURE_EX(Start_ == lowerBound, TInvalidIpRangeException() << "Cannot create IP range from start address "
-        << Start_ << " with prefix length " << (ui32)len);
+    if (checkLowerBound) {
+        Y_ENSURE_EX(Start_ == lowerBound, TInvalidIpRangeException() << "Cannot create IP range from start address "
+            << Start_ << " with prefix length " << (ui32)len);
+    }
 
+    Start_ = lowerBound;
     End_ = UpperBoundForPrefix(Start_, len);
 
     return *this;
 }
 
+TIpAddressRange::TIpAddressRangeBuilder& TIpAddressRange::TIpAddressRangeBuilder::WithPrefix(ui8 len) {
+    return WithPrefixImpl(len, true);
+}
+
+TIpAddressRange::TIpAddressRangeBuilder& TIpAddressRange::TIpAddressRangeBuilder::WithMaskedPrefix(ui8 len) {
+    return WithPrefixImpl(len, false);
+}
+
 void TIpAddressRange::Init(TIpv6Address from, TIpv6Address to) {
     Start_ = from;
     End_ = to;
@@ -230,7 +241,7 @@ TIpAddressRange TIpAddressRange::FromCidrString(const TStringBuf str) {
     ythrow TInvalidIpRangeException() << "Cannot parse " << str << " as a CIDR string";
 }
 
-TMaybe<TIpAddressRange> TIpAddressRange::TryFromCidrString(const TStringBuf str) {
+TMaybe<TIpAddressRange> TIpAddressRange::TryFromCidrStringImpl(const TStringBuf str, bool compact) {
     auto idx = str.rfind('/');
     if (idx == TStringBuf::npos) {
         return Nothing();
@@ -246,8 +257,25 @@ TMaybe<TIpAddressRange> TIpAddressRange::TryFromCidrString(const TStringBuf str)
         return Nothing();
     }
 
-    return TIpAddressRange::From(address)
-        .WithPrefix(prefixLen);
+    return compact ?
+        TIpAddressRange::From(address).WithMaskedPrefix(prefixLen) :
+        TIpAddressRange::From(address).WithPrefix(prefixLen);
+}
+
+TMaybe<TIpAddressRange> TIpAddressRange::TryFromCidrString(const TStringBuf str) {
+    return TryFromCidrStringImpl(str, false);
+}
+
+TIpAddressRange TIpAddressRange::FromCompactString(const TStringBuf str) {
+    if (auto result = TryFromCompactString(str)) {
+        return *result;
+    }
+
+    ythrow TInvalidIpRangeException() << "Cannot parse " << str << " as a CIDR string";
+}
+
+TMaybe<TIpAddressRange> TIpAddressRange::TryFromCompactString(const TStringBuf str) {
+    return TryFromCidrStringImpl(str, true);
 }
 
 TIpAddressRange TIpAddressRange::FromRangeString(const TStringBuf str) {
diff --git a/library/cpp/ipmath/ipmath.h b/library/cpp/ipmath/ipmath.h
index f1c15e341ee..8c24bd471cd 100644
--- a/library/cpp/ipmath/ipmath.h
+++ b/library/cpp/ipmath/ipmath.h
@@ -26,7 +26,7 @@ public:
     static TIpAddressRangeBuilder From(const TStringBuf from);
 
     /**
-     * Parses a string tormatted in Classless Iter-Domain Routing (CIDR) notation.
+     * Parses a string formatted in Classless Inter-Domain Routing (CIDR) notation.
      * @param str a CIDR-formatted string, e.g. "192.168.0.0/16"
      * @return a new TIpAddressRange
      * @throws TInvalidIpRangeException if the string cannot be parsed.
@@ -35,6 +35,15 @@ public:
     static TMaybe<TIpAddressRange> TryFromCidrString(const TStringBuf str);
 
     /**
+     * Parses a string formatted in compact Classless Inter-Domain Routing (CIDR) notation with node address.
+     * @param str a CIDR-formatted string with node address, e.g. "192.168.1.24/16"
+     * @return a new TIpAddressRange
+     * @throws TInvalidIpRangeException if the string cannot be parsed.
+     */
+    static TIpAddressRange FromCompactString(const TStringBuf str);
+    static TMaybe<TIpAddressRange> TryFromCompactString(const TStringBuf str);
+
+    /**
      * Parses a string formatted as two dash-separated addresses.
      * @param str a CIDR-formatted string, e.g. "192.168.0.0-192.168.0.2"
      * @return a new TIpAddressRange
@@ -99,6 +108,8 @@ public:
 private:
     void Init(TIpv6Address, TIpv6Address);
 
+    static TMaybe<TIpAddressRange> TryFromCidrStringImpl(const TStringBuf str, bool compact);
+
     TIpv6Address Start_;
     TIpv6Address End_;
 };
@@ -129,8 +140,11 @@ public:
     TIpAddressRangeBuilder& To(TIpv6Address);
 
     TIpAddressRangeBuilder& WithPrefix(ui8 len);
+    TIpAddressRangeBuilder& WithMaskedPrefix(ui8 len);
 
 private:
+    TIpAddressRangeBuilder& WithPrefixImpl(ui8 len, bool checkLowerBound);
+
     TIpv6Address Start_;
     TIpv6Address End_;
 };
diff --git a/library/cpp/ipmath/ipmath_ut.cpp b/library/cpp/ipmath/ipmath_ut.cpp
index 179f2325034..8179c1bd672 100644
--- a/library/cpp/ipmath/ipmath_ut.cpp
+++ b/library/cpp/ipmath/ipmath_ut.cpp
@@ -35,6 +35,7 @@ public:
     UNIT_TEST(IpRangeFromIpv6);
     UNIT_TEST(FullIpRange);
     UNIT_TEST(IpRangeFromCidr);
+    UNIT_TEST(IpRangeFromCompact);
     UNIT_TEST(IpRangeFromIpv4Builder);
     UNIT_TEST(IpRangeFromInvalidIpv4);
     UNIT_TEST(IpRangeFromInvalidIpv6);
@@ -145,6 +146,15 @@ public:
         ASSERT_THROW(TIpAddressRange::FromCidrString("::/150"), TInvalidIpRangeException);
     }
 
+    void IpRangeFromCompact() {
+        // FromCidrString disallows node addresses
+        EXPECT_THROW(TIpAddressRange::FromCidrString("10.10.36.12/12"), TInvalidIpRangeException);
+
+        // FromCompactString allows to use node address instead of network address (suffix of zeroes is not required)
+        ASSERT_THAT(TIpAddressRange::FromCompactString("10.10.36.12/12"), Eq(TIpAddressRange::FromCidrString("10.0.0.0/12")));
+        ASSERT_THAT(TIpAddressRange::FromCompactString("abcd:ef01:2345::/24"), Eq(TIpAddressRange::FromCidrString("abcd:ef00::/24")));
+    }
+
     void RangeFromRangeString() {
         {
             auto range = TIpAddressRange::FromRangeString("10.0.0.0-10.0.0.3");