Extract common funcs

cristalhq · Apr 15, 2024 · 8f2edab · 8f2edab
1 parent 16eb750
commit 8f2edab
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 145 deletions.
diff --git a/common.go b/common.go
@@ -2,12 +2,22 @@ package paseto
 
 import (
 	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/hmac"
+	"crypto/sha512"
 	"crypto/subtle"
 	"encoding/base64"
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"io"
 	"strings"
+
+	"golang.org/x/crypto/blake2b"
+	"golang.org/x/crypto/chacha20"
+	"golang.org/x/crypto/chacha20poly1305"
+	"golang.org/x/crypto/hkdf"
 )
 
 func pae(pieces ...[]byte) []byte {
@@ -110,6 +120,64 @@ func buildToken(header, body, footer []byte) string {
 	return string(token)
 }
 
+func doHMACSHA384(key, b []byte) []byte {
+	h := hmac.New(sha512.New384, key)
+	h.Write(b)
+	return h.Sum(nil)
+}
+
+func doBLAKE2b(size int, key, b []byte) []byte {
+	h, err := blake2b.New(size, key)
+	if err != nil {
+		panic(err)
+	}
+	h.Write(b)
+	return h.Sum(nil)
+}
+
+func doAES256CTR(key, nonce, m []byte) []byte {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		panic(err)
+	}
+
+	dst := make([]byte, len(m))
+	ciph := cipher.NewCTR(block, nonce)
+	ciph.XORKeyStream(dst, m)
+	return dst
+}
+
+func doCHACHA20(key, nonce, m []byte) []byte {
+	ciph, err := chacha20.NewUnauthenticatedCipher(key, nonce)
+	if err != nil {
+		panic(err)
+	}
+
+	dst := make([]byte, len(m))
+	ciph.XORKeyStream(dst, m)
+	return dst
+}
+
+func doOpenCHACHA(key, nonce, c, a []byte) ([]byte, error) {
+	aead, err := chacha20poly1305.NewX(key)
+	if err != nil {
+		panic(err)
+	}
+	return aead.Open(nil, nonce, c, a)
+}
+
+func doSealCHACHA(key, nonce, m, a []byte) []byte {
+	aead, err := chacha20poly1305.NewX(key)
+	if err != nil {
+		panic(err)
+	}
+	return aead.Seal(nil, nonce, m, a)
+}
+
+func doHKDF(key, salt, info []byte) io.Reader {
+	return hkdf.New(sha512.New384, key, salt, info)
+}
+
 func b64Decode(dst, src []byte) (n int, err error) {
 	return base64.RawURLEncoding.Decode(dst, src)
 }

diff --git a/v1loc.go b/v1loc.go
@@ -1,17 +1,12 @@
 package paseto
 
 import (
-	"crypto/aes"
-	"crypto/cipher"
 	"crypto/hmac"
 	"crypto/rand"
-	"crypto/sha512"
 	"errors"
 	"fmt"
 	"io"
 	"strings"
-
-	"golang.org/x/crypto/hkdf"
 )
 
 const (
@@ -55,9 +50,7 @@ func V1Encrypt(key []byte, payload, footer any, randBytes []byte) (string, error
 	}
 
 	// step 4.
-	hash := hmac.New(sha512.New384, b)
-	hash.Write(m)
-	n := hash.Sum(nil)[:v1locNonce]
+	n := doHMACSHA384(b, m)[:v1locNonce]
 
 	// step 5.
 	ek, ak, err := v1locSplitKey(k, n[:v1locNonceH])
@@ -66,21 +59,13 @@ func V1Encrypt(key []byte, payload, footer any, randBytes []byte) (string, error
 	}
 
 	// step 6.
-	block, err := aes.NewCipher(ek)
-	if err != nil {
-		return "", fmt.Errorf("create aes cipher: %w", err)
-	}
-	c := make([]byte, len(m))
-	ciph := cipher.NewCTR(block, n[v1locNonceH:])
-	ciph.XORKeyStream(c, m)
+	c := doAES256CTR(ek, n[v1locNonceH:], m)
 
 	// step 7.
 	preAuth := pae(h, n, c, f)
 
 	// step 8.
-	hasher := hmac.New(sha512.New384, ak)
-	hasher.Write(preAuth)
-	t := hasher.Sum(nil)
+	t := doHMACSHA384(ak, preAuth)
 
 	// step 9.
 	body := make([]byte, 0, len(n)+len(c)+len(t))
@@ -133,50 +118,41 @@ func V1Decrypt(token string, key []byte, payload, footer any) error {
 	preAuth := pae(h, n, c, f)
 
 	// step 7.
-	hasher := hmac.New(sha512.New384, ak)
-	hasher.Write(preAuth)
-	t2 := hasher.Sum(nil)
+	t2 := doHMACSHA384(ak, preAuth)
 
 	// step 8.
 	if !hmac.Equal(t2, t) {
 		return ErrInvalidTokenAuth
 	}
 
 	// step 9.
-	block, err := aes.NewCipher(ek)
-	if err != nil {
-		return fmt.Errorf("create aes cipher: %w", err)
-	}
-
-	decryptedPayload := make([]byte, len(c))
-	ciph := cipher.NewCTR(block, n[v1locNonceH:])
-	ciph.XORKeyStream(decryptedPayload, c)
+	p := doAES256CTR(ek, n[v1locNonceH:], c)
 
 	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)
 		}
 	}
 	return nil
 }
 
 func v1locSplitKey(key, salt []byte) ([]byte, []byte, error) {
-	eReader := hkdf.New(sha512.New384, key, salt, []byte("paseto-encryption-key"))
-	aReader := hkdf.New(sha512.New384, key, salt, []byte("paseto-auth-key-for-aead"))
+	er := doHKDF(key, salt, []byte("paseto-encryption-key"))
+	ar := doHKDF(key, salt, []byte("paseto-auth-key-for-aead"))
 
 	ek := make([]byte, 32)
 	ak := make([]byte, 32)
 
-	if _, err := io.ReadFull(eReader, ek); err != nil {
+	if _, err := io.ReadFull(er, ek); err != nil {
 		return nil, nil, err
 	}
-	if _, err := io.ReadFull(aReader, ak); err != nil {
+	if _, err := io.ReadFull(ar, ak); err != nil {
 		return nil, nil, err
 	}
 	return ek, ak, nil

diff --git a/v2loc.go b/v2loc.go
@@ -2,12 +2,10 @@ package paseto
 
 import (
 	"crypto/rand"
-	"crypto/subtle"
 	"errors"
 	"fmt"
 	"io"
 
-	"golang.org/x/crypto/blake2b"
 	"golang.org/x/crypto/chacha20poly1305"
 )
 
@@ -34,7 +32,7 @@ func V2Encrypt(key []byte, payload, footer any, randBytes []byte) (string, error
 	f := footerBytes
 
 	// step 1.
-	if subtle.ConstantTimeEq(int32(len(k)), v2locKey) != 1 {
+	if !constTimeEq(int32(len(k)), v2locKey) {
 		return "", errors.New("bad key")
 	}
 
@@ -51,25 +49,18 @@ func V2Encrypt(key []byte, payload, footer any, randBytes []byte) (string, error
 	}
 
 	// step 4.
-	hasher, err := blake2b.New(v2locNonce, b)
-	if err != nil {
-		return "", fmt.Errorf("create blake2b hash: %w", err)
-	}
-	hasher.Write(m)
-	n := hasher.Sum(nil)
+	n := doBLAKE2b(v2locNonce, b, m)
 
 	// 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)
-	}
-	c := aead.Seal(m[:0], n, m, preAuth)
+	c := doSealCHACHA(k, n, m, preAuth)
 
 	// step 7.
-	body := append(n, c...)
+	body := make([]byte, 0, len(n)+len(c))
+	body = append(body, n...)
+	body = append(body, c...)
 
 	return buildToken(h, body, f), nil
 }
@@ -80,7 +71,7 @@ func V2Decrypt(token string, key []byte, payload, footer any) error {
 	k := key
 
 	// step 1.
-	if subtle.ConstantTimeEq(int32(len(k)), v2locKey) != 1 {
+	if !constTimeEq(int32(len(k)), v2locKey) {
 		return errors.New("bad key")
 	}
 
@@ -104,12 +95,7 @@ func V2Decrypt(token string, key []byte, payload, footer any) error {
 	preAuth := pae(h, n, f)
 
 	// step 6.
-	aead, err := chacha20poly1305.NewX(k)
-	if err != nil {
-		return fmt.Errorf("create chacha20poly1305 cipher: %w", err)
-	}
-
-	p, err := aead.Open(c[:0], n, c, preAuth)
+	p, err := doOpenCHACHA(k, n, c, preAuth)
 	if err != nil {
 		return ErrInvalidTokenAuth
 	}

diff --git a/v3loc.go b/v3loc.go
@@ -1,18 +1,12 @@
 package paseto
 
 import (
-	"crypto/aes"
-	"crypto/cipher"
 	"crypto/hmac"
 	"crypto/rand"
-	"crypto/sha512"
-	"crypto/subtle"
 	"errors"
 	"fmt"
 	"io"
 	"strings"
-
-	"golang.org/x/crypto/hkdf"
 )
 
 const (
@@ -41,7 +35,7 @@ func V3Encrypt(key []byte, payload, footer any, implicit string, randBytes []byt
 	i := []byte(implicit)
 
 	// step 1.
-	if subtle.ConstantTimeEq(int32(len(k)), v3locKey) != 1 {
+	if !constTimeEq(int32(len(k)), v3locKey) {
 		return "", errors.New("bad key")
 	}
 
@@ -64,21 +58,13 @@ func V3Encrypt(key []byte, payload, footer any, implicit string, randBytes []byt
 	}
 
 	// step 5.
-	block, err := aes.NewCipher(ek)
-	if err != nil {
-		return "", fmt.Errorf("create aes cipher: %w", err)
-	}
-	c := make([]byte, len(m))
-	ciph := cipher.NewCTR(block, n2)
-	ciph.XORKeyStream(c, m)
+	c := doAES256CTR(ek, n2, m)
 
 	// step 6.
 	preAuth := pae(h, n, c, f, i)
 
 	// step 7.
-	hasher := hmac.New(sha512.New384, ak)
-	hasher.Write(preAuth)
-	t := hasher.Sum(nil)
+	t := doHMACSHA384(ak, preAuth)
 
 	// step 7.
 	body := make([]byte, 0, len(n)+len(c)+len(t))
@@ -96,7 +82,7 @@ func V3Decrypt(token string, key []byte, payload, footer any, implicit string) e
 	i := []byte(implicit)
 
 	// step 1.
-	if subtle.ConstantTimeEq(int32(len(k)), v3locKey) != 1 {
+	if !constTimeEq(int32(len(k)), v3locKey) {
 		return errors.New("bad key")
 	}
 
@@ -129,24 +115,15 @@ func V3Decrypt(token string, key []byte, payload, footer any, implicit string) e
 	preAuth := pae(h, n, c, f, i)
 
 	// step 7.
-	hasher := hmac.New(sha512.New384, ak)
-	hasher.Write(preAuth)
-	t2 := hasher.Sum(nil)
+	t2 := doHMACSHA384(ak, preAuth)
 
 	// step 8.
 	if !hmac.Equal(t, t2) {
 		return ErrInvalidTokenAuth
 	}
 
 	// step 9.
-	block, err := aes.NewCipher(ek)
-	if err != nil {
-		return fmt.Errorf("create aes cipher: %w", err)
-	}
-
-	p := make([]byte, len(c))
-	ciph := cipher.NewCTR(block, n2)
-	ciph.XORKeyStream(p, c)
+	p := doAES256CTR(ek, n2, c)
 
 	// step 7.
 	if payload != nil {
@@ -164,19 +141,19 @@ func V3Decrypt(token string, key []byte, payload, footer any, implicit string) e
 }
 
 func v3locSplitKey(key, n []byte) (ek, ak, n2 []byte, err error) {
-	eKDF := hkdf.New(sha512.New384, key, nil, append([]byte("paseto-encryption-key"), n...))
-	aKDF := hkdf.New(sha512.New384, key, nil, append([]byte("paseto-auth-key-for-aead"), n...))
+	er := doHKDF(key, nil, append([]byte("paseto-encryption-key"), n...))
+	ar := doHKDF(key, nil, append([]byte("paseto-auth-key-for-aead"), n...))
 
 	tmp := make([]byte, v3locKDF)
-	if _, err := io.ReadFull(eKDF, tmp); err != nil {
-		return nil, nil, nil, fmt.Errorf("unable to generate encryption key from seed: %w", err)
+	if _, err := io.ReadFull(er, tmp); err != nil {
+		return nil, nil, nil, err
 	}
 
 	ek, n2 = tmp[:v3locKey], tmp[v3locKey:]
 
 	ak = make([]byte, v3locKDF)
-	if _, err := io.ReadFull(aKDF, ak); err != nil {
-		return nil, nil, nil, fmt.Errorf("unable to generate authentication key from seed: %w", err)
+	if _, err := io.ReadFull(ar, ak); err != nil {
+		return nil, nil, nil, err
 	}
 	return ek, n2, ak, nil
 }