Merge "Code Health: Make OperationLimiter generic"

3 files changed, 57 insertions, 17 deletions

diff --git a/DnsProxyListener.cpp b/DnsProxyListener.cpp
index 716c8e60..55210842 100644
--- a/DnsProxyListener.cpp
+++ b/DnsProxyListener.cpp
@@ -81,6 +81,16 @@ constexpr int MAX_QUERIES_PER_UID = 256;
 
 android::netdutils::OperationLimiter<uid_t> queryLimiter(MAX_QUERIES_PER_UID);
 
+bool startQueryLimiter(uid_t uid) {
+    const int globalLimit =
+            android::net::Experiments::getInstance()->getFlag("max_queries_global", INT_MAX);
+    return queryLimiter.start(uid, globalLimit);
+}
+
+void endQueryLimiter(uid_t uid) {
+    queryLimiter.finish(uid);
+}
+
 void logArguments(int argc, char** argv) {
     if (!WOULD_LOG(VERBOSE)) return;
     for (int i = 0; i < argc; i++) {
@@ -788,14 +798,14 @@ void DnsProxyListener::GetAddrInfoHandler::doDns64Synthesis(int32_t* rv, addrinf
     if (ipv6WantedButNoData) {
         // If caller wants IPv6 answers but no data, try to query IPv4 answers for synthesis
         const uid_t uid = mClient->getUid();
-        if (queryLimiter.start(uid)) {
+        if (startQueryLimiter(uid)) {
             const char* host = mHost.starts_with('^') ? nullptr : mHost.c_str();
             const char* service = mService.starts_with('^') ? nullptr : mService.c_str();
             mHints->ai_family = AF_INET;
             // Don't need to do freeaddrinfo(res) before starting new DNS lookup because previous
             // DNS lookup is failed with error EAI_NODATA.
             *rv = resolv_getaddrinfo(host, service, mHints.get(), &mNetContext, res, event);
-            queryLimiter.finish(uid);
+            endQueryLimiter(uid);
             if (*rv) {
                 *rv = EAI_NODATA;  // return original error code
                 return;
@@ -829,7 +839,7 @@ void DnsProxyListener::GetAddrInfoHandler::run() {
     int32_t rv = 0;
     NetworkDnsEventReported event;
     initDnsEvent(&event, mNetContext);
-    if (queryLimiter.start(uid)) {
+    if (startQueryLimiter(uid)) {
         const char* host = mHost.starts_with('^') ? nullptr : mHost.c_str();
         const char* service = mService.starts_with('^') ? nullptr : mService.c_str();
         if (evaluate_domain_name(mNetContext, host)) {
@@ -837,7 +847,7 @@ void DnsProxyListener::GetAddrInfoHandler::run() {
         } else {
             rv = EAI_SYSTEM;
         }
-        queryLimiter.finish(uid);
+        endQueryLimiter(uid);
     } else {
         // Note that this error code is currently not passed down to the client.
         // android_getaddrinfo_proxy() returns EAI_NODATA on any error.
@@ -1038,14 +1048,14 @@ void DnsProxyListener::ResNSendHandler::run() {
     int ansLen = -1;
     NetworkDnsEventReported event;
     initDnsEvent(&event, mNetContext);
-    if (queryLimiter.start(uid)) {
+    if (startQueryLimiter(uid)) {
         if (evaluate_domain_name(mNetContext, rr_name.c_str())) {
             ansLen = resolv_res_nsend(&mNetContext, {msg.data(), msgLen}, ansBuf, &rcode,
                                       static_cast<ResNsendFlags>(mFlags), &event);
         } else {
             ansLen = -EAI_SYSTEM;
         }
-        queryLimiter.finish(uid);
+        endQueryLimiter(uid);
     } else {
         LOG(WARNING) << "ResNSendHandler::run: resnsend: from UID " << uid
                      << ", max concurrent queries reached";
@@ -1214,10 +1224,10 @@ void DnsProxyListener::GetHostByNameHandler::doDns64Synthesis(int32_t* rv, hoste
 
     // If caller wants IPv6 answers but no data, try to query IPv4 answers for synthesis
     const uid_t uid = mClient->getUid();
-    if (queryLimiter.start(uid)) {
+    if (startQueryLimiter(uid)) {
         const char* name = mName.starts_with('^') ? nullptr : mName.c_str();
         *rv = resolv_gethostbyname(name, AF_INET, hbuf, buf, buflen, &mNetContext, hpp, event);
-        queryLimiter.finish(uid);
+        endQueryLimiter(uid);
         if (*rv) {
             *rv = EAI_NODATA;  // return original error code
             return;
@@ -1245,7 +1255,7 @@ void DnsProxyListener::GetHostByNameHandler::run() {
     int32_t rv = 0;
     NetworkDnsEventReported event;
     initDnsEvent(&event, mNetContext);
-    if (queryLimiter.start(uid)) {
+    if (startQueryLimiter(uid)) {
         const char* name = mName.starts_with('^') ? nullptr : mName.c_str();
         if (evaluate_domain_name(mNetContext, name)) {
             rv = resolv_gethostbyname(name, mAf, &hbuf, tmpbuf, sizeof tmpbuf, &mNetContext, &hp,
@@ -1253,7 +1263,7 @@ void DnsProxyListener::GetHostByNameHandler::run() {
         } else {
             rv = EAI_SYSTEM;
         }
-        queryLimiter.finish(uid);
+        endQueryLimiter(uid);
     } else {
         rv = EAI_MEMORY;
         LOG(ERROR) << "GetHostByNameHandler::run: from UID " << uid
@@ -1376,12 +1386,12 @@ void DnsProxyListener::GetHostByAddrHandler::doDns64ReverseLookup(hostent* hbuf,
     }
 
     const uid_t uid = mClient->getUid();
-    if (queryLimiter.start(uid)) {
+    if (startQueryLimiter(uid)) {
         // Remove NAT64 prefix and do reverse DNS query
         struct in_addr v4addr = {.s_addr = v6addr.s6_addr32[3]};
         resolv_gethostbyaddr(&v4addr, sizeof(v4addr), AF_INET, hbuf, buf, buflen, &mNetContext, hpp,
                              event);
-        queryLimiter.finish(uid);
+        endQueryLimiter(uid);
         if (*hpp) {
             // Replace IPv4 address with original queried IPv6 address in place. The space has
             // reserved by dns_gethtbyaddr() and netbsd_gethostent_r() in
@@ -1407,7 +1417,7 @@ void DnsProxyListener::GetHostByAddrHandler::run() {
     int32_t rv = 0;
     NetworkDnsEventReported event;
     initDnsEvent(&event, mNetContext);
-    if (queryLimiter.start(uid)) {
+    if (startQueryLimiter(uid)) {
         // From Android U, evaluate_domain_name() is not only for OEM customization, but also tells
         // DNS resolver whether the UID can send DNS on the specified network. The function needs
         // to be called even when there is no domain name to evaluate (GetHostByAddr). This is
@@ -1420,7 +1430,7 @@ void DnsProxyListener::GetHostByAddrHandler::run() {
             rv = resolv_gethostbyaddr(&mAddress, mAddressLen, mAddressFamily, &hbuf, tmpbuf,
                                       sizeof tmpbuf, &mNetContext, &hp, &event);
         }
-        queryLimiter.finish(uid);
+        endQueryLimiter(uid);
     } else {
         rv = EAI_MEMORY;
         LOG(ERROR) << "GetHostByAddrHandler::run: from UID " << uid
diff --git a/OperationLimiter.h b/OperationLimiter.h
index 332157a7..24f4dc3c 100644
--- a/OperationLimiter.h
+++ b/OperationLimiter.h
@@ -56,10 +56,8 @@ class OperationLimiter {
     //
     // Note: each successful start(key) must be matched by exactly one call to
     // finish(key).
-    bool start(KeyType key) EXCLUDES(mMutex) {
+    bool start(KeyType key, int globalLimit = INT_MAX) EXCLUDES(mMutex) {
         std::lock_guard lock(mMutex);
-        int globalLimit =
-                android::net::Experiments::getInstance()->getFlag("max_queries_global", INT_MAX);
         if (globalLimit < mLimitPerKey) {
             LOG(ERROR) << "Misconfiguration on max_queries_global " << globalLimit;
             globalLimit = INT_MAX;
diff --git a/OperationLimiterTest.cpp b/OperationLimiterTest.cpp
index 90ab9d9a..ad8c8143 100644
--- a/OperationLimiterTest.cpp
+++ b/OperationLimiterTest.cpp
@@ -74,5 +74,37 @@ TEST_F(OperationLimiterTest, destroyWithActiveOperations) {
             "" /* "active operations */);
 }
 
+TEST_F(OperationLimiterTest, globalLimits) {
+    OperationLimiter<int> limiter(1);
+
+    EXPECT_TRUE(limiter.start(42, 2));
+
+    // Calling with a different key is okay.
+    EXPECT_TRUE(limiter.start(43, 2));
+
+    // Global limit reached... calling with a different key should have no effect.
+    EXPECT_FALSE(limiter.start(44, 2));
+
+    // Global limit extended... calling with a different key is available again.
+    EXPECT_TRUE(limiter.start(44, 4));
+
+    // Per-key limit reached.
+    EXPECT_FALSE(limiter.start(44, 4));
+
+    // Global limit is still available.
+    EXPECT_TRUE(limiter.start(45, 4));
+
+    // Global limit reached again.
+    EXPECT_FALSE(limiter.start(46, 4));
+
+    // Shrink global limit.
+    EXPECT_FALSE(limiter.start(46, 3));
+
+    // Finish all pending operations
+    for (const auto& key : {42, 43, 44, 45}) {
+        limiter.finish(key);
+    }
+}
+
 }  // namespace netdutils
 }  // namespace android