From 6196b7e658a85892a16e9b3b269501ec7169fc6c Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 6 Jul 2020 12:10:39 -0700
Subject: [PATCH] wgengine/magicsock: change API to not permit disco key
 changes

Generate the disco key ourselves and give out the public half instead.

Fixes #525
---
 ipn/local.go                         |  4 +---
 wgengine/magicsock/magicsock.go      | 21 ++++++++++-----------
 wgengine/magicsock/magicsock_test.go |  7 +++----
 wgengine/userspace.go                |  4 ++--
 wgengine/watchdog.go                 |  6 +++---
 wgengine/wgengine.go                 |  5 ++---
 6 files changed, 21 insertions(+), 26 deletions(-)

diff --git a/ipn/local.go b/ipn/local.go
index 3ba92186a..6571d84cc 100644
--- a/ipn/local.go
+++ b/ipn/local.go
@@ -361,9 +361,7 @@ func (b *LocalBackend) Start(opts Options) error {
 
 	var discoPublic tailcfg.DiscoKey
 	if controlclient.Debug.Disco {
-		discoPrivate := key.NewPrivate()
-		b.e.SetDiscoPrivateKey(discoPrivate)
-		discoPublic = tailcfg.DiscoKey(discoPrivate.Public())
+		discoPublic = b.e.DiscoPublicKey()
 	}
 
 	var err error
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 4c08f144b..deef60953 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -504,19 +504,18 @@ func (c *Conn) SetNetInfoCallback(fn func(*tailcfg.NetInfo)) {
 	}
 }
 
-// SetDiscoPrivateKey sets the discovery key.
-func (c *Conn) SetDiscoPrivateKey(k key.Private) {
+// DiscoPublicKey returns the discovery public key.
+func (c *Conn) DiscoPublicKey() tailcfg.DiscoKey {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	if !c.discoPrivate.IsZero() && c.discoPrivate != k {
-		// TODO: support changing a key at runtime; need to
-		// clear a bunch of maps at least
-		panic("unsupported")
-	}
-	c.discoPrivate = k
-	c.discoPublic = tailcfg.DiscoKey(k.Public())
-	c.discoShort = c.discoPublic.ShortString()
-	c.logf("magicsock: set disco key = %v", c.discoShort)
+	if c.discoPrivate.IsZero() {
+		priv := key.NewPrivate()
+		c.discoPrivate = priv
+		c.discoPublic = tailcfg.DiscoKey(priv.Public())
+		c.discoShort = c.discoPublic.ShortString()
+		c.logf("magicsock: disco key = %v", c.discoShort)
+	}
+	return c.discoPublic
 }
 
 // c.mu must NOT be held.
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index b5e7d5c55..a090f900b 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -855,12 +855,11 @@ func initAddrSet(as *AddrSet) {
 }
 
 func TestDiscoMessage(t *testing.T) {
-	peer1Priv := key.NewPrivate()
-	peer1Pub := peer1Priv.Public()
-
 	c := newConn()
 	c.logf = t.Logf
-	c.SetDiscoPrivateKey(key.NewPrivate())
+
+	peer1Pub := c.DiscoPublicKey()
+	peer1Priv := c.discoPrivate
 	c.endpointOfDisco = map[tailcfg.DiscoKey]*discoEndpoint{
 		tailcfg.DiscoKey(peer1Pub): &discoEndpoint{
 			// ...
diff --git a/wgengine/userspace.go b/wgengine/userspace.go
index 11e181d4b..46b5cac65 100644
--- a/wgengine/userspace.go
+++ b/wgengine/userspace.go
@@ -827,8 +827,8 @@ func (e *userspaceEngine) SetNetworkMap(nm *controlclient.NetworkMap) {
 	e.magicConn.SetNetworkMap(nm)
 }
 
-func (e *userspaceEngine) SetDiscoPrivateKey(k key.Private) {
-	e.magicConn.SetDiscoPrivateKey(k)
+func (e *userspaceEngine) DiscoPublicKey() tailcfg.DiscoKey {
+	return e.magicConn.DiscoPublicKey()
 }
 
 func (e *userspaceEngine) UpdateStatus(sb *ipnstate.StatusBuilder) {
diff --git a/wgengine/watchdog.go b/wgengine/watchdog.go
index 4e39d83bf..60826ac43 100644
--- a/wgengine/watchdog.go
+++ b/wgengine/watchdog.go
@@ -14,7 +14,6 @@ import (
 	"tailscale.com/control/controlclient"
 	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/tailcfg"
-	"tailscale.com/types/key"
 	"tailscale.com/wgengine/filter"
 	"tailscale.com/wgengine/router"
 	"tailscale.com/wgengine/tsdns"
@@ -101,8 +100,9 @@ func (e *watchdogEngine) SetDERPMap(m *tailcfg.DERPMap) {
 func (e *watchdogEngine) SetNetworkMap(nm *controlclient.NetworkMap) {
 	e.watchdog("SetNetworkMap", func() { e.wrap.SetNetworkMap(nm) })
 }
-func (e *watchdogEngine) SetDiscoPrivateKey(k key.Private) {
-	e.watchdog("SetDiscoPrivateKey", func() { e.wrap.SetDiscoPrivateKey(k) })
+func (e *watchdogEngine) DiscoPublicKey() (k tailcfg.DiscoKey) {
+	e.watchdog("DiscoPublicKey", func() { k = e.wrap.DiscoPublicKey() })
+	return k
 }
 func (e *watchdogEngine) Close() {
 	e.watchdog("Close", e.wrap.Close)
diff --git a/wgengine/wgengine.go b/wgengine/wgengine.go
index 8a1aaa833..da6e67239 100644
--- a/wgengine/wgengine.go
+++ b/wgengine/wgengine.go
@@ -12,7 +12,6 @@ import (
 	"tailscale.com/control/controlclient"
 	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/tailcfg"
-	"tailscale.com/types/key"
 	"tailscale.com/wgengine/filter"
 	"tailscale.com/wgengine/router"
 	"tailscale.com/wgengine/tsdns"
@@ -117,9 +116,9 @@ type Engine interface {
 	// new NetInfo summary is available.
 	SetNetInfoCallback(NetInfoCallback)
 
-	// SetDiscoPrivateKey sets the private key used for path discovery
+	// DiscoPublicKey gets the public key used for path discovery
 	// messages.
-	SetDiscoPrivateKey(key.Private)
+	DiscoPublicKey() tailcfg.DiscoKey
 
 	// UpdateStatus populates the network state using the provided
 	// status builder.