From 60895dec587ce3022b5eb9feaf2e966991ccdb72 Mon Sep 17 00:00:00 2001
From: Oleg Kovalov <oleg@hey.com>
Date: Sun, 14 Apr 2024 16:38:40 +0200
Subject: [PATCH] Rework v2.local (#10)

---
 v2loc.go      | 107 ++++++++++++++++++++++++++++++++------------------
 v2loc_test.go |   4 +-
 2 files changed, 70 insertions(+), 41 deletions(-)

diff --git a/v2loc.go b/v2loc.go
index ae7db90..1c459de 100644
--- a/v2loc.go
+++ b/v2loc.go
@@ -2,6 +2,8 @@ package paseto
 
 import (
 	"crypto/rand"
+	"crypto/subtle"
+	"errors"
 	"fmt"
 	"io"
 
@@ -10,18 +12,12 @@ import (
 )
 
 const (
-	v2LocHeader = "v2.local."
-	v2NonceSize = chacha20poly1305.NonceSizeX
+	v2locHeader = "v2.local."
+	v2locKey    = 32
+	v2locNonce  = chacha20poly1305.NonceSizeX
 )
 
 func V2Encrypt(key []byte, payload, footer any, randBytes []byte) (string, error) {
-	if randBytes == nil {
-		randBytes = make([]byte, v2NonceSize)
-		if _, err := io.ReadFull(rand.Reader, randBytes); err != nil {
-			return "", fmt.Errorf("read from crypto/rand.Reader: %w", err)
-		}
-	}
-
 	payloadBytes, err := toBytes(payload)
 	if err != nil {
 		return "", fmt.Errorf("encode payload: %w", err)
@@ -32,68 +28,101 @@ func V2Encrypt(key []byte, payload, footer any, randBytes []byte) (string, error
 		return "", fmt.Errorf("encode footer: %w", err)
 	}
 
-	hash, err := blake2b.New(v2NonceSize, randBytes)
+	// step 0.
+	m := payloadBytes
+	k := key
+	f := footerBytes
+
+	// step 1.
+	if subtle.ConstantTimeEq(int32(len(k)), v2locKey) != 1 {
+		return "", errors.New("bad key")
+	}
+
+	// step 2.
+	h := []byte(v2locHeader)
+
+	// step 3.
+	b := randBytes
+	if b == nil {
+		b = make([]byte, v2locNonce)
+		if _, err := io.ReadFull(rand.Reader, b); err != nil {
+			return "", fmt.Errorf("read from crypto/rand.Reader: %w", err)
+		}
+	}
+
+	// step 4.
+	hasher, err := blake2b.New(v2locNonce, b)
 	if err != nil {
 		return "", fmt.Errorf("create blake2b hash: %w", err)
 	}
-	if _, err := hash.Write(payloadBytes); err != nil {
-		return "", fmt.Errorf("hash payload: %w", err)
-	}
-	nonce := hash.Sum(nil)
+	hasher.Write(m)
+	n := hasher.Sum(nil)
+
+	// step 5.
+	preAuth := pae(h, n, f)
 
-	aead, err := chacha20poly1305.NewX(key)
+	// step 6.
+	aead, err := chacha20poly1305.NewX(k)
 	if err != nil {
 		return "", fmt.Errorf("create chacha20poly1305 cipher: %w", err)
 	}
+	c := aead.Seal(m[:0], n, m, preAuth)
 
-	preAuth := pae([]byte(v2LocHeader), nonce, footerBytes)
+	// step 7.
+	body := append(n, c...)
 
-	encryptedPayload := aead.Seal(
-		payloadBytes[:0],
-		nonce,
-		payloadBytes,
-		preAuth,
-	)
-	body := append(nonce, encryptedPayload...)
-
-	return buildToken([]byte(v2LocHeader), body, footerBytes), nil
+	return buildToken(h, body, f), nil
 }
 
 func V2Decrypt(token string, key []byte, payload, footer any) error {
-	body, footerBytes, err := splitToken(token, v2LocHeader)
+	// step 0.
+	m := token
+	k := key
+
+	// step 1.
+	if subtle.ConstantTimeEq(int32(len(k)), v2locKey) != 1 {
+		return errors.New("bad key")
+	}
+
+	// step 2.
+	// TODO: ?
+
+	// step 3.
+	h := []byte(v2locHeader)
+
+	// step 4.
+	body, footerBytes, err := splitToken(m, v2locHeader)
 	if err != nil {
 		return fmt.Errorf("decode token: %w", err)
 	}
-	if len(body) < v2NonceSize {
+	if len(body) < v2locNonce {
 		return ErrIncorrectTokenFormat
 	}
+	n, c, f := body[:v2locNonce], body[v2locNonce:], footerBytes
 
-	aead, err := chacha20poly1305.NewX(key)
+	// step 5.
+	preAuth := pae(h, n, f)
+
+	// step 6.
+	aead, err := chacha20poly1305.NewX(k)
 	if err != nil {
 		return fmt.Errorf("create chacha20poly1305 cipher: %w", err)
 	}
 
-	nonce, encryptedPayload := body[:v2NonceSize], body[v2NonceSize:]
-	preAuth := pae([]byte(v2LocHeader), nonce, footerBytes)
-
-	decryptedPayload, err := aead.Open(
-		encryptedPayload[:0],
-		nonce,
-		encryptedPayload,
-		preAuth,
-	)
+	p, err := aead.Open(c[:0], n, c, preAuth)
 	if err != nil {
 		return ErrInvalidTokenAuth
 	}
 
+	// step 7.
 	if payload != nil {
-		if err := fromBytes(decryptedPayload, payload); err != nil {
+		if err := fromBytes(p, payload); err != nil {
 			return fmt.Errorf("decode payload: %w", err)
 		}
 	}
 
 	if footer != nil {
-		if err := fromBytes(footerBytes, footer); err != nil {
+		if err := fromBytes(f, footer); err != nil {
 			return fmt.Errorf("decode footer: %w", err)
 		}
 	}
diff --git a/v2loc_test.go b/v2loc_test.go
index 2385e8d..605a1f5 100644
--- a/v2loc_test.go
+++ b/v2loc_test.go
@@ -10,7 +10,7 @@ func TestV2Loc_Encrypt(t *testing.T) {
 	testCases := loadGoldenFile("testdata/v2.json")
 
 	for _, tc := range testCases.Tests {
-		if tc.Key == "" || !strings.HasPrefix(tc.Token, v2LocHeader) {
+		if tc.Key == "" || !strings.HasPrefix(tc.Token, v2locHeader) {
 			continue
 		}
 
@@ -33,7 +33,7 @@ func TestV2Loc_Decrypt(t *testing.T) {
 	testCases := loadGoldenFile("testdata/v2.json")
 
 	for _, tc := range testCases.Tests {
-		if tc.Key == "" || !strings.HasPrefix(tc.Token, v2LocHeader) {
+		if tc.Key == "" || !strings.HasPrefix(tc.Token, v2locHeader) {
 			continue
 		}