encrypt/encrypt_stream_test.go at master · things-go/encrypt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
package encrypt

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"io"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var encdec = []struct {
	enc func(block cipher.Block, iv []byte) cipher.Stream
	dec func(block cipher.Block, iv []byte) cipher.Stream
}{
	{cipher.NewCFBEncrypter, cipher.NewCFBDecrypter},
	{cipher.NewCTR, cipher.NewCTR},
	{cipher.NewOFB, cipher.NewOFB},
}

func TestBlockStream(t *testing.T) {
	t.Run("encrypt decrypt", func(t *testing.T) {
		plainText := []byte("hello world,this is golang language. welcome")
		for _, v := range encdec {
			for _, keySize := range aesKeySizes {
				key := make([]byte, keySize)
				_, err := io.ReadFull(rand.Reader, key)
				require.NoError(t, err)

				blk, err := NewStreamCipher(key, aes.NewCipher,
					WithStreamRandIV(RandIV),
					WithStreamCodec(v.enc, v.dec),
				)
				require.NoError(t, err)

				assert.Equal(t, aes.BlockSize, blk.BlockSize())

				cipherText, err := blk.Encrypt(plainText)
				require.NoError(t, err)
				want, err := blk.Decrypt(cipherText)
				require.NoError(t, err)
				assert.Equal(t, want, plainText)

				cipherText, err = blk.Encrypt(plainText)
				require.NoError(t, err)
				want, err = blk.Decrypt(cipherText)
				require.NoError(t, err)
				assert.Equal(t, want, plainText)
			}
		}
	})

	t.Run("invalid cipher text", func(t *testing.T) {
		key := make([]byte, 16)
		_, err := io.ReadFull(rand.Reader, key)
		require.NoError(t, err)

		blk, err := aes.NewCipher(key)
		require.NoError(t, err)

		stream := blockStream{
			block:      blk,
			newEncrypt: cipher.NewCFBEncrypter,
			newDecrypt: cipher.NewCFBDecrypter,
		}

		_, err = stream.Encrypt([]byte{})
		require.EqualError(t, err, ErrInputInvalidLength.Error())

		_, err = stream.Decrypt(key[:len(key)-1])
		require.EqualError(t, err, ErrInputNotMoreABlock.Error())
	})
	t.Run("invalid cipher", func(t *testing.T) {
		key := make([]byte, 16)
		_, err := io.ReadFull(rand.Reader, key)
		require.NoError(t, err)

		_, err = NewStreamCipher(key, mockErrorNewCipher)
		require.Error(t, err)
	})
	t.Run("invalid iv function or length", func(t *testing.T) {
		plainText := []byte("hello world,this is golang language. welcome")
		key := make([]byte, 16)
		_, err := io.ReadFull(rand.Reader, key)
		require.NoError(t, err)

		bc, err := NewStreamCipher(key, aes.NewCipher, WithStreamRandIV(func(block cipher.Block) ([]byte, error) {
			return nil, errors.New("invalid iv")
		}))
		require.NoError(t, err)

		_, err = bc.Encrypt(plainText)
		require.Error(t, err)
	})
}