From cfdb8626738daf799fd689b62de3891500789238 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 19 Jul 2022 14:04:08 -0700
Subject: [PATCH] types/key: add naclbox shared key wrapper type + Seal method

So the control plane can stop doing precomputations on each naclbox
message.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 types/key/machine.go      | 27 +++++++++++++++++++++++++++
 types/key/machine_test.go | 20 ++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/types/key/machine.go b/types/key/machine.go
index 9687300f5..74eb24b5f 100644
--- a/types/key/machine.go
+++ b/types/key/machine.go
@@ -105,6 +105,33 @@ func (k MachinePrivate) SealTo(p MachinePublic, cleartext []byte) (ciphertext []
 	return box.Seal(nonce[:], cleartext, &nonce, &p.k, &k.k)
 }
 
+// SharedKey returns the precomputed Nacl box shared key between k and p.
+func (k MachinePrivate) SharedKey(p MachinePublic) MachinePrecomputedSharedKey {
+	var shared MachinePrecomputedSharedKey
+	box.Precompute(&shared.k, &p.k, &k.k)
+	return shared
+}
+
+// MachinePrecomputedSharedKey is a precomputed shared NaCl box shared key.
+type MachinePrecomputedSharedKey struct {
+	k [32]byte
+}
+
+// Seal wraps cleartext into a NaCl box (see
+// golang.org/x/crypto/nacl) using the shared key k as generated
+// by MachinePrivate.SharedKey.
+//
+// The returned ciphertext is a 24-byte nonce concatenated with the
+// box value.
+func (k MachinePrecomputedSharedKey) Seal(cleartext []byte) (ciphertext []byte) {
+	if k == (MachinePrecomputedSharedKey{}) {
+		panic("can't seal with zero keys")
+	}
+	var nonce [24]byte
+	rand(nonce[:])
+	return box.SealAfterPrecomputation(nonce[:], cleartext, &nonce, &k.k)
+}
+
 // OpenFrom opens the NaCl box ciphertext, which must be a value
 // created by SealTo, and returns the inner cleartext if ciphertext is
 // a valid box from p to k.
diff --git a/types/key/machine_test.go b/types/key/machine_test.go
index 88b73fb49..1ded5da86 100644
--- a/types/key/machine_test.go
+++ b/types/key/machine_test.go
@@ -90,3 +90,23 @@ func TestMachineSerialization(t *testing.T) {
 		t.Error("json serialization doesn't roundtrip")
 	}
 }
+
+func TestSealViaSharedKey(t *testing.T) {
+	// encrypt a message from a to b
+	a := NewMachine()
+	b := NewMachine()
+	apub, bpub := a.Public(), b.Public()
+
+	shared := a.SharedKey(bpub)
+
+	const clear = "the eagle flies at midnight"
+	enc := shared.Seal([]byte(clear))
+
+	back, ok := b.OpenFrom(apub, enc)
+	if !ok {
+		t.Fatal("failed to decrypt")
+	}
+	if string(back) != clear {
+		t.Errorf("got %q; want cleartext %q", back, clear)
+	}
+}