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 <[email protected]>

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{} }

tailcfg/tailcfg_test.go

@@ -14,6 +14,7 @@ import (
 
 	"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)
+	}
+}

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:`) {