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
diff --git a/libnetdutils/Android.bp b/libnetdutils/Android.bp
index ba3c3c9..9dd6cfb 100644
--- a/libnetdutils/Android.bp
+++ b/libnetdutils/Android.bp
@@ -16,6 +16,7 @@
shared_libs: [
"libbase",
"libbinder",
+ "liblog",
],
export_shared_lib_headers: [
"libbase",
@@ -29,6 +30,7 @@
"BackoffSequenceTest.cpp",
"FdTest.cpp",
"MemBlockTest.cpp",
+ "OperationLimiterTest.cpp",
"SliceTest.cpp",
"StatusTest.cpp",
"SyscallsTest.cpp",
@@ -39,5 +41,8 @@
"-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 0000000..8d72d75
--- /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 0000000..633536b
--- /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