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>
This commit is contained in:
Brad Fitzpatrick 2021-05-24 14:54:47 -07:00 committed by Brad Fitzpatrick
parent b340beff8e
commit e66d4e4c81
3 changed files with 46 additions and 13 deletions

View File

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

View File

@ -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)
}
}

View File

@ -10,7 +10,6 @@
package wgkey
import (
"bytes"
"crypto/rand"
"crypto/subtle"
"encoding/base64"
@ -76,6 +75,7 @@ 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:`) {