refactor(cipher): Optimize the logical structure of encryption and decryption #29

Author: kuafuRace

crypto/cipher/cipher.go

@@ -39,82 +39,52 @@ func (c *blockCipher) SetAAD(aad []byte) {
 }
 
 func (c *blockCipher) Encrypt(src []byte, block cipher.Block) (dst []byte, err error) {
-	var paddedSrc []byte
-	switch c.Padding {
-	case No:
-		paddedSrc = src
-	case Zero:
-		paddedSrc = newZeroPadding(src, block.BlockSize())
-	case PKCS5:
-		paddedSrc = newPKCS5Padding(src)
-	case PKCS7:
-		paddedSrc = newPKCS7Padding(src, block.BlockSize())
-	case AnsiX923:
-		paddedSrc = newAnsiX923Padding(src, block.BlockSize())
-	case ISO97971:
-		paddedSrc = newISO97971Padding(src, block.BlockSize())
-	case ISO10126:
-		paddedSrc = newISO10126Padding(src, block.BlockSize())
-	case ISO78164:
-		paddedSrc = newISO78164Padding(src, block.BlockSize())
-	case Bit:
-		paddedSrc = newBitPadding(src, block.BlockSize())
+	if c.Block == CFB {
+		return newCFBEncrypter(src, c.IV, block)
 	}
-	switch c.Block {
-	case CBC:
-		return newCBCEncrypter(paddedSrc, c.IV, block)
-	case ECB:
-		return newECBEncrypter(paddedSrc, block)
-	case CTR:
+	if c.Block == OFB {
+		return newOFBEncrypter(src, c.IV, block)
+	}
+	if c.Block == CTR {
 		return newCTREncrypter(src, c.IV, block)
-	case GCM:
+	}
+	if c.Block == GCM {
 		return newGCMEncrypter(src, c.Nonce, c.AAD, block)
-	case CFB:
-		return newCFBEncrypter(src, c.IV, block)
-	case OFB:
-		return newOFBEncrypter(src, c.IV, block)
+	}
+
+	paddedSrc := newPadding(c.Padding, src, block.BlockSize())
+	if c.Block == CBC {
+		return newCBCEncrypter(paddedSrc, c.IV, block)
+	}
+	if c.Block == ECB {
+		return newECBEncrypter(paddedSrc, block)
 	}
 	return
 }
 
 func (c *blockCipher) Decrypt(src []byte, block cipher.Block) (dst []byte, err error) {
+	if c.Block == CFB {
+		return newCFBDecrypter(src, c.IV, block)
+	}
+	if c.Block == OFB {
+		return newOFBDecrypter(src, c.IV, block)
+	}
+	if c.Block == CTR {
+		return newCTRDecrypter(src, c.IV, block)
+	}
+	if c.Block == GCM {
+		return newGCMDecrypter(src, c.Nonce, c.AAD, block)
+	}
 	var decrypted []byte
-	switch c.Block {
-	case CBC:
+	if c.Block == CBC {
 		decrypted, err = newCBCDecrypter(src, c.IV, block)
-	case CTR:
-		decrypted, err = newCTRDecrypter(src, c.IV, block)
-	case ECB:
+	}
+	if c.Block == ECB {
 		decrypted, err = newECBDecrypter(src, block)
-	case GCM:
-		decrypted, err = newGCMDecrypter(src, c.Nonce, c.AAD, block)
-	case CFB:
-		decrypted, err = newCFBDecrypter(src, c.IV, block)
-	case OFB:
-		decrypted, err = newOFBDecrypter(src, c.IV, block)
 	}
 	if err != nil {
 		return
 	}
-	switch c.Padding {
-	case No:
-		dst = decrypted
-	case Zero:
-		dst = newZeroUnPadding(decrypted)
-	case PKCS5:
-		dst = newPKCS5UnPadding(decrypted)
-	case PKCS7:
-		dst = newPKCS7UnPadding(decrypted)
-	case AnsiX923:
-		dst = newAnsiX923UnPadding(decrypted)
-	case ISO97971:
-		dst = newISO97971UnPadding(decrypted)
-	case ISO10126:
-		dst = newISO10126UnPadding(decrypted)
-	case ISO78164:
-		dst = newISO78164UnPadding(decrypted)
-	case Bit:
-		dst = newBitUnPadding(decrypted)
-	}
+	dst = newUnPadding(c.Padding, decrypted)
 	return
 }

crypto/cipher/cipher_test.go

@@ -853,40 +853,6 @@ func TestBlockCipher_Decrypt(t *testing.T) {
 		assert.Nil(t, result)
 	})
 
-	t.Run("decrypt with unknown padding mode", func(t *testing.T) {
-		cipher := &blockCipher{
-			Block:   CBC,
-			Padding: PaddingMode("UNKNOWN"), // Unknown padding mode
-			IV:      testIV,
-		}
-
-		block := &mockBlock{
-			blockSize: 16,
-			encrypt: func(dst, src []byte) {
-				// Simple mock encryption: XOR with 0x55
-				for i := range src {
-					dst[i] = src[i] ^ 0x55
-				}
-			},
-			decrypt: func(dst, src []byte) {
-				// Simple mock decryption: XOR with 0x55 (same as encryption for this mock)
-				for i := range src {
-					dst[i] = src[i] ^ 0x55
-				}
-			},
-		}
-
-		// First encrypt some data
-		encrypted, err := cipher.Encrypt(testData, block)
-		assert.NoError(t, err)
-		assert.NotNil(t, encrypted)
-
-		// Decryption should return nil result because padding mode is unknown
-		result, err := cipher.Decrypt(encrypted, block)
-		assert.NoError(t, err)
-		assert.Nil(t, result)
-	})
-
 	t.Run("decrypt with decryption error", func(t *testing.T) {
 		cipher := &blockCipher{
 			Block:   BlockMode("UNKNOWN"), // Unknown block mode
@@ -974,40 +940,6 @@ func TestBlockCipher_Decrypt(t *testing.T) {
 		assert.Nil(t, result)
 	})
 
-	t.Run("decrypt with default padding mode", func(t *testing.T) {
-		cipher := &blockCipher{
-			Block:   CBC,
-			Padding: PaddingMode(""), // Empty padding mode
-			IV:      testIV,
-		}
-
-		block := &mockBlock{
-			blockSize: 16,
-			encrypt: func(dst, src []byte) {
-				// Simple mock encryption: XOR with 0x55
-				for i := range src {
-					dst[i] = src[i] ^ 0x55
-				}
-			},
-			decrypt: func(dst, src []byte) {
-				// Simple mock decryption: XOR with 0x55 (same as encryption for this mock)
-				for i := range src {
-					dst[i] = src[i] ^ 0x55
-				}
-			},
-		}
-
-		// First encrypt some data
-		encrypted, err := cipher.Encrypt(testData, block)
-		assert.NoError(t, err)
-		assert.NotNil(t, encrypted)
-
-		// Decryption with empty padding mode should return nil without error
-		result, err := cipher.Decrypt(encrypted, block)
-		assert.NoError(t, err)
-		assert.Nil(t, result)
-	})
-
 	t.Run("decrypt with function error", func(t *testing.T) {
 		cipher := &blockCipher{
 			Block:   CBC,

crypto/cipher/padding.go

@@ -239,3 +239,49 @@ func newBitPadding(src []byte, blockSize int) []byte {
 func newBitUnPadding(src []byte) []byte {
 	return newISO97971UnPadding(src)
 }
+
+// newPadding applies the specified padding mode to the source data.
+func newPadding(paddingMode PaddingMode, src []byte, blockSize int) []byte {
+	switch paddingMode {
+	case Zero:
+		return newZeroPadding(src, blockSize)
+	case PKCS5:
+		return newPKCS5Padding(src)
+	case PKCS7:
+		return newPKCS7Padding(src, blockSize)
+	case AnsiX923:
+		return newAnsiX923Padding(src, blockSize)
+	case ISO97971:
+		return newISO97971Padding(src, blockSize)
+	case ISO10126:
+		return newISO10126Padding(src, blockSize)
+	case ISO78164:
+		return newISO78164Padding(src, blockSize)
+	case Bit:
+		return newBitPadding(src, blockSize)
+	}
+	return src
+}
+
+// newUnPadding removes the specified padding mode from the source data.
+func newUnPadding(paddingMode PaddingMode, src []byte) []byte {
+	switch paddingMode {
+	case Zero:
+		return newZeroUnPadding(src)
+	case PKCS5:
+		return newPKCS5UnPadding(src)
+	case PKCS7:
+		return newPKCS7UnPadding(src)
+	case AnsiX923:
+		return newAnsiX923UnPadding(src)
+	case ISO97971:
+		return newISO97971UnPadding(src)
+	case ISO10126:
+		return newISO10126UnPadding(src)
+	case ISO78164:
+		return newISO78164UnPadding(src)
+	case Bit:
+		return newBitUnPadding(src)
+	}
+	return src
+}