extract function newCipherBlock

Author: briskt

apikey.go

@@ -8,13 +8,12 @@ import (
 	"encoding/base64"
 	"encoding/hex"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"time"
 
-	"github.com/pkg/errors"
-
 	"golang.org/x/crypto/bcrypt"
 )
 
@@ -72,16 +71,9 @@ func (k *ApiKey) IsCorrect(given string) error {
 
 // EncryptData uses the Secret to AES encrypt an arbitrary data block. It does not encrypt the key itself.
 func (k *ApiKey) EncryptData(plaintext []byte) ([]byte, error) {
-	var sec []byte
-	var err error
-	sec, err = base64.StdEncoding.DecodeString(k.Secret)
-	if err != nil {
-		sec = []byte(k.Secret)
-	}
-	// create cipher block with api secret as aes key
-	block, err := aes.NewCipher(sec)
+	block, err := newCipherBlock(k.Secret)
 	if err != nil {
-		return []byte{}, err
+		return nil, err
 	}
 
 	// byte array to hold encrypted content
@@ -103,16 +95,9 @@ func (k *ApiKey) EncryptData(plaintext []byte) ([]byte, error) {
 
 // DecryptData uses the Secret to AES decrypt an arbitrary data block. It does not decrypt the key itself.
 func (k *ApiKey) DecryptData(ciphertext []byte) ([]byte, error) {
-	var sec []byte
-	var err error
-	sec, err = base64.StdEncoding.DecodeString(k.Secret)
-	if err != nil {
-		sec = []byte(k.Secret)
-	}
-
-	block, err := aes.NewCipher(sec)
+	block, err := newCipherBlock(k.Secret)
 	if err != nil {
-		return []byte{}, errors.Wrap(err, "failed to create new cipher")
+		return nil, err
 	}
 
 	// plaintext must be as long as ciphertext minus the length of the IV, which is the same as the AES block size
@@ -131,16 +116,9 @@ func (k *ApiKey) DecryptData(ciphertext []byte) ([]byte, error) {
 // DecryptLegacy uses the Secret to AES decrypt an arbitrary data block. This is intended only for legacy data such
 // as U2F keys.
 func (k *ApiKey) DecryptLegacy(ciphertext []byte) ([]byte, error) {
-	var sec []byte
-	var err error
-	sec, err = base64.StdEncoding.DecodeString(k.Secret)
-	if err != nil {
-		sec = []byte(k.Secret)
-	}
-
-	block, err := aes.NewCipher(sec)
+	block, err := newCipherBlock(k.Secret)
 	if err != nil {
-		return []byte{}, errors.Wrap(err, "failed to create new cipher")
+		return nil, err
 	}
 
 	// data was encrypted, then base64 encoded, then joined with a :, need to split
@@ -302,3 +280,20 @@ func NewApiKey(email string) (ApiKey, error) {
 	}
 	return key, nil
 }
+
+// newCipherBlock creates a new cipher.Block from a base64-encoded AES key. If the string is not valid base64 data, it
+// will be interpreted as binary data.
+func newCipherBlock(key string) (cipher.Block, error) {
+	var sec []byte
+	var err error
+	sec, err = base64.StdEncoding.DecodeString(key)
+	if err != nil {
+		sec = []byte(key)
+	}
+
+	block, err := aes.NewCipher(sec)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create new cipher: %w", err)
+	}
+	return block, nil
+}

apikey_test.go

@@ -2,8 +2,12 @@ package mfa
 
 import (
 	"bytes"
+	"crypto/aes"
+	"crypto/rand"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"regexp"
 	"testing"
@@ -309,3 +313,46 @@ func (ms *MfaSuite) TestNewApiKey() {
 	ms.NoError(err)
 	ms.Regexp(regexp.MustCompile("[a-f0-9]{40}"), got)
 }
+
+func (ms *MfaSuite) TestNewCipherBlock() {
+	random := make([]byte, 32)
+	_, err := io.ReadFull(rand.Reader, random)
+	ms.NoError(err)
+
+	tests := []struct {
+		name    string
+		key     string
+		wantErr bool
+	}{
+		{
+			name:    "key too short",
+			key:     "0123456789012345678901234567890",
+			wantErr: true,
+		},
+		{
+			name:    "key too long",
+			key:     "012345678901234567890123456789012",
+			wantErr: true,
+		},
+		{
+			name: "raw",
+			key:  string(random),
+		},
+		{
+			name: "base64",
+			key:  base64.StdEncoding.EncodeToString(random),
+		},
+	}
+	for _, tt := range tests {
+		ms.Run(tt.name, func() {
+			got, err := newCipherBlock(tt.key)
+			if tt.wantErr {
+				ms.Error(err)
+				return
+			}
+
+			ms.NoError(err)
+			ms.Equal(aes.BlockSize, got.BlockSize())
+		})
+	}
+}