DoH: Make session resumption as a flag

Use a flag, doh_session_resumption, to control whether or not to
enable session resumption. If the flag is unset, session resumption
is disabled.

Because the value of the flag is cached in the DoH client, if
there are some networks existing before setting a new value to
the flag, those network will still use the old value of the flag
rather than the new one until private DNS settings changes.

Bug: 205922811
Test: cd packages/modules/DnsResolver && atest
Test: changed the flag and checked that DnsResolver behaved as
      expected
Change-Id: I6fb24d4251b6e0fc163d169bb9cc68703bb76812

8 files changed, 66 insertions, 14 deletions

diff --git a/Experiments.h b/Experiments.h
index d03759e1..ea267e82 100644
--- a/Experiments.h
+++ b/Experiments.h
@@ -65,6 +65,7 @@ class Experiments {
             "doh_query_timeout_ms",
             "doh_probe_timeout_ms",
             "doh_idle_timeout_ms",
+            "doh_session_resumption",
             "mdns_resolution",
     };
     // This value is used in updateInternal as the default value if any flags can't be found.
diff --git a/PrivateDnsConfiguration.cpp b/PrivateDnsConfiguration.cpp
index 0fef5fe7..dd0be8f4 100644
--- a/PrivateDnsConfiguration.cpp
+++ b/PrivateDnsConfiguration.cpp
@@ -505,9 +505,12 @@ int PrivateDnsConfiguration::setDoh(int32_t netId, uint32_t mark,
                         getTimeoutFromFlag("doh_probe_timeout_ms", kDohProbeDefaultTimeoutMs),
                 .idle_timeout_ms =
                         getTimeoutFromFlag("doh_idle_timeout_ms", kDohIdleDefaultTimeoutMs),
+                .use_session_resumption =
+                        Experiments::getInstance()->getFlag("doh_session_resumption", 0) == 1,
         };
         LOG(DEBUG) << __func__ << ": probe_timeout_ms=" << flags.probe_timeout_ms
-                   << ", idle_timeout_ms=" << flags.idle_timeout_ms;
+                   << ", idle_timeout_ms=" << flags.idle_timeout_ms
+                   << ", use_session_resumption=" << flags.use_session_resumption;
 
         return doh_net_new(mDohDispatcher, netId, dohId.httpsTemplate.c_str(), dohId.host.c_str(),
                            dohId.ipAddr.c_str(), mark, caCert.c_str(), &flags);
diff --git a/doh.h b/doh.h
index af7c6b34..75ecf10a 100644
--- a/doh.h
+++ b/doh.h
@@ -53,6 +53,7 @@ struct DohDispatcher;
 struct FeatureFlags {
     uint64_t probe_timeout_ms;
     uint64_t idle_timeout_ms;
+    bool use_session_resumption;
 };
 
 using ValidationCallback = void (*)(uint32_t net_id, bool success, const char* ip_addr,
diff --git a/doh/ffi.rs b/doh/ffi.rs
index 9ecc0e8d..65ca012e 100644
--- a/doh/ffi.rs
+++ b/doh/ffi.rs
@@ -42,6 +42,7 @@ pub type TagSocketCallback = extern "C" fn(sock: RawFd);
 pub struct FeatureFlags {
     probe_timeout_ms: uint64_t,
     idle_timeout_ms: uint64_t,
+    use_session_resumption: bool,
 }
 
 fn wrap_validation_callback(validation_fn: ValidationCallback) -> ValidationReporter {
@@ -231,6 +232,7 @@ pub unsafe extern "C" fn doh_net_new(
             sk_mark,
             cert_path,
             idle_timeout_ms: flags.idle_timeout_ms,
+            use_session_resumption: flags.use_session_resumption,
         },
         timeout: Duration::from_millis(flags.probe_timeout_ms),
     };
@@ -381,6 +383,7 @@ mod tests {
             sk_mark: 0,
             cert_path: None,
             idle_timeout_ms: 0,
+            use_session_resumption: true,
         };
 
         wrap_validation_callback(success_cb)(&info, true).await;
diff --git a/doh/network/driver.rs b/doh/network/driver.rs
index 67c32ae7..375ac2c1 100644
--- a/doh/network/driver.rs
+++ b/doh/network/driver.rs
@@ -185,7 +185,8 @@ impl Driver {
         }
 
         if !self.connection.wait_for_live().await {
-            let session = self.connection.session();
+            let session =
+                if self.info.use_session_resumption { self.connection.session() } else { None };
             // Try reconnecting
             self.connection =
                 build_connection(&self.info, &self.tag_socket, &mut self.config, session).await?;
diff --git a/doh/network/mod.rs b/doh/network/mod.rs
index 21e47ec3..19be8643 100644
--- a/doh/network/mod.rs
+++ b/doh/network/mod.rs
@@ -48,6 +48,7 @@ pub struct ServerInfo {
     pub sk_mark: u32,
     pub cert_path: Option<String>,
     pub idle_timeout_ms: u64,
+    pub use_session_resumption: bool,
 }
 
 #[derive(Debug)]
diff --git a/tests/doh_ffi_test.cpp b/tests/doh_ffi_test.cpp
index 954e1ced..d858a760 100644
--- a/tests/doh_ffi_test.cpp
+++ b/tests/doh_ffi_test.cpp
@@ -56,6 +56,7 @@ TEST(DoHFFITest, SmokeTest) {
     const FeatureFlags flags = {
             .probe_timeout_ms = TIMEOUT_MS,
             .idle_timeout_ms = TIMEOUT_MS,
+            .use_session_resumption = true,
     };
 
     // TODO: Use a local server instead of dns.google.
diff --git a/tests/resolv_private_dns_test.cpp b/tests/resolv_private_dns_test.cpp
index 931d60cb..1791bb64 100644
--- a/tests/resolv_private_dns_test.cpp
+++ b/tests/resolv_private_dns_test.cpp
@@ -51,11 +51,14 @@ using android::netdutils::IPSockAddr;
 using android::netdutils::ScopedAddrinfo;
 using android::netdutils::Stopwatch;
 using std::chrono::milliseconds;
+using std::this_thread::sleep_for;
 
 const std::string kDohFlag("persist.device_config.netd_native.doh");
 const std::string kDohQueryTimeoutFlag("persist.device_config.netd_native.doh_query_timeout_ms");
 const std::string kDohProbeTimeoutFlag("persist.device_config.netd_native.doh_probe_timeout_ms");
 const std::string kDohIdleTimeoutFlag("persist.device_config.netd_native.doh_idle_timeout_ms");
+const std::string kDohSessionResumptionFlag(
+        "persist.device_config.netd_native.doh_session_resumption");
 
 constexpr int MAXPACKET = (8 * 1024);
 
@@ -217,7 +220,7 @@ class BaseTest : public ::testing::Test {
     }
 
     bool hasUncaughtPrivateDnsValidation(const std::string& serverAddr) {
-        std::this_thread::sleep_for(std::chrono::milliseconds(200));
+        sleep_for(milliseconds(200));
         return sUnsolicitedEventListener->findValidationRecord(
                        serverAddr, IDnsResolverUnsolicitedEventListener::PROTOCOL_DOT) ||
                sUnsolicitedEventListener->findValidationRecord(
@@ -690,7 +693,7 @@ TEST_F(PrivateDnsDohTest, TemporaryConnectionStalled) {
     Stopwatch s;
     int fd = resNetworkQuery(TEST_NETID, kQueryHostname, ns_c_in, ns_t_a,
                              ANDROID_RESOLV_NO_CACHE_LOOKUP);
-    std::this_thread::sleep_for(std::chrono::milliseconds(connectionStalledTimeMs));
+    sleep_for(milliseconds(connectionStalledTimeMs));
     EXPECT_TRUE(doh.block_sending(false));
 
     expectAnswersValid(fd, AF_INET, kQueryAnswerA);
@@ -714,7 +717,7 @@ TEST_F(PrivateDnsDohTest, ExcessDnsRequests) {
     EXPECT_TRUE(doh.setMaxIdleTimeout(initial_max_idle_timeout_ms));
     // Sleep a while to avoid binding socket failed.
     // TODO: Make DohFrontend retry binding sockets.
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    sleep_for(milliseconds(100));
     ASSERT_TRUE(doh.startServer());
 
     auto parcel = DnsResponderClient::GetDefaultResolverParamsParcel();
@@ -767,7 +770,7 @@ TEST_F(PrivateDnsDohTest, ExcessDnsRequests) {
 
     // Sleep a while to wait for DoH and DoT validation.
     // TODO: Extend WaitForDohValidation() to support passing a netId.
-    std::this_thread::sleep_for(std::chrono::milliseconds(200));
+    sleep_for(milliseconds(200));
     EXPECT_TRUE(dot_ipv6.waitForQueries(1));
 
     int fd = resNetworkQuery(TEST_NETID_2, kQueryHostname, ns_c_in, ns_t_aaaa,
@@ -793,7 +796,7 @@ TEST_F(PrivateDnsDohTest, RunOutOfDataLimit) {
     EXPECT_TRUE(doh.setMaxBufferSize(initial_max_data));
     // Sleep a while to avoid binding socket failed.
     // TODO: Make DohFrontend retry binding sockets.
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    sleep_for(milliseconds(100));
     ASSERT_TRUE(doh.startServer());
 
     const auto parcel = DnsResponderClient::GetDefaultResolverParamsParcel();
@@ -815,7 +818,7 @@ TEST_F(PrivateDnsDohTest, RunOutOfDataLimit) {
             expectAnswersValid(fd, AF_INET, kQueryAnswerA);
         });
     }
-    std::this_thread::sleep_for(std::chrono::milliseconds(500));
+    sleep_for(milliseconds(500));
     EXPECT_TRUE(doh.block_sending(false));
 
     // In current implementation, the fifth DoH query will get blocked and result in timeout.
@@ -844,7 +847,7 @@ TEST_F(PrivateDnsDohTest, RunOutOfStreams) {
     ASSERT_TRUE(doh.stopServer());
     EXPECT_TRUE(doh.setMaxStreamsBidi(initial_max_streams_bidi));
     // Sleep a while to avoid binding socket failed.
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    sleep_for(milliseconds(100));
     ASSERT_TRUE(doh.startServer());
 
     const auto parcel = DnsResponderClient::GetDefaultResolverParamsParcel();
@@ -866,7 +869,7 @@ TEST_F(PrivateDnsDohTest, RunOutOfStreams) {
             expectAnswersValid(fd, AF_INET, kQueryAnswerA);
         });
     }
-    std::this_thread::sleep_for(std::chrono::milliseconds(500));
+    sleep_for(milliseconds(500));
     EXPECT_TRUE(doh.block_sending(false));
 
     for (std::thread& thread : threads) {
@@ -884,7 +887,7 @@ TEST_F(PrivateDnsDohTest, ReconnectAfterIdleTimeout) {
     ASSERT_TRUE(doh.stopServer());
     EXPECT_TRUE(doh.setMaxIdleTimeout(initial_max_idle_timeout_ms));
     // Sleep a while to avoid binding socket failed.
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    sleep_for(milliseconds(100));
     ASSERT_TRUE(doh.startServer());
 
     const auto parcel = DnsResponderClient::GetDefaultResolverParamsParcel();
@@ -898,7 +901,7 @@ TEST_F(PrivateDnsDohTest, ReconnectAfterIdleTimeout) {
 
     for (int i = 0; i < 5; i++) {
         SCOPED_TRACE(fmt::format("Round: {}", i));
-        std::this_thread::sleep_for(std::chrono::milliseconds(initial_max_idle_timeout_ms + 500));
+        sleep_for(milliseconds(initial_max_idle_timeout_ms + 500));
 
         // As the connection is closed, the DnsResolver will reconnect to the DoH server
         // for this DNS request.
@@ -909,7 +912,6 @@ TEST_F(PrivateDnsDohTest, ReconnectAfterIdleTimeout) {
 
     EXPECT_NO_FAILURE(expectQueries(0 /* dns */, 0 /* dot */, 5 /* doh */));
     EXPECT_EQ(doh.connections(), 6);
-    EXPECT_EQ(doh.resumedConnections(), 5);
 }
 
 // Tests that the experiment flag doh_idle_timeout_ms is effective.
@@ -945,9 +947,48 @@ TEST_F(PrivateDnsDohTest, ConnectionIdleTimer) {
     EXPECT_EQ(doh.connections(), 1);
 
     // Expect that the DoH connection gets disconnected while sleeping.
-    std::this_thread::sleep_for(std::chrono::milliseconds(connection_idle_timeout + tolerance_ms));
+    sleep_for(milliseconds(connection_idle_timeout + tolerance_ms));
 
     EXPECT_NO_FAILURE(sendQueryAndCheckResult());
     EXPECT_NO_FAILURE(expectQueries(0 /* dns */, 0 /* dot */, 4 /* doh */));
     EXPECT_EQ(doh.connections(), 2);
 }
+
+// Tests that the flag "doh_session_resumption" works as expected.
+TEST_F(PrivateDnsDohTest, SessionResumption) {
+    const int initial_max_idle_timeout_ms = 1000;
+    for (const auto& flag : {"0", "1"}) {
+        auto sp = make_unique<ScopedSystemProperties>(kDohSessionResumptionFlag, flag);
+        resetNetwork();
+
+        ASSERT_TRUE(doh.stopServer());
+        EXPECT_TRUE(doh.setMaxIdleTimeout(initial_max_idle_timeout_ms));
+        // Sleep a while to avoid binding socket failed.
+        sleep_for(milliseconds(100));
+        ASSERT_TRUE(doh.startServer());
+
+        const auto parcel = DnsResponderClient::GetDefaultResolverParamsParcel();
+        ASSERT_TRUE(mDnsClient.SetResolversFromParcel(parcel));
+        EXPECT_TRUE(WaitForDohValidation(test::kDefaultListenAddr, true));
+        EXPECT_TRUE(WaitForDotValidation(test::kDefaultListenAddr, true));
+        EXPECT_TRUE(dot.waitForQueries(1));
+        dot.clearQueries();
+        doh.clearQueries();
+        dns.clearQueries();
+
+        for (int i = 0; i < 2; i++) {
+            SCOPED_TRACE(fmt::format("Round: {}", i));
+            sleep_for(milliseconds(initial_max_idle_timeout_ms + 500));
+
+            // As the connection is closed, the DnsResolver will reconnect to the DoH server
+            // for this DNS request.
+            int fd = resNetworkQuery(TEST_NETID, kQueryHostname, ns_c_in, ns_t_a,
+                                     ANDROID_RESOLV_NO_CACHE_LOOKUP);
+            expectAnswersValid(fd, AF_INET, kQueryAnswerA);
+        }
+
+        EXPECT_NO_FAILURE(expectQueries(0 /* dns */, 0 /* dot */, 2 /* doh */));
+        EXPECT_EQ(doh.connections(), 3);
+        EXPECT_EQ(doh.resumedConnections(), (flag == "1" ? 2 : 0));
+    }
+}