From e66d4e4c8131f89e85f372f6f41ff9580fedf140 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 24 May 2021 14:54:47 -0700
Subject: [PATCH] tailcfg, types/wgkey: add AppendTo methods on some types

Add MarshalText-like appending variants. Like:
https://pkg.go.dev/inet.af/netaddr#IP.AppendTo

To be used by @josharian's pending deephash optimizations.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 tailcfg/tailcfg.go      | 12 +++++++++---
 tailcfg/tailcfg_test.go | 23 +++++++++++++++++++++++
 types/wgkey/key.go      | 24 ++++++++++++++----------
 3 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/tailcfg/tailcfg.go b/tailcfg/tailcfg.go
index f1a9e3749..87836d10c 100644
--- a/tailcfg/tailcfg.go
+++ b/tailcfg/tailcfg.go
@@ -1026,11 +1026,16 @@ func (k MachineKey) MarshalText() ([]byte, error)     { return keyMarshalText("m
 func (k MachineKey) HexString() string                { return fmt.Sprintf("%x", k[:]) }
 func (k *MachineKey) UnmarshalText(text []byte) error { return keyUnmarshalText(k[:], "mkey:", text) }
 
-func keyMarshalText(prefix string, k [32]byte) []byte {
-	buf := make([]byte, len(prefix)+64)
+func appendKey(base []byte, prefix string, k [32]byte) []byte {
+	ret := append(base, make([]byte, len(prefix)+64)...)
+	buf := ret[len(base):]
 	copy(buf, prefix)
 	hex.Encode(buf[len(prefix):], k[:])
-	return buf
+	return ret
+}
+
+func keyMarshalText(prefix string, k [32]byte) []byte {
+	return appendKey(nil, prefix, k)
 }
 
 func keyUnmarshalText(dst []byte, prefix string, text []byte) error {
@@ -1061,6 +1066,7 @@ func (k DiscoKey) String() string                   { return fmt.Sprintf("discok
 func (k DiscoKey) MarshalText() ([]byte, error)     { return keyMarshalText("discokey:", k), nil }
 func (k *DiscoKey) UnmarshalText(text []byte) error { return keyUnmarshalText(k[:], "discokey:", text) }
 func (k DiscoKey) ShortString() string              { return fmt.Sprintf("d:%x", k[:8]) }
+func (k DiscoKey) AppendTo(b []byte) []byte         { return appendKey(b, "discokey:", k) }
 
 // IsZero reports whether k is the zero value.
 func (k DiscoKey) IsZero() bool { return k == DiscoKey{} }
diff --git a/tailcfg/tailcfg_test.go b/tailcfg/tailcfg_test.go
index 7632f3ccb..b1aad3857 100644
--- a/tailcfg/tailcfg_test.go
+++ b/tailcfg/tailcfg_test.go
@@ -14,6 +14,7 @@
 
 	"inet.af/netaddr"
 	"tailscale.com/types/wgkey"
+	"tailscale.com/version"
 )
 
 func fieldsOf(t reflect.Type) (fields []string) {
@@ -528,3 +529,25 @@ func BenchmarkKeyMarshalText(b *testing.B) {
 		sinkBytes = keyMarshalText("prefix", k)
 	}
 }
+
+func TestAppendKeyAllocs(t *testing.T) {
+	if version.IsRace() {
+		t.Skip("skipping in race detector") // append(b, make([]byte, N)...) not optimized in compiler with race
+	}
+	var k [32]byte
+	n := int(testing.AllocsPerRun(1000, func() {
+		sinkBytes = keyMarshalText("prefix", k)
+	}))
+	if n != 1 {
+		t.Fatalf("allocs = %v; want 1", n)
+	}
+}
+
+func TestDiscoKeyAppend(t *testing.T) {
+	d := DiscoKey{1: 1, 2: 2}
+	got := string(d.AppendTo([]byte("foo")))
+	want := "foodiscokey:0001020000000000000000000000000000000000000000000000000000000000"
+	if got != want {
+		t.Errorf("got %q; want %q", got, want)
+	}
+}
diff --git a/types/wgkey/key.go b/types/wgkey/key.go
index 3e2252010..15ddba39d 100644
--- a/types/wgkey/key.go
+++ b/types/wgkey/key.go
@@ -10,7 +10,6 @@
 package wgkey
 
 import (
-	"bytes"
 	"crypto/rand"
 	"crypto/subtle"
 	"encoding/base64"
@@ -72,10 +71,11 @@ func ParsePrivateHex(v string) (Private, error) {
 	return pk, nil
 }
 
-func (k Key) Base64() string    { return base64.StdEncoding.EncodeToString(k[:]) }
-func (k Key) String() string    { return k.ShortString() }
-func (k Key) HexString() string { return hex.EncodeToString(k[:]) }
-func (k Key) Equal(k2 Key) bool { return subtle.ConstantTimeCompare(k[:], k2[:]) == 1 }
+func (k Key) Base64() string           { return base64.StdEncoding.EncodeToString(k[:]) }
+func (k Key) String() string           { return k.ShortString() }
+func (k Key) HexString() string        { return hex.EncodeToString(k[:]) }
+func (k Key) Equal(k2 Key) bool        { return subtle.ConstantTimeCompare(k[:], k2[:]) == 1 }
+func (k Key) AppendTo(b []byte) []byte { return appendKey(b, "", k) }
 
 func (k *Key) ShortString() string {
 	// The goal here is to generate "[" + base64.StdEncoding.EncodeToString(k[:])[:5] + "]".
@@ -178,13 +178,17 @@ func (k *Private) Public() Key {
 	return (Key)(p)
 }
 
-func (k Private) MarshalText() ([]byte, error) {
-	// TODO(josharian): use encoding/hex instead?
-	buf := new(bytes.Buffer)
-	fmt.Fprintf(buf, `privkey:%x`, k[:])
-	return buf.Bytes(), nil
+func appendKey(base []byte, prefix string, k [32]byte) []byte {
+	ret := append(base, make([]byte, len(prefix)+64)...)
+	buf := ret[len(base):]
+	copy(buf, prefix)
+	hex.Encode(buf[len(prefix):], k[:])
+	return ret
 }
 
+func (k Private) MarshalText() ([]byte, error) { return appendKey(nil, "privkey:", k), nil }
+func (k Private) AppendTo(b []byte) []byte     { return appendKey(b, "privkey:", k) }
+
 func (k *Private) UnmarshalText(b []byte) error {
 	s := string(b)
 	if !strings.HasPrefix(s, `privkey:`) {