From eb5ff89fc24b99ab661464990fcaa5f835a7f651 Mon Sep 17 00:00:00 2001
From: armfazh <armfazh@cloudflare.com>
Date: Tue, 21 Feb 2023 02:11:19 -0800
Subject: [PATCH] Allow concatenating and reuse payloads.

---
 cipher/ascon/ascon.go      | 43 +++++++++++++++++++++++------
 cipher/ascon/ascon_test.go | 56 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 8 deletions(-)

diff --git a/cipher/ascon/ascon.go b/cipher/ascon/ascon.go
index 2164eeea..62648f0c 100644
--- a/cipher/ascon/ascon.go
+++ b/cipher/ascon/ascon.go
@@ -28,7 +28,16 @@ const (
 type Mode int
 
 // KeySize is 16 for Ascon128 and Ascon128a, or 20 for Ascon80pq.
-func (m Mode) KeySize() int { v := int(m) >> 2; return KeySize&^v | KeySize80pq&v }
+func (m Mode) KeySize() int {
+	switch m {
+	case Ascon128, Ascon128a, Ascon80pq:
+		v := int(m) >> 2
+		return KeySize&^v | KeySize80pq&v
+	default:
+		panic(ErrMode)
+	}
+}
+
 func (m Mode) String() string {
 	switch m {
 	case Ascon128:
@@ -105,15 +114,15 @@ func (a *Cipher) Seal(dst, nonce, plaintext, additionalData []byte) []byte {
 	}
 
 	ptLen := len(plaintext)
-	output := make([]byte, ptLen+TagSize)
-	ciphertext, tag := output[:ptLen], output[ptLen:]
+	ret, out := sliceForAppend(dst, ptLen+TagSize)
+	ciphertext, tag := out[:ptLen], out[ptLen:]
 
 	a.initialize(nonce)
 	a.assocData(additionalData)
 	a.procText(plaintext, ciphertext, true)
 	a.finalize(tag)
 
-	return output
+	return ret
 }
 
 // Open decrypts and authenticates ciphertext, authenticates the
@@ -136,7 +145,8 @@ func (a *Cipher) Open(dst, nonce, ciphertext, additionalData []byte) ([]byte, er
 	}
 
 	ptLen := len(ciphertext) - TagSize
-	plaintext := make([]byte, ptLen)
+	ret, out := sliceForAppend(dst, ptLen)
+	plaintext := out[:ptLen]
 	ciphertext, tag0 := ciphertext[:ptLen], ciphertext[ptLen:]
 	tag1 := (&[TagSize]byte{})[:]
 
@@ -149,7 +159,7 @@ func (a *Cipher) Open(dst, nonce, ciphertext, additionalData []byte) ([]byte, er
 		return nil, ErrDecryption
 	}
 
-	return plaintext, nil
+	return ret, nil
 }
 
 func abs(x int) int { m := uint(x >> (bits.UintSize - 1)); return int((uint(x) + m) ^ m) }
@@ -220,8 +230,10 @@ func (a *Cipher) procText(in, out []byte, enc bool) {
 	for i := 0; i < len(in); i++ {
 		off := 56 - (8 * (i % 8))
 		si := byte((a.s[i/8] >> off) & 0xFF)
-		out[i] = si ^ in[i]
-		ss := (in[i] &^ mask8) | (out[i] & mask8)
+		inB := in[i]
+		outB := si ^ inB
+		out[i] = outB
+		ss := inB&^mask8 | outB&mask8
 		a.s[i/8] = (a.s[i/8] &^ (0xFF << off)) | uint64(ss)<<off
 	}
 	a.s[len(in)/8] ^= uint64(0x80) << (56 - 8*(len(in)%8))
@@ -287,6 +299,21 @@ func (a *Cipher) perm(n int) {
 	a.s[0], a.s[1], a.s[2], a.s[3], a.s[4] = x0, x1, x2, x3, x4
 }
 
+// sliceForAppend takes a slice and a requested number of bytes. It returns a
+// slice with the contents of the given slice followed by that many bytes and a
+// second slice that aliases into it and contains only the extra bytes. If the
+// original slice has sufficient capacity then no allocation is performed.
+func sliceForAppend(in []byte, n int) (head, tail []byte) {
+	if total := len(in) + n; cap(in) >= total {
+		head = in[:total]
+	} else {
+		head = make([]byte, total)
+		copy(head, in)
+	}
+	tail = head[len(in):]
+	return
+}
+
 var (
 	ErrKeySize    = errors.New("ascon: bad key size")
 	ErrNonceSize  = errors.New("ascon: bad nonce size")
diff --git a/cipher/ascon/ascon_test.go b/cipher/ascon/ascon_test.go
index 04e6c046..80a8e049 100644
--- a/cipher/ascon/ascon_test.go
+++ b/cipher/ascon/ascon_test.go
@@ -127,6 +127,62 @@ func TestBadInputs(t *testing.T) {
 	test.CheckIsErr(t, err, "should panic due to bad ciphertext")
 }
 
+func TestAPI(t *testing.T) {
+	key, _ := hexa.DecodeString("000102030405060708090A0B0C0D0E0F")
+	nonce, _ := hexa.DecodeString("000102030405060708090A0B0C0D0E0F")
+	c, _ := ascon.New(key, ascon.Ascon128)
+	pt := []byte("helloworld")
+	ct, _ := hexa.DecodeString("d4e663d29cd60a693c20f890982e167d266f940b93b586945065")
+
+	t.Run("append", func(t *testing.T) {
+		prefix := [5]byte{0x1F, 0x2F, 0x3F, 0x4F, 0x5F}
+		prefixAndCt := c.Seal(prefix[:], nonce, pt, nil)
+		got := prefixAndCt
+		want := append(append([]byte{}, prefix[:]...), ct...)
+		if !bytes.Equal(got, want) {
+			test.ReportError(t, got, want)
+		}
+
+		ciphertext := prefixAndCt[len(prefix):]
+		prefix = [5]byte{0x11, 0x22, 0x33, 0x44, 0x55}
+		prefixAndPt, err := c.Open(prefix[:], nonce, ciphertext, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		got = prefixAndPt
+		want = append(append([]byte{}, prefix[:]...), pt...)
+		if !bytes.Equal(got, want) {
+			test.ReportError(t, got, want)
+		}
+	})
+	t.Run("reuse", func(t *testing.T) {
+		ptWithCap := make([]byte, len(pt), len(pt)+100)
+		copy(ptWithCap, pt)
+		// reusing the input to store the ciphertext.
+		ciphertext := c.Seal(ptWithCap[:0], nonce, ptWithCap, nil)
+		got := ciphertext
+		want := ct
+		if !bytes.Equal(got, want) {
+			test.ReportError(t, got, want)
+		}
+		test.CheckOk(&ptWithCap[0] == &ciphertext[0], "should have same address", t)
+
+		ctWithCap := make([]byte, len(ct), len(ct)+100)
+		copy(ctWithCap, ct)
+		// reusing the input to store the plaintext.
+		plaintext, err := c.Open(ctWithCap[:0], nonce, ctWithCap, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		got = plaintext
+		want = pt
+		if !bytes.Equal(got, want) {
+			test.ReportError(t, got, want)
+		}
+		test.CheckOk(&ctWithCap[0] == &plaintext[0], "should have same address", t)
+	})
+}
+
 func BenchmarkAscon(b *testing.B) {
 	for _, mode := range []ascon.Mode{ascon.Ascon128, ascon.Ascon128a, ascon.Ascon80pq} {
 		for _, length := range []int{64, 1350, 8 * 1024} {