DoH: Support Early Data

Early Data can be enabled by the flag `doh_early_data`.

As 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 old flag until private
DNS setting changes.

Bug: 235763732
Test: atest
Test: atest doh_ffi_test
Test: manual test on wifi.
  1. adb shell tc qdisc add dev wlan0 root netem delay 500ms
  2. Observed the first DNS query latency on a subsequent connection:
     - when doh_early_data=0/doh_session_resumption=0, it took ~1500 ms
     - when doh_early_data=1/doh_session_resumption=1, it took ~500 ms
Change-Id: I06bc5f9aa006e357dc3ecf04a693126a55782e87

7 files changed, 63 insertions, 30 deletions

diff --git a/doh/config.rs b/doh/config.rs
index 1f91a151..bcc21184 100644
--- a/doh/config.rs
+++ b/doh/config.rs
@@ -63,6 +63,9 @@ impl Config {
             }
             None => config.verify_peer(false),
         }
+        if key.enable_early_data {
+            config.enable_early_data();
+        }
 
         // Some of these configs are necessary, or the server can't respond the HTTP/3 request.
         config.set_max_idle_timeout(key.max_idle_timeout);
@@ -126,6 +129,7 @@ pub struct Cache {
 pub struct Key {
     pub cert_path: Option<String>,
     pub max_idle_timeout: u64,
+    pub enable_early_data: bool,
 }
 
 impl Cache {
@@ -174,13 +178,15 @@ impl Cache {
 #[test]
 fn create_quiche_config() {
     assert!(
-        Config::from_key(&Key { cert_path: None, max_idle_timeout: 1000 }).is_ok(),
+        Config::from_key(&Key { cert_path: None, max_idle_timeout: 1000, enable_early_data: true })
+            .is_ok(),
         "quiche config without cert creating failed"
     );
     assert!(
         Config::from_key(&Key {
             cert_path: Some("data/local/tmp/".to_string()),
-            max_idle_timeout: 1000
+            max_idle_timeout: 1000,
+            enable_early_data: true,
         })
         .is_ok(),
         "quiche config with cert creating failed"
@@ -191,38 +197,53 @@ fn create_quiche_config() {
 fn shared_cache() {
     let cache_a = Cache::new();
     let cache_b = cache_a.clone();
-    let config_a = cache_a.get(&Key { cert_path: None, max_idle_timeout: 1000 }).unwrap();
+    let config_a = cache_a
+        .get(&Key { cert_path: None, max_idle_timeout: 1000, enable_early_data: true })
+        .unwrap();
     assert_eq!(Arc::strong_count(&config_a.0), 2);
-    let _config_b = cache_b.get(&Key { cert_path: None, max_idle_timeout: 1000 }).unwrap();
+    let _config_b = cache_b
+        .get(&Key { cert_path: None, max_idle_timeout: 1000, enable_early_data: true })
+        .unwrap();
     assert_eq!(Arc::strong_count(&config_a.0), 3);
 }
 
 #[test]
 fn different_keys() {
     let cache = Cache::new();
-    let key_a = Key { cert_path: None, max_idle_timeout: 1000 };
-    let key_b = Key { cert_path: Some("a".to_string()), max_idle_timeout: 1000 };
-    let key_c = Key { cert_path: Some("a".to_string()), max_idle_timeout: 5000 };
+    let key_a = Key { cert_path: None, max_idle_timeout: 1000, enable_early_data: false };
+    let key_b =
+        Key { cert_path: Some("a".to_string()), max_idle_timeout: 1000, enable_early_data: false };
+    let key_c =
+        Key { cert_path: Some("a".to_string()), max_idle_timeout: 5000, enable_early_data: false };
+    let key_d =
+        Key { cert_path: Some("a".to_string()), max_idle_timeout: 5000, enable_early_data: true };
     let config_a = cache.get(&key_a).unwrap();
     let config_b = cache.get(&key_b).unwrap();
     let _config_b = cache.get(&key_b).unwrap();
     let config_c = cache.get(&key_c).unwrap();
     let _config_c = cache.get(&key_c).unwrap();
+    let config_d = cache.get(&key_d).unwrap();
+    let _config_d = cache.get(&key_d).unwrap();
 
     assert_eq!(Arc::strong_count(&config_a.0), 1);
     assert_eq!(Arc::strong_count(&config_b.0), 2);
+    assert_eq!(Arc::strong_count(&config_c.0), 2);
 
-    // config_c was most recently created, so it should have an extra strong reference due to
+    // config_d was most recently created, so it should have an extra strong reference due to
     // keep-alive in the cache.
-    assert_eq!(Arc::strong_count(&config_c.0), 3);
+    assert_eq!(Arc::strong_count(&config_d.0), 3);
 }
 
 #[test]
 fn lifetimes() {
     let cache = Cache::new();
-    let key_a = Key { cert_path: Some("a".to_string()), max_idle_timeout: 1000 };
-    let key_b = Key { cert_path: Some("b".to_string()), max_idle_timeout: 1000 };
-    let config_none = cache.get(&Key { cert_path: None, max_idle_timeout: 1000 }).unwrap();
+    let key_a =
+        Key { cert_path: Some("a".to_string()), max_idle_timeout: 1000, enable_early_data: true };
+    let key_b =
+        Key { cert_path: Some("b".to_string()), max_idle_timeout: 1000, enable_early_data: true };
+    let config_none = cache
+        .get(&Key { cert_path: None, max_idle_timeout: 1000, enable_early_data: true })
+        .unwrap();
     let config_a = cache.get(&key_a).unwrap();
     let config_b = cache.get(&key_b).unwrap();
 
@@ -268,7 +289,9 @@ fn lifetimes() {
 #[tokio::test]
 async fn quiche_connect() {
     use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
-    let mut config = Config::from_key(&Key { cert_path: None, max_idle_timeout: 10 }).unwrap();
+    let mut config =
+        Config::from_key(&Key { cert_path: None, max_idle_timeout: 10, enable_early_data: true })
+            .unwrap();
     let socket_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 42));
     let conn_id = quiche::ConnectionId::from_ref(&[]);
     quiche::connect(None, &conn_id, socket_addr, config.take().await.deref_mut()).unwrap();
diff --git a/doh/connection/driver.rs b/doh/connection/driver.rs
index 1ba8faa7..f452d0a0 100644
--- a/doh/connection/driver.rs
+++ b/doh/connection/driver.rs
@@ -194,6 +194,19 @@ impl Driver {
     }
 
     async fn drive_once(mut self) -> Result<Self> {
+        // If the QUIC connection is live, but the HTTP/3 is not, try to bring it up
+        if self.quiche_conn.is_established() || self.quiche_conn.is_in_early_data() {
+            info!(
+                "Connection {} established on network {}",
+                self.quiche_conn.trace_id(),
+                self.net_id
+            );
+            let h3_config = h3::Config::new()?;
+            let h3_conn = h3::Connection::with_transport(&mut self.quiche_conn, &h3_config)?;
+            self = H3Driver::new(self, h3_conn).drive().await?;
+            let _ = self.status_tx.send(Status::QUIC);
+        }
+
         let timer = optional_timeout(self.quiche_conn.timeout(), self.net_id);
         select! {
             // If a quiche timer would fire, call their callback
@@ -210,19 +223,6 @@ impl Driver {
         // Any of the actions in the select could require us to send packets to the peer
         self.flush_tx().await?;
 
-        // If the QUIC connection is live, but the HTTP/3 is not, try to bring it up
-        if self.quiche_conn.is_established() {
-            info!(
-                "Connection {} established on network {}",
-                self.quiche_conn.trace_id(),
-                self.net_id
-            );
-            let h3_config = h3::Config::new()?;
-            let h3_conn = h3::Connection::with_transport(&mut self.quiche_conn, &h3_config)?;
-            self = H3Driver::new(self, h3_conn).drive().await?;
-            let _ = self.status_tx.send(Status::QUIC);
-        }
-
         // If the connection has entered draining state (the server is closing the connection),
         // tell the status watcher not to use the connection. Besides, per Quiche document,
         // the connection should not be dropped until is_closed() returns true.
@@ -285,7 +285,8 @@ impl H3Driver {
         }
         select! {
             // Only attempt to enqueue new requests if we have no buffered request and aren't
-            // closing
+            // closing. Maybe limit the number of in-flight queries if the handshake
+            // still hasn't finished.
             msg = self.driver.request_rx.recv(), if !self.driver.closing && self.buffered_request.is_none() => {
                 match msg {
                     Some(request) => self.handle_request(request)?,
@@ -321,8 +322,8 @@ impl H3Driver {
     }
 
     fn handle_request(&mut self, request: Request) -> Result<()> {
-        info!("Handling DNS request on network {}, stats=[{:?}], peer_streams_left_bidi={}, peer_streams_left_uni={}",
-                self.driver.net_id, self.driver.quiche_conn.stats(), self.driver.quiche_conn.peer_streams_left_bidi(), self.driver.quiche_conn.peer_streams_left_uni());
+        info!("Handling DNS request on network {}, is_in_early_data={}, stats=[{:?}], peer_streams_left_bidi={}, peer_streams_left_uni={}",
+                self.driver.net_id, self.driver.quiche_conn.is_in_early_data(), self.driver.quiche_conn.stats(), self.driver.quiche_conn.peer_streams_left_bidi(), self.driver.quiche_conn.peer_streams_left_uni());
         // If the request has already timed out, don't issue it to the server.
         if let Some(expiry) = request.expiry {
             if BootTime::now() > expiry {
diff --git a/doh/connection/mod.rs b/doh/connection/mod.rs
index edd79cf9..8634014d 100644
--- a/doh/connection/mod.rs
+++ b/doh/connection/mod.rs
@@ -142,6 +142,8 @@ impl Connection {
         let scid = new_scid();
         let mut quiche_conn =
             quiche::connect(server_name, &quiche::ConnectionId::from_ref(&scid), to, config)?;
+
+        // We will fall back to a full handshake if the session is expired.
         if let Some(session) = session {
             debug!("Setting session");
             quiche_conn.set_session(&session)?;
diff --git a/doh/dispatcher/driver.rs b/doh/dispatcher/driver.rs
index 5e0359cb..aaea68cc 100644
--- a/doh/dispatcher/driver.rs
+++ b/doh/dispatcher/driver.rs
@@ -117,6 +117,7 @@ impl Driver {
                 let key = config::Key {
                     cert_path: info.cert_path.clone(),
                     max_idle_timeout: info.idle_timeout_ms,
+                    enable_early_data: info.enable_early_data,
                 };
                 let config = self.config_cache.get(&key)?;
                 vacant.insert(
diff --git a/doh/dispatcher/mod.rs b/doh/dispatcher/mod.rs
index 18439c8e..6dff3e14 100644
--- a/doh/dispatcher/mod.rs
+++ b/doh/dispatcher/mod.rs
@@ -87,7 +87,9 @@ impl Dispatcher {
             .build()?;
         let join_handle = runtime.spawn(async {
             let result = Driver::new(cmd_receiver, validation, tagger).drive().await;
-            if let Err(ref e) = result { error!("Dispatcher driver exited due to {:?}", e) }
+            if let Err(ref e) = result {
+                error!("Dispatcher driver exited due to {:?}", e)
+            }
             result
         });
         Ok(Dispatcher { cmd_sender, join_handle, runtime })
diff --git a/doh/ffi.rs b/doh/ffi.rs
index 65ca012e..6e9f64ad 100644
--- a/doh/ffi.rs
+++ b/doh/ffi.rs
@@ -43,6 +43,7 @@ pub struct FeatureFlags {
     probe_timeout_ms: uint64_t,
     idle_timeout_ms: uint64_t,
     use_session_resumption: bool,
+    enable_early_data: bool,
 }
 
 fn wrap_validation_callback(validation_fn: ValidationCallback) -> ValidationReporter {
@@ -233,6 +234,7 @@ pub unsafe extern "C" fn doh_net_new(
             cert_path,
             idle_timeout_ms: flags.idle_timeout_ms,
             use_session_resumption: flags.use_session_resumption,
+            enable_early_data: flags.enable_early_data,
         },
         timeout: Duration::from_millis(flags.probe_timeout_ms),
     };
@@ -384,6 +386,7 @@ mod tests {
             cert_path: None,
             idle_timeout_ms: 0,
             use_session_resumption: true,
+            enable_early_data: true,
         };
 
         wrap_validation_callback(success_cb)(&info, true).await;
diff --git a/doh/network/mod.rs b/doh/network/mod.rs
index 19be8643..7e39f60b 100644
--- a/doh/network/mod.rs
+++ b/doh/network/mod.rs
@@ -49,6 +49,7 @@ pub struct ServerInfo {
     pub cert_path: Option<String>,
     pub idle_timeout_ms: u64,
     pub use_session_resumption: bool,
+    pub enable_early_data: bool,
 }
 
 #[derive(Debug)]