From aafc13d6786bd99cf1f18c3161b3febddc4669c9 Mon Sep 17 00:00:00 2001
From: Sara Golemon <sara.golemon@mongodb.com>
Date: Wed, 21 Jun 2017 14:38:53 +0000
Subject: [PATCH] SERVER-29169 ServerAddressRestriction

---
 src/mongo/db/auth/address_restriction.cpp     |  3 +++
 src/mongo/db/auth/address_restriction.h       | 19 +++++++++++++++++++
 .../db/auth/address_restriction_test.cpp      | 13 +++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/src/mongo/db/auth/address_restriction.cpp b/src/mongo/db/auth/address_restriction.cpp
index a92dd4f07bd..fb7d626c659 100644
--- a/src/mongo/db/auth/address_restriction.cpp
+++ b/src/mongo/db/auth/address_restriction.cpp
@@ -34,5 +34,8 @@ namespace address_restriction_detail {
 constexpr StringData ClientSource::label;
 constexpr StringData ClientSource::field;
 
+constexpr StringData ServerAddress::label;
+constexpr StringData ServerAddress::field;
+
 }  // address_restriction_detail
 }  // mongo
diff --git a/src/mongo/db/auth/address_restriction.h b/src/mongo/db/auth/address_restriction.h
index de504638cd8..4f96c4f21df 100644
--- a/src/mongo/db/auth/address_restriction.h
+++ b/src/mongo/db/auth/address_restriction.h
@@ -48,6 +48,13 @@ struct ClientSource {
         return environment.getClientSource();
     }
 };
+struct ServerAddress {
+    static constexpr auto label = "Server address "_sd;
+    static constexpr auto const field = "serverAddress"_sd;
+    static auto addr(const RestrictionEnvironment& environment) {
+        return environment.getServerAddress();
+    }
+};
 
 // Represents a restriction based on client or server address
 template <typename T>
@@ -139,6 +146,8 @@ private:
 
 using ClientSourceRestriction =
     address_restriction_detail::AddressRestriction<address_restriction_detail::ClientSource>;
+using ServerAddressRestriction =
+    address_restriction_detail::AddressRestriction<address_restriction_detail::ServerAddress>;
 
 template <>
 inline BSONObjBuilder& BSONObjBuilderValueStream::operator<<<ClientSourceRestriction>(
@@ -150,4 +159,14 @@ inline BSONObjBuilder& BSONObjBuilderValueStream::operator<<<ClientSourceRestric
     return *_builder;
 }
 
+template <>
+inline BSONObjBuilder& BSONObjBuilderValueStream::operator<<<ServerAddressRestriction>(
+    ServerAddressRestriction value) {
+    BSONObjBuilder b;
+    value.appendToBuilder(&b);
+    _builder->append(_fieldName, b.obj());
+    _fieldName = StringData();
+    return *_builder;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/auth/address_restriction_test.cpp b/src/mongo/db/auth/address_restriction_test.cpp
index e4e6186206c..19a6c3c780d 100644
--- a/src/mongo/db/auth/address_restriction_test.cpp
+++ b/src/mongo/db/auth/address_restriction_test.cpp
@@ -57,6 +57,15 @@ TEST(AddressRestrictionTest, toAndFromString) {
             expected << "{\"clientSource\": \"" << p.output << "\"}";
             ASSERT_EQUALS(actual.str(), expected.str());
         }
+
+        {
+            const ServerAddressRestriction sar(CIDR(p.input));
+            std::ostringstream actual;
+            actual << sar;
+            std::ostringstream expected;
+            expected << "{\"serverAddress\": \"" << p.output << "\"}";
+            ASSERT_EQUALS(actual.str(), expected.str());
+        }
     }
 }
 
@@ -160,6 +169,10 @@ TEST(AddressRestrictionTest, contains) {
         const ClientSourceRestriction csr(CIDR(p.range));
         ASSERT_EQ(csr.validate(rec).isOK(), p.valid);
         ASSERT_FALSE(csr.validate(res).isOK());
+
+        const ServerAddressRestriction sar(CIDR(p.range));
+        ASSERT_EQ(sar.validate(res).isOK(), p.valid);
+        ASSERT_FALSE(sar.validate(rec).isOK());
     }
 }