Injecting handshake relevant statistics into statsd.

|sampling_rate_denom| and |round_trips| are not handled in this change.

event_metrics {
  data {
    aggregated_atom_info {
      atom {
        network_dns_handshake_reported {
          protocol: PROTO_DOH
          result: HR_SUCCESS
          cause: HC_SERVER_PROBE
          network_type: NT_WIFI
          private_dns_mode: PDM_STRICT
          latency_micros: 42596
          bytes_sent: 761
          bytes_received: 6420
          round_trips: -1
          tls_session_cache_hit: false
          tls_version: 3
          hostname_verification: true
          quic_version: 1
          server_index: -1
          sampling_rate_denom: -1
        }
      }
      elapsed_timestamp_nanos: 195891466159
    }
  }
}

Bug: 249206898
Test: atest resolv_integration_test
Change-Id: I312b61e224722e6bc90e29a05bc2103ebbd6dd73

13 files changed, 309 insertions, 33 deletions

diff --git a/Android.bp b/Android.bp
index 11d05970..20d38b51 100644
--- a/Android.bp
+++ b/Android.bp
@@ -236,15 +236,11 @@ cc_library {
         "libcrypto",
         "liblog", //Used by libstatslog_resolv
         "libssl",
+        "libstatssocket",
     ],
     header_libs: [
         "libnetdbinder_utils_headers",
     ],
-    runtime_libs: [
-        // Causes the linkerconfig to create a namespace link from resolv to the
-        // libstatssocket library within the statsd apex
-        "libstatssocket",
-    ],
     export_include_dirs: ["include"],
 
     product_variables: {
@@ -349,6 +345,7 @@ doh_rust_deps = [
     "liblibc",
     "liblog_rust",
     "libring",
+    "libstatslog_rust",
     "libthiserror",
     "libtokio",
     "liburl",
@@ -395,6 +392,11 @@ rust_ffi_static {
 
     rlibs: doh_rust_deps + ["libquiche_static"],
     prefer_rlib: true,
+
+    shared_libs: [
+        "libstatssocket",
+    ],
+
     // TODO(b/194022174), for unit tests to run on the Android 10 platform,
     // libunwind must be statically linked.
     whole_static_libs: ["libunwind"],
@@ -425,6 +427,7 @@ rust_ffi_static {
         "liblog_rust",
         "libquiche_static",
         "libring",
+        "libstatslog_rust",
         "libthiserror",
         "libtokio",
         "liburl",
diff --git a/PrivateDnsConfiguration.cpp b/PrivateDnsConfiguration.cpp
index 08576999..a7a8d9c3 100644
--- a/PrivateDnsConfiguration.cpp
+++ b/PrivateDnsConfiguration.cpp
@@ -568,6 +568,9 @@ int PrivateDnsConfiguration::setDoh(int32_t netId, uint32_t mark,
         return 0;
     }
 
+    const NetworkType networkType = resolv_get_network_types_for_net(netId);
+    const PrivateDnsStatus status = getStatusLocked(netId);
+
     const auto getTimeoutFromFlag = [&](const std::string_view key, int defaultValue) -> uint64_t {
         static constexpr int kMinTimeoutMs = 1000;
         uint64_t timeout = Experiments::getInstance()->getFlag(key, defaultValue);
@@ -628,8 +631,10 @@ int PrivateDnsConfiguration::setDoh(int32_t netId, uint32_t mark,
                    << ", use_session_resumption=" << flags.use_session_resumption
                    << ", enable_early_data=" << flags.enable_early_data;
 
+        const PrivateDnsModes privateDnsMode = convertEnumType(status.mode);
         return doh_net_new(mDohDispatcher, netId, dohId.httpsTemplate.c_str(), dohId.host.c_str(),
-                           dohId.ipAddr.c_str(), mark, caCert.c_str(), &flags);
+                           dohId.ipAddr.c_str(), mark, caCert.c_str(), &flags, networkType,
+                           privateDnsMode);
     }
 
     LOG(INFO) << __func__ << ": No suitable DoH server found";
diff --git a/doh.h b/doh.h
index 4fa8a5d7..6f82bea5 100644
--- a/doh.h
+++ b/doh.h
@@ -92,7 +92,7 @@ void doh_dispatcher_delete(DohDispatcher* doh);
 /// `url`, `domain`, `ip_addr`, `cert_path` are null terminated strings.
 int32_t doh_net_new(DohDispatcher* doh, uint32_t net_id, const char* url, const char* domain,
                     const char* ip_addr, uint32_t sk_mark, const char* cert_path,
-                    const FeatureFlags* flags);
+                    const FeatureFlags* flags, uint32_t network_type, uint32_t private_dns_mode);
 
 /// Sends a DNS query via the network associated to the given |net_id| and waits for the response.
 /// The return code should be either one of the public constant RESULT_* to indicate the error or
diff --git a/doh/connection/driver.rs b/doh/connection/driver.rs
index 315fc130..6fb86b4e 100644
--- a/doh/connection/driver.rs
+++ b/doh/connection/driver.rs
@@ -17,12 +17,14 @@
 
 use crate::boot_time;
 use crate::boot_time::BootTime;
+use crate::metrics::log_handshake_event_stats;
 use log::{debug, info, warn};
 use quiche::h3;
 use std::collections::HashMap;
 use std::default::Default;
 use std::future;
 use std::io;
+use std::time::Instant;
 use thiserror::Error;
 use tokio::net::UdpSocket;
 use tokio::select;
@@ -30,6 +32,50 @@ use tokio::sync::{mpsc, oneshot, watch};
 
 use super::Status;
 
+#[derive(Copy, Clone, Debug)]
+pub enum Cause {
+    Probe,
+    Reconnect,
+    Retry,
+}
+
+#[derive(Clone)]
+#[allow(dead_code)]
+pub enum HandshakeResult {
+    Unknown,
+    Success,
+    Timeout,
+    TlsFail,
+    ServerUnreachable,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct HandshakeInfo {
+    pub cause: Cause,
+    pub sent_bytes: u64,
+    pub recv_bytes: u64,
+    pub elapsed: u128,
+    pub quic_version: u32,
+    pub network_type: u32,
+    pub private_dns_mode: u32,
+    pub session_hit_checker: bool,
+}
+
+impl std::fmt::Display for HandshakeInfo {
+    #[inline]
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(
+            f,
+            "cause={:?}, sent_bytes={}, recv_bytes={}, quic_version={}, session_hit_checker={}",
+            self.cause,
+            self.sent_bytes,
+            self.recv_bytes,
+            self.quic_version,
+            self.session_hit_checker
+        )
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum Error {
     #[error("network IO error: {0}")]
@@ -92,6 +138,8 @@ struct Driver {
     // if we poll on a dead receiver in a select! it will immediately return None. As a result, we
     // need this to gate whether or not to include .recv() in our select!
     closing: bool,
+    handshake_info: HandshakeInfo,
+    connection_start: Instant,
 }
 
 struct H3Driver {
@@ -121,8 +169,9 @@ pub async fn drive(
     quiche_conn: quiche::Connection,
     socket: UdpSocket,
     net_id: u32,
+    handshake_info: HandshakeInfo,
 ) -> Result<()> {
-    Driver::new(request_rx, status_tx, quiche_conn, socket, net_id).drive().await
+    Driver::new(request_rx, status_tx, quiche_conn, socket, net_id, handshake_info).drive().await
 }
 
 impl Driver {
@@ -132,6 +181,7 @@ impl Driver {
         quiche_conn: quiche::Connection,
         socket: UdpSocket,
         net_id: u32,
+        handshake_info: HandshakeInfo,
     ) -> Self {
         Self {
             request_rx,
@@ -141,10 +191,13 @@ impl Driver {
             buffer: Box::new([0; MAX_UDP_PACKET_SIZE]),
             net_id,
             closing: false,
+            handshake_info,
+            connection_start: Instant::now(),
         }
     }
 
     async fn drive(mut self) -> Result<()> {
+        self.connection_start = Instant::now();
         // Prime connection
         self.flush_tx().await?;
         loop {
@@ -202,6 +255,13 @@ impl Driver {
                 self.quiche_conn.trace_id(),
                 self.net_id
             );
+            self.handshake_info.elapsed = self.connection_start.elapsed().as_micros();
+            // In Stats, sent_bytes implements the way that omits the length of padding data
+            // append to the datagram.
+            self.handshake_info.sent_bytes = self.quiche_conn.stats().sent_bytes;
+            self.handshake_info.recv_bytes = self.quiche_conn.stats().recv_bytes;
+            self.handshake_info.quic_version = quiche::PROTOCOL_VERSION;
+            log_handshake_event_stats(HandshakeResult::Success, self.handshake_info);
             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?;
@@ -213,7 +273,20 @@ impl Driver {
             // If a quiche timer would fire, call their callback
             _ = timer => {
                 info!("Driver: Timer expired on network {}", self.net_id);
-                self.quiche_conn.on_timeout()
+                self.quiche_conn.on_timeout();
+
+                if !self.quiche_conn.is_established() && self.quiche_conn.is_closed() {
+                    info!(
+                        "Connection {} timeouted on network {}",
+                        self.quiche_conn.trace_id(),
+                        self.net_id
+                    );
+                    self.handshake_info.elapsed = self.connection_start.elapsed().as_micros();
+                    log_handshake_event_stats(
+                        HandshakeResult::Timeout,
+                        self.handshake_info,
+                    );
+                }
             }
             // If we got packets from our peer, pass them to quiche
             Ok((size, from)) = self.socket.recv_from(self.buffer.as_mut()) => {
@@ -221,6 +294,7 @@ impl Driver {
                 debug!("Received {} bytes on network {}", size, self.net_id);
             }
         };
+
         // Any of the actions in the select could require us to send packets to the peer
         self.flush_tx().await?;
 
diff --git a/doh/connection/mod.rs b/doh/connection/mod.rs
index 8634014d..73a125f3 100644
--- a/doh/connection/mod.rs
+++ b/doh/connection/mod.rs
@@ -16,6 +16,9 @@
 //! Module providing an async abstraction around a quiche HTTP/3 connection
 
 use crate::boot_time::BootTime;
+use crate::connection::driver::Cause;
+use crate::connection::driver::HandshakeInfo;
+use crate::network::ServerInfo;
 use crate::network::SocketTagger;
 use log::{debug, error, warn};
 use quiche::h3;
@@ -27,7 +30,7 @@ use tokio::net::UdpSocket;
 use tokio::sync::{mpsc, oneshot, watch};
 use tokio::task;
 
-mod driver;
+pub mod driver;
 
 pub use driver::Stream;
 use driver::{drive, Request};
@@ -129,29 +132,40 @@ impl Connection {
     const MAX_PENDING_REQUESTS: usize = 10;
     /// Create a new connection with a background task handling IO.
     pub async fn new(
-        server_name: Option<&str>,
-        to: SocketAddr,
-        socket_mark: u32,
-        net_id: u32,
+        info: &ServerInfo,
         tag_socket: &SocketTagger,
         config: &mut quiche::Config,
         session: Option<Vec<u8>>,
+        cause: Cause,
     ) -> Result<Self> {
+        let server_name = info.domain.as_deref();
+        let to = info.peer_addr;
+        let socket_mark = info.sk_mark;
+        let net_id = info.net_id;
         let (request_tx, request_rx) = mpsc::channel(Self::MAX_PENDING_REQUESTS);
         let (status_tx, status_rx) = watch::channel(Status::QUIC);
         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)?;
         }
-
+        let handshake_info = HandshakeInfo {
+            cause,
+            sent_bytes: 0,
+            recv_bytes: 0,
+            elapsed: 0,
+            quic_version: 0,
+            network_type: info.network_type,
+            private_dns_mode: info.private_dns_mode,
+            session_hit_checker: quiche_conn.session().is_some(),
+        };
         let socket = build_socket(to, socket_mark, tag_socket).await?;
         let driver = async move {
-            let result = drive(request_rx, status_tx, quiche_conn, socket, net_id).await;
+            let result =
+                drive(request_rx, status_tx, quiche_conn, socket, net_id, handshake_info).await;
             if let Err(ref e) = result {
                 warn!("Connection driver returns some Err: {:?}", e);
             }
diff --git a/doh/doh.rs b/doh/doh.rs
index 5ceb5e81..42152a2c 100644
--- a/doh/doh.rs
+++ b/doh/doh.rs
@@ -22,4 +22,5 @@ mod connection;
 mod dispatcher;
 mod encoding;
 mod ffi;
+mod metrics;
 mod network;
diff --git a/doh/doh_test_superset_for_fuzzer.rs b/doh/doh_test_superset_for_fuzzer.rs
index 93c5985e..60315127 100644
--- a/doh/doh_test_superset_for_fuzzer.rs
+++ b/doh/doh_test_superset_for_fuzzer.rs
@@ -23,6 +23,7 @@ mod connection;
 mod dispatcher;
 mod encoding;
 mod ffi;
+mod metrics;
 mod network;
 /// The Rust FFI bindings to C APIs for implementation of doh frontend.
 mod tests;
diff --git a/doh/ffi.rs b/doh/ffi.rs
index 6e9f64ad..e29c2152 100644
--- a/doh/ffi.rs
+++ b/doh/ffi.rs
@@ -190,6 +190,8 @@ pub unsafe extern "C" fn doh_net_new(
     sk_mark: libc::uint32_t,
     cert_path: *const c_char,
     flags: &FeatureFlags,
+    network_type: uint32_t,
+    private_dns_mode: uint32_t,
 ) -> int32_t {
     let (url, domain, ip_addr, cert_path) = match (
         std::ffi::CStr::from_ptr(url).to_str(),
@@ -235,6 +237,8 @@ pub unsafe extern "C" fn doh_net_new(
             idle_timeout_ms: flags.idle_timeout_ms,
             use_session_resumption: flags.use_session_resumption,
             enable_early_data: flags.enable_early_data,
+            network_type,
+            private_dns_mode,
         },
         timeout: Duration::from_millis(flags.probe_timeout_ms),
     };
@@ -387,6 +391,8 @@ mod tests {
             idle_timeout_ms: 0,
             use_session_resumption: true,
             enable_early_data: true,
+            network_type: 2,
+            private_dns_mode: 3,
         };
 
         wrap_validation_callback(success_cb)(&info, true).await;
diff --git a/doh/metrics.rs b/doh/metrics.rs
new file mode 100644
index 00000000..9b7a96b1
--- /dev/null
+++ b/doh/metrics.rs
@@ -0,0 +1,157 @@
+// Copyright 2022, 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.
+
+//! This module provides convenience functions for doh logging.
+
+use crate::connection::driver::Cause;
+use crate::connection::driver::HandshakeInfo;
+use crate::connection::driver::HandshakeResult;
+use statslog_rust::network_dns_handshake_reported::{
+    Cause as StatsdCause, NetworkDnsHandshakeReported, NetworkType as StatsdNetworkType,
+    PrivateDnsMode as StatsdPrivateDnsMode, Protocol as StatsdProtocol, Result as StatsdResult,
+};
+
+const CELLULAR: u32 = 1;
+const WIFI: u32 = 2;
+const BLUETOOTH: u32 = 3;
+const ETHERNET: u32 = 4;
+const VPN: u32 = 5;
+const WIFI_AWARE: u32 = 6;
+const LOWPAN: u32 = 7;
+const CELLULAR_VPN: u32 = 8;
+const WIFI_VPN: u32 = 9;
+const BLUETOOTH_VPN: u32 = 10;
+const ETHERNET_VPN: u32 = 11;
+const WIFI_CELLULAR_VPN: u32 = 12;
+
+const OFF: u32 = 1;
+const OPPORTUNISTIC: u32 = 2;
+const STRICT: u32 = 3;
+
+const TLS1_3_VERSION: u32 = 3;
+
+fn create_default_handshake_atom() -> NetworkDnsHandshakeReported {
+    NetworkDnsHandshakeReported {
+        protocol: StatsdProtocol::ProtoUnknown,
+        result: StatsdResult::HrUnknown,
+        cause: StatsdCause::HcUnknown,
+        network_type: StatsdNetworkType::NtUnknown,
+        private_dns_mode: StatsdPrivateDnsMode::PdmUnknown,
+        latency_micros: -1,
+        bytes_sent: -1,
+        bytes_received: -1,
+        round_trips: -1,
+        tls_session_cache_hit: false,
+        tls_version: -1,
+        hostname_verification: false,
+        quic_version: -1,
+        server_index: -1,
+        sampling_rate_denom: -1,
+    }
+}
+
+fn construct_handshake_event_stats(
+    result: HandshakeResult,
+    handshake_info: HandshakeInfo,
+) -> NetworkDnsHandshakeReported {
+    let mut handshake_event_atom = create_default_handshake_atom();
+    handshake_event_atom.protocol = StatsdProtocol::ProtoDoh;
+    handshake_event_atom.result = match result {
+        HandshakeResult::Success => StatsdResult::HrSuccess,
+        HandshakeResult::Timeout => StatsdResult::HrTimeout,
+        _ => StatsdResult::HrUnknown,
+    };
+    handshake_event_atom.cause = match handshake_info.cause {
+        Cause::Probe => StatsdCause::HcServerProbe,
+        Cause::Reconnect => StatsdCause::HcReconnectAfterIdle,
+        Cause::Retry => StatsdCause::HcRetryAfterError,
+    };
+    handshake_event_atom.network_type = match handshake_info.network_type {
+        CELLULAR => StatsdNetworkType::NtCellular,
+        WIFI => StatsdNetworkType::NtWifi,
+        BLUETOOTH => StatsdNetworkType::NtBluetooth,
+        ETHERNET => StatsdNetworkType::NtEthernet,
+        VPN => StatsdNetworkType::NtVpn,
+        WIFI_AWARE => StatsdNetworkType::NtWifiAware,
+        LOWPAN => StatsdNetworkType::NtLowpan,
+        CELLULAR_VPN => StatsdNetworkType::NtCellularVpn,
+        WIFI_VPN => StatsdNetworkType::NtWifiVpn,
+        BLUETOOTH_VPN => StatsdNetworkType::NtBluetoothVpn,
+        ETHERNET_VPN => StatsdNetworkType::NtEthernetVpn,
+        WIFI_CELLULAR_VPN => StatsdNetworkType::NtWifiCellularVpn,
+        _ => StatsdNetworkType::NtUnknown,
+    };
+    handshake_event_atom.private_dns_mode = match handshake_info.private_dns_mode {
+        OFF => StatsdPrivateDnsMode::PdmOff,
+        OPPORTUNISTIC => StatsdPrivateDnsMode::PdmOpportunistic,
+        STRICT => StatsdPrivateDnsMode::PdmStrict,
+        _ => StatsdPrivateDnsMode::PdmUnknown,
+    };
+    handshake_event_atom.latency_micros = handshake_info.elapsed as i32;
+    handshake_event_atom.bytes_sent = handshake_info.sent_bytes as i32;
+    handshake_event_atom.bytes_received = handshake_info.recv_bytes as i32;
+    handshake_event_atom.tls_session_cache_hit = handshake_info.session_hit_checker;
+    handshake_event_atom.tls_version = TLS1_3_VERSION as i32;
+    handshake_event_atom.hostname_verification = matches!(handshake_info.private_dns_mode, STRICT);
+    handshake_event_atom.quic_version = handshake_info.quic_version as i32;
+    handshake_event_atom
+}
+
+/// Log hankshake events via statsd API.
+pub fn log_handshake_event_stats(result: HandshakeResult, handshake_info: HandshakeInfo) {
+    let handshake_event_stats = construct_handshake_event_stats(result, handshake_info);
+
+    let logging_result = handshake_event_stats.stats_write();
+    if let Err(e) = logging_result {
+        log::error!("Error in logging handshake event. {:?}", e);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_metrics_write() {
+        let handshake_info = HandshakeInfo {
+            cause: Cause::Retry,
+            network_type: WIFI,
+            private_dns_mode: STRICT,
+            elapsed: 42596,
+            sent_bytes: 761,
+            recv_bytes: 6420,
+            session_hit_checker: false,
+            quic_version: 1,
+        };
+        let result = HandshakeResult::Timeout;
+        let handshake_event_stats = construct_handshake_event_stats(result, handshake_info);
+        assert_eq!(handshake_event_stats.protocol as i32, StatsdProtocol::ProtoDoh as i32);
+        assert_eq!(handshake_event_stats.result as i32, HandshakeResult::Timeout as i32);
+        assert_eq!(handshake_event_stats.cause as i32, StatsdCause::HcRetryAfterError as i32);
+        assert_eq!(handshake_event_stats.network_type as i32, StatsdNetworkType::NtWifi as i32);
+        assert_eq!(
+            handshake_event_stats.private_dns_mode as i32,
+            StatsdPrivateDnsMode::PdmStrict as i32
+        );
+        assert_eq!(handshake_event_stats.latency_micros, 42596);
+        assert_eq!(handshake_event_stats.bytes_sent, 761);
+        assert_eq!(handshake_event_stats.bytes_received, 6420);
+        assert_eq!(handshake_event_stats.round_trips, -1);
+        assert!(!handshake_event_stats.tls_session_cache_hit);
+        assert!(handshake_event_stats.hostname_verification);
+        assert_eq!(handshake_event_stats.quic_version, 1);
+        assert_eq!(handshake_event_stats.server_index, -1);
+        assert_eq!(handshake_event_stats.sampling_rate_denom, -1);
+    }
+}
diff --git a/doh/network/driver.rs b/doh/network/driver.rs
index 118e5758..cad5584e 100644
--- a/doh/network/driver.rs
+++ b/doh/network/driver.rs
@@ -18,6 +18,7 @@
 
 use crate::boot_time::{timeout, BootTime, Duration};
 use crate::config::Config;
+use crate::connection::driver::Cause;
 use crate::connection::Connection;
 use crate::dispatcher::{QueryError, Response};
 use crate::encoding;
@@ -76,18 +77,10 @@ async fn build_connection(
     tag_socket: &SocketTagger,
     config: &mut Config,
     session: Option<Vec<u8>>,
+    cause: Cause,
 ) -> Result<Connection> {
     use std::ops::DerefMut;
-    Ok(Connection::new(
-        info.domain.as_deref(),
-        info.peer_addr,
-        info.sk_mark,
-        info.net_id,
-        tag_socket,
-        config.take().await.deref_mut(),
-        session,
-    )
-    .await?)
+    Ok(Connection::new(info, tag_socket, config.take().await.deref_mut(), session, cause).await?)
 }
 
 impl Driver {
@@ -101,7 +94,8 @@ impl Driver {
     ) -> Result<(Self, mpsc::Sender<Command>, watch::Receiver<Status>)> {
         let (command_tx, command_rx) = mpsc::channel(Self::MAX_BUFFERED_COMMANDS);
         let (status_tx, status_rx) = watch::channel(Status::Unprobed);
-        let connection = build_connection(&info, &tag_socket, &mut config, None).await?;
+        let connection =
+            build_connection(&info, &tag_socket, &mut config, None, Cause::Probe).await?;
         Ok((
             Self { info, config, connection, status_tx, command_rx, validation, tag_socket },
             command_tx,
@@ -132,8 +126,14 @@ impl Driver {
             debug!("Network is currently failed, reconnecting");
             // If our network is currently failed, it may be due to issues with the connection.
             // Re-establish before re-probing
-            self.connection =
-                build_connection(&self.info, &self.tag_socket, &mut self.config, None).await?;
+            self.connection = build_connection(
+                &self.info,
+                &self.tag_socket,
+                &mut self.config,
+                None,
+                Cause::Retry,
+            )
+            .await?;
             self.status_tx.send(Status::Unprobed)?;
         }
         if self.status_tx.borrow().is_live() {
@@ -189,8 +189,14 @@ impl Driver {
             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?;
+            self.connection = build_connection(
+                &self.info,
+                &self.tag_socket,
+                &mut self.config,
+                session,
+                Cause::Reconnect,
+            )
+            .await?;
         }
         let request = encoding::dns_request(&query.query, &self.info.url)?;
         let stream_fut = self.connection.query(request, Some(query.expiry)).await?;
diff --git a/doh/network/mod.rs b/doh/network/mod.rs
index 7e39f60b..6d9fbab6 100644
--- a/doh/network/mod.rs
+++ b/doh/network/mod.rs
@@ -50,6 +50,8 @@ pub struct ServerInfo {
     pub idle_timeout_ms: u64,
     pub use_session_resumption: bool,
     pub enable_early_data: bool,
+    pub network_type: u32,
+    pub private_dns_mode: u32,
 }
 
 #[derive(Debug)]
diff --git a/tests/Android.bp b/tests/Android.bp
index f3ec842c..a5ae57b6 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -102,6 +102,7 @@ cc_test {
     ],
     shared_libs: [
         "libbinder_ndk",
+        "libstatssocket",
     ],
     static_libs: [
         "dnsresolver_aidl_interface-lateststable-ndk",
@@ -245,6 +246,7 @@ cc_test {
     ],
     shared_libs: [
         "libbinder_ndk",
+        "libstatssocket",
     ],
     static_libs: [
         "dnsresolver_aidl_interface-lateststable-ndk",
@@ -326,9 +328,11 @@ cc_test {
         "libgmock",
         "libnetdutils",
         "libssl",
+        "stats_proto",
     ],
     shared_libs: [
         "libnetd_client",
+        "libstatssocket",
     ],
 }
 
@@ -348,6 +352,7 @@ cc_defaults {
     ],
     shared_libs: [
         "libbinder_ndk",
+        "libstatssocket",
     ],
     static_libs: [
         "dnsresolver_aidl_interface-lateststable-ndk",
diff --git a/tests/doh_ffi_test.cpp b/tests/doh_ffi_test.cpp
index b91c59ed..0e51402a 100644
--- a/tests/doh_ffi_test.cpp
+++ b/tests/doh_ffi_test.cpp
@@ -109,7 +109,9 @@ TEST_F(DoHFFITest, SmokeTest) {
     // sk_mark doesn't matter here because this test doesn't have permission to set sk_mark.
     // The DNS packet would be sent via default network.
     EXPECT_EQ(doh_net_new(doh, dnsNetId, "https://dns.google/dns-query", /* domain */ "", server_ip,
-                          /* sk_mark */ 0, /* cert_path */ "", &flags),
+                          /* sk_mark */ 0, /* cert_path */ "", &flags,
+                          /* NetworkType::NT_WIFI */ 3,
+                          /* PrivateDnsMode::STRICT */ 2),
               0);
     {
         std::unique_lock<std::mutex> lk(m);