Limit the number of outstanding DNS queries by UID

Test: system/netd/tests/runtests.sh
Test: run netdutils_test on marlin:
    [----------] 3 tests from OperationLimiter
    [ RUN      ] OperationLimiter.limits
    [       OK ] OperationLimiter.limits (0 ms)
    [ RUN      ] OperationLimiter.finishWithoutStart
    [       OK ] OperationLimiter.finishWithoutStart (3 ms)
    [ RUN      ] OperationLimiter.destroyWithActiveOperations
    [       OK ] OperationLimiter.destroyWithActiveOperations (1 ms)
    [----------] 3 tests from OperationLimiter (6 ms total)
Bug: 79674503
Change-Id: I5f11f0ed6b6f2479921d90a919d17dfd7b7f5788
Merged-In: I5f11f0ed6b6f2479921d90a919d17dfd7b7f5788

5 files changed, 225 insertions, 15 deletions

diff --git a/libnetdutils/Android.bp b/libnetdutils/Android.bp
index ba3c3c9c..9dd6cfb2 100644
--- a/libnetdutils/Android.bp
+++ b/libnetdutils/Android.bp
@@ -16,6 +16,7 @@ cc_library {
     shared_libs: [
         "libbase",
         "libbinder",
+        "liblog",
     ],
     export_shared_lib_headers: [
         "libbase",
@@ -29,6 +30,7 @@ cc_test {
         "BackoffSequenceTest.cpp",
         "FdTest.cpp",
         "MemBlockTest.cpp",
+        "OperationLimiterTest.cpp",
         "SliceTest.cpp",
         "StatusTest.cpp",
         "SyscallsTest.cpp",
@@ -39,5 +41,8 @@ cc_test {
         "-Wno-error=unused-variable",
     ],
     static_libs: ["libgmock"],
-    shared_libs: ["libnetdutils"],
+    shared_libs: [
+        "libbase",
+        "libnetdutils",
+    ],
 }
diff --git a/libnetdutils/OperationLimiterTest.cpp b/libnetdutils/OperationLimiterTest.cpp
new file mode 100644
index 00000000..8d72d752
--- /dev/null
+++ b/libnetdutils/OperationLimiterTest.cpp
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "netdutils/OperationLimiter.h"
+
+#include <gtest/gtest-spi.h>
+
+namespace android {
+namespace netdutils {
+
+TEST(OperationLimiter, limits) {
+    OperationLimiter<int> limiter(3);
+
+    EXPECT_TRUE(limiter.start(42));
+    EXPECT_TRUE(limiter.start(42));
+    EXPECT_TRUE(limiter.start(42));
+
+    // Limit reached... calling any number of times should have no effect.
+    EXPECT_FALSE(limiter.start(42));
+    EXPECT_FALSE(limiter.start(42));
+    EXPECT_FALSE(limiter.start(42));
+
+    // Finishing a single operations is enough for starting a new one...
+    limiter.finish(42);
+    EXPECT_TRUE(limiter.start(42));
+
+    // ...but not two!
+    EXPECT_FALSE(limiter.start(42));
+
+    // Different ids should still have quota...
+    EXPECT_TRUE(limiter.start(666));
+    limiter.finish(666);
+
+    // Finish all pending operations
+    limiter.finish(42);
+    limiter.finish(42);
+    limiter.finish(42);
+}
+
+TEST(OperationLimiter, finishWithoutStart) {
+    OperationLimiter<int> limiter(1);
+
+    // Will output a LOG(FATAL_WITHOUT_ABORT), but we have no way to probe this.
+    limiter.finish(42);
+
+    // This will ensure that the finish() above didn't set a negative value.
+    EXPECT_TRUE(limiter.start(42));
+    EXPECT_FALSE(limiter.start(42));
+}
+
+TEST(OperationLimiter, destroyWithActiveOperations) {
+    // The death message doesn't seem to be captured on Android.
+    EXPECT_DEBUG_DEATH({
+        OperationLimiter<int> limiter(3);
+        limiter.start(42);
+    }, "" /* "active operations */);
+}
+
+}  // namespace netdutils
+}  // namespace android
diff --git a/libnetdutils/include/netdutils/OperationLimiter.h b/libnetdutils/include/netdutils/OperationLimiter.h
new file mode 100644
index 00000000..633536b1
--- /dev/null
+++ b/libnetdutils/include/netdutils/OperationLimiter.h
@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef NETUTILS_OPERATIONLIMITER_H
+#define NETUTILS_OPERATIONLIMITER_H
+
+#include <mutex>
+#include <unordered_map>
+
+#include <android-base/logging.h>
+#include <android-base/thread_annotations.h>
+
+namespace android {
+namespace netdutils {
+
+// Tracks the number of operations in progress on behalf of a particular key or
+// ID, rejecting further attempts to start new operations after a configurable
+// limit has been reached.
+//
+// The intended usage pattern is:
+//     OperationLimiter<UserId> connections_per_user;
+//     ...
+//     // Before opening a new connection
+//     if (!limiter.start(user)) {
+//         return error;
+//     } else {
+//         // open the connection
+//         // ...do some work...
+//         // close the connection
+//         limiter.finish(user);
+//     }
+//
+// This class is thread-safe.
+template<typename KeyType>
+class OperationLimiter {
+public:
+    explicit OperationLimiter(int limit) : mLimitPerKey(limit) {}
+
+    ~OperationLimiter() {
+        DCHECK(mCounters.empty())
+                << "Destroying OperationLimiter with active operations";
+    }
+
+    // Returns false if |key| has reached the maximum number of concurrent
+    // operations, otherwise increments the counter and returns true.
+    //
+    // Note: each successful start(key) must be matched by exactly one call to
+    // finish(key).
+    bool start(KeyType key) EXCLUDES(mMutex) {
+        std::lock_guard<std::mutex> lock(mMutex);
+        auto& cnt = mCounters[key];  // operator[] creates new entries as needed.
+        if (cnt >= mLimitPerKey) {
+            // Oh, no!
+            return false;
+        }
+        ++cnt;
+        return true;
+    }
+
+    // Decrements the number of operations in progress accounted to |key|.
+    // See usage notes on start().
+    void finish(KeyType key) EXCLUDES(mMutex) {
+        std::lock_guard<std::mutex> lock(mMutex);
+        auto it = mCounters.find(key);
+        if (it == mCounters.end()) {
+            LOG(FATAL_WITHOUT_ABORT) << "Decremented non-existent counter for key=" << key;
+            return;
+        }
+        auto& cnt = it->second;
+        --cnt;
+        if (cnt <= 0) {
+            // Cleanup counters once they drop down to zero.
+            mCounters.erase(it);
+        }
+    }
+
+private:
+    // Protects access to the map below.
+    std::mutex mMutex;
+
+    // Tracks the number of outstanding queries by key.
+    std::unordered_map<KeyType, int> mCounters GUARDED_BY(mMutex);
+
+    // Maximum number of outstanding queries from a single key.
+    const int mLimitPerKey;
+};
+
+}  // namespace netdutils
+}  // namespace android
+
+#endif // NETUTILS_OPERATIONLIMITER_H
diff --git a/server/DnsProxyListener.cpp b/server/DnsProxyListener.cpp
index dd7b4e2b..9babac2c 100644
--- a/server/DnsProxyListener.cpp
+++ b/server/DnsProxyListener.cpp
@@ -41,6 +41,7 @@
 #include <cutils/log.h>
 #include <cutils/misc.h>
 #include <netdutils/Slice.h>
+#include <netdutils/OperationLimiter.h>
 #include <utils/String16.h>
 #include <sysutils/SocketClient.h>
 
@@ -74,6 +75,10 @@ constexpr const char CONNECTIVITY_USE_RESTRICTED_NETWORKS[] =
 constexpr const char NETWORK_BYPASS_PRIVATE_DNS[] =
     "android.permission.NETWORK_BYPASS_PRIVATE_DNS";
 
+// Limits the number of outstanding DNS queries by client UID.
+constexpr int MAX_QUERIES_PER_UID = 256;
+android::netdutils::OperationLimiter<uid_t> queryLimiter(MAX_QUERIES_PER_UID);
+
 void logArguments(int argc, char** argv) {
     for (int i = 0; i < argc; i++) {
         ALOGD("argv[%i]=%s", i, argv[i]);
@@ -353,7 +358,17 @@ void DnsProxyListener::GetAddrInfoHandler::run() {
     struct addrinfo* result = NULL;
     Stopwatch s;
     maybeFixupNetContext(&mNetContext);
-    uint32_t rv = android_getaddrinfofornetcontext(mHost, mService, mHints, &mNetContext, &result);
+    const uid_t uid = mClient->getUid();
+    uint32_t rv = 0;
+    if (queryLimiter.start(uid)) {
+        rv = android_getaddrinfofornetcontext(mHost, mService, mHints, &mNetContext, &result);
+        queryLimiter.finish(uid);
+    } else {
+        // Note that this error code is currently not passed down to the client.
+        // android_getaddrinfo_proxy() returns EAI_NODATA on any error.
+        rv = EAI_MEMORY;
+        ALOGE("getaddrinfo: from UID %d, max concurrent queries reached", uid);
+    }
     const int latencyMs = lround(s.timeTaken());
 
     if (rv) {
@@ -463,7 +478,7 @@ int DnsProxyListener::GetAddrInfoCmd::runCommand(SocketClient *cli,
     int ai_protocol = atoi(argv[6]);
     unsigned netId = strtoul(argv[7], NULL, 10);
     const bool useLocalNameservers = checkAndClearUseLocalNameserversFlag(&netId);
-    uid_t uid = cli->getUid();
+    const uid_t uid = cli->getUid();
 
     android_net_context netcontext;
     mDnsProxyListener->mNetCtrl->getNetworkContext(netId, uid, &netcontext);
@@ -563,16 +578,23 @@ DnsProxyListener::GetHostByNameHandler::~GetHostByNameHandler() {
 
 void DnsProxyListener::GetHostByNameHandler::run() {
     if (DBG) {
-        ALOGD("DnsProxyListener::GetHostByNameHandler::run\n");
+        ALOGD("DnsProxyListener::GetHostByNameHandler::run");
     }
 
     Stopwatch s;
     maybeFixupNetContext(&mNetContext);
-    struct hostent* hp = android_gethostbynamefornetcontext(mName, mAf, &mNetContext);
+    const uid_t uid = mClient->getUid();
+    struct hostent* hp = nullptr;
+    if (queryLimiter.start(uid)) {
+        hp = android_gethostbynamefornetcontext(mName, mAf, &mNetContext);
+        queryLimiter.finish(uid);
+    } else {
+        ALOGE("gethostbyname: from UID %d, max concurrent queries reached", uid);
+    }
     const int latencyMs = lround(s.timeTaken());
 
     if (DBG) {
-        ALOGD("GetHostByNameHandler::run gethostbyname errno: %s hp->h_name = %s, name_len = %zu\n",
+        ALOGD("GetHostByNameHandler::run gethostbyname errno: %s hp->h_name = %s, name_len = %zu",
                 hp ? "success" : strerror(errno),
                 (hp && hp->h_name) ? hp->h_name : "null",
                 (hp && hp->h_name) ? strlen(hp->h_name) + 1 : 0);
@@ -587,7 +609,7 @@ void DnsProxyListener::GetHostByNameHandler::run() {
     }
 
     if (!success) {
-        ALOGW("GetHostByNameHandler: Error writing DNS result to client\n");
+        ALOGW("GetHostByNameHandler: Error writing DNS result to client");
     }
 
     if (mNetdEventListener != nullptr) {
@@ -625,7 +647,7 @@ void DnsProxyListener::GetHostByNameHandler::run() {
                 mNetdEventListener->onDnsEvent(mNetContext.dns_netid,
                                                INetdEventListener::EVENT_GETHOSTBYNAME,
                                                h_errno, latencyMs, String16(mName), ip_addrs,
-                                               total_ip_addr_count, mClient->getUid());
+                                               total_ip_addr_count, uid);
                 break;
         }
     }
@@ -706,16 +728,22 @@ DnsProxyListener::GetHostByAddrHandler::~GetHostByAddrHandler() {
 
 void DnsProxyListener::GetHostByAddrHandler::run() {
     if (DBG) {
-        ALOGD("DnsProxyListener::GetHostByAddrHandler::run\n");
+        ALOGD("DnsProxyListener::GetHostByAddrHandler::run");
     }
 
-    // NOTE gethostbyaddr should take a void* but bionic thinks it should be char*
     maybeFixupNetContext(&mNetContext);
-    struct hostent* hp = android_gethostbyaddrfornetcontext(
-            (char*) mAddress, mAddressLen, mAddressFamily, &mNetContext);
+    const uid_t uid = mClient->getUid();
+    struct hostent* hp = nullptr;
+    if (queryLimiter.start(uid)) {
+        hp = android_gethostbyaddrfornetcontext(
+                mAddress, mAddressLen, mAddressFamily, &mNetContext);
+        queryLimiter.finish(uid);
+    } else {
+        ALOGE("gethostbyaddr: from UID %d, max concurrent queries reached", uid);
+    }
 
     if (DBG) {
-        ALOGD("GetHostByAddrHandler::run gethostbyaddr errno: %s hp->h_name = %s, name_len = %zu\n",
+        ALOGD("GetHostByAddrHandler::run gethostbyaddr errno: %s hp->h_name = %s, name_len = %zu",
                 hp ? "success" : strerror(errno),
                 (hp && hp->h_name) ? hp->h_name : "null",
                 (hp && hp->h_name) ? strlen(hp->h_name) + 1 : 0);
@@ -730,7 +758,7 @@ void DnsProxyListener::GetHostByAddrHandler::run() {
     }
 
     if (!success) {
-        ALOGW("GetHostByAddrHandler: Error writing DNS result to client\n");
+        ALOGW("GetHostByAddrHandler: Error writing DNS result to client");
     }
     mClient->decRef();
 }
diff --git a/tests/dns_tls_test.cpp b/tests/dns_tls_test.cpp
index 7820338c..bb5bfe56 100644
--- a/tests/dns_tls_test.cpp
+++ b/tests/dns_tls_test.cpp
@@ -135,7 +135,7 @@ public:
     bool query(uint16_t id, const Slice query) override {
         // Return the response immediately (asynchronously).
         std::thread(&IDnsTlsSocketObserver::onResponse, mObserver, make_echo(id, query)).detach();
-	return true;
+        return true;
     }
 private:
     IDnsTlsSocketObserver* const mObserver;