From c751a21876f95d09c7a9061e4ec9f5c22ffbeeff Mon Sep 17 00:00:00 2001
From: Andrea Gottardo <andrea@gottardo.me>
Date: Thu, 3 Oct 2024 13:30:07 -0700
Subject: [PATCH] net/dns: close idle DoH connections when entering sleep mode

Updates tailscale/tailscale#3363
Updates tailscale/tailscale#6148

Provides a facility for the iOS code (and later Android) to signal the beginning of device-wide sleep mode to the LocalBackend, and wires it up to the DNS forwarder, to early-close any open DoH connections when the device is about to enter sleep mode (we don't want a single TCP keepalive to wake up the device again seconds later).

Signed-off-by: Andrea Gottardo <andrea@gottardo.me>
---
 ipn/ipnlocal/local.go         | 11 +++++++++++
 net/dns/manager.go            |  8 ++++++++
 net/dns/resolver/forwarder.go | 12 ++++++++++++
 net/dns/resolver/tsdns.go     |  8 ++++++++
 4 files changed, 39 insertions(+)

diff --git a/ipn/ipnlocal/local.go b/ipn/ipnlocal/local.go
index 8fc78a36b..831a98349 100644
--- a/ipn/ipnlocal/local.go
+++ b/ipn/ipnlocal/local.go
@@ -2146,6 +2146,17 @@ func (b *LocalBackend) Start(opts ipn.Options) error {
 	return nil
 }
 
+// OnSleepChange is invoked by the system when transitioning into or out of sleep mode.
+// This function is used to pause or resume non-essential network activity during sleep,
+// with an eye towards power savings.
+func (b *LocalBackend) OnSleepChange(isSleeping bool) {
+	b.logf("OnSleepChange(isSleeping=%v)", isSleeping)
+	dnsManager, ok := b.sys.DNSManager.GetOK()
+	if ok {
+		dnsManager.OnSleepChange(isSleeping)
+	}
+}
+
 // invalidPacketFilterWarnable is a Warnable to warn the user that the control server sent an invalid packet filter.
 var invalidPacketFilterWarnable = health.Register(&health.Warnable{
 	Code:     "invalid-packet-filter",
diff --git a/net/dns/manager.go b/net/dns/manager.go
index 51a0fa12c..fe76e35a9 100644
--- a/net/dns/manager.go
+++ b/net/dns/manager.go
@@ -127,6 +127,14 @@ func (m *Manager) GetBaseConfig() (OSConfig, error) {
 	return m.os.GetBaseConfig()
 }
 
+// OnSleepChange is called by the backend when the device enters or leaves sleep mode.
+// We use this to trigger behaviors needed to provide battery savings.
+func (m *Manager) OnSleepChange(isSleeping bool) {
+	if isSleeping {
+		m.resolver.OnSleepChange(isSleeping)
+	}
+}
+
 // setLocked sets the DNS configuration.
 //
 // m.mu must be held.
diff --git a/net/dns/resolver/forwarder.go b/net/dns/resolver/forwarder.go
index 846ca3d5e..9dca1155c 100644
--- a/net/dns/resolver/forwarder.go
+++ b/net/dns/resolver/forwarder.go
@@ -273,6 +273,18 @@ func (f *forwarder) Close() error {
 	return nil
 }
 
+// CloseIdleConnections closes any idle connections to the upstream
+// DoH servers. It is desirable to call this when the device enters
+// sleep mode, when we know that no DNS queries will be made for a
+// while, so that we can preserve battery life.
+func (f *forwarder) CloseIdleConnections() {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	for _, c := range f.dohClient {
+		c.Transport.(*http.Transport).CloseIdleConnections()
+	}
+}
+
 // resolversWithDelays maps from a set of DNS server names to a slice of a type
 // that included a startDelay, upgrading any well-known DoH (DNS-over-HTTP)
 // servers in the process, insert a DoH lookup first before UDP fallbacks.
diff --git a/net/dns/resolver/tsdns.go b/net/dns/resolver/tsdns.go
index d196ad4d6..f9dc61306 100644
--- a/net/dns/resolver/tsdns.go
+++ b/net/dns/resolver/tsdns.go
@@ -260,6 +260,14 @@ func (r *Resolver) SetMissingUpstreamRecovery(f func()) {
 	r.forwarder.missingUpstreamRecovery = f
 }
 
+// OnSleepChange asks the forwarder to close any idle connections to the upstream
+// DoH servers when the device enters sleep mode.
+func (r *Resolver) OnSleepChange(isSleeping bool) {
+	if isSleeping {
+		r.forwarder.CloseIdleConnections()
+	}
+}
+
 func (r *Resolver) TestOnlySetHook(hook func(Config)) { r.saveConfigForTests = hook }
 
 func (r *Resolver) SetConfig(cfg Config) error {