fix: コード最適化、テスト追加

contentencoding.go

@@ -44,13 +44,11 @@ func (i ContentEncoding) overhead(gcm cipher.AEAD) int {
 }
 
 func (i ContentEncoding) calculateRecordPadSize(pad, baseRecordSize int) int {
-	padSize := i.Padding()
-
 	// Pad so that at least one data byte is in a block.
-	recordPad := min(baseRecordSize-1, pad)
+	recordPad := min(pad, baseRecordSize-1)
 
 	if i != AES128GCM {
-		recordPad = min((1<<(padSize*8))-1, recordPad)
+		recordPad = min(recordPad, (1<<(i.Padding()*8))-1)
 	}
 	if pad > 0 && recordPad == 0 {
 		recordPad++ // Deal with perverse case of rs=overhead+1 with padding.
@@ -63,14 +61,11 @@ func (i ContentEncoding) calculateCipherBlockEnd(gcm cipher.AEAD, start, content
 	tagSize := gcm.Overhead()
 	if i != AES128GCM {
 		blockSize += tagSize
+		if start+blockSize == contentLen {
+			return 0, ErrTruncated
+		}
 	}
-	end := start + blockSize
-
-	if i != AES128GCM && end == contentLen {
-		return 0, ErrTruncated
-	}
-
-	end = min(end, contentLen)
+	end := min(start+blockSize, contentLen)
 	if end-start <= tagSize {
 		return 0, ErrTruncated
 	}
@@ -103,8 +98,8 @@ func (i ContentEncoding) appendPadding(plaintext []byte, pad int, last bool) ([]
 func (i ContentEncoding) unpad(plaintext []byte, last bool) ([]byte, error) {
 	switch i {
 	case AES128GCM:
-		for i := len(plaintext) - 1; i >= 0; i-- {
-			c := plaintext[i]
+		for j := len(plaintext) - 1; j >= 0; j-- {
+			c := plaintext[j]
 			switch {
 			case c == 0:
 				continue
@@ -113,7 +108,7 @@ func (i ContentEncoding) unpad(plaintext []byte, last bool) ([]byte, error) {
 			case !last && c != 1:
 				return nil, ErrInvalidPaddingNonLast
 			default:
-				return plaintext[:i], nil
+				return plaintext[:j], nil
 			}
 		}
 
@@ -136,14 +131,13 @@ func (i ContentEncoding) unpad(plaintext []byte, last bool) ([]byte, error) {
 }
 
 // isLastBlock returns true when last block
-func (i ContentEncoding) isLastBlock(pad, contentLen, blockEnd int) bool {
-	var last bool
-	switch i {
-	case AES128GCM:
-		last = blockEnd >= contentLen
-	default:
-		// The > here ensures that we write out a padding-only block at the end of a buffer.
-		last = blockEnd > contentLen
+func (i ContentEncoding) isLastBlock(pad, plaintextLen, blockEnd int) bool {
+	if pad != 0 {
+		return false
+	}
+	if i == AES128GCM {
+		return blockEnd >= plaintextLen
 	}
-	return last && pad == 0
+	// The > here ensures that we write out a padding-only block at the end of a buffer.
+	return blockEnd > plaintextLen
 }

contentencoding_test.go

@@ -16,4 +16,55 @@ import (
 func TestContentEncoding_Padding(t *testing.T) {
 	assert.Equal(t, 2, AESGCM.Padding())
 	assert.Equal(t, 1, AES128GCM.Padding())
+	assert.Equal(t, 0, ContentEncoding("unknown").Padding())
+}
+
+func TestContentEncoding_Unpad_Errors(t *testing.T) {
+	// AES128GCM error cases
+	_, err := AES128GCM.unpad([]byte{0, 0, 0}, true)
+	assert.Equal(t, ErrAllZeroPlaintext, err)
+
+	_, err = AES128GCM.unpad([]byte{0, 0, 3}, true)
+	assert.Equal(t, ErrInvalidPaddingLast, err)
+
+	_, err = AES128GCM.unpad([]byte{0, 0, 3}, false)
+	assert.Equal(t, ErrInvalidPaddingNonLast, err)
+
+	// Other (AESGCM) error cases
+	_, err = ContentEncoding("unknown").unpad([]byte{0, 0, 0}, true)
+	assert.Contains(t, err.Error(), "unknown padding size 0")
+
+	_, err = AESGCM.unpad([]byte{0, 10}, true)
+	assert.Contains(t, err.Error(), "padding exceeds block size: 12")
+}
+
+func TestContentEncoding_AppendPadding_Error(t *testing.T) {
+	_, err := AESGCM.appendPadding([]byte("test"), -1, true)
+	assert.Error(t, err)
+
+	_, err = AESGCM.appendPadding([]byte("test"), 70000, true)
+	assert.Error(t, err)
+}
+
+func TestContentEncoding_IsLastBlock(t *testing.T) {
+	assert.True(t, AES128GCM.isLastBlock(0, 10, 10))
+	assert.False(t, AES128GCM.isLastBlock(1, 10, 10))
+	assert.True(t, AESGCM.isLastBlock(0, 10, 11))
+	assert.False(t, AESGCM.isLastBlock(0, 10, 10))
+}
+
+func TestContentEncoding_Padding_Misc(t *testing.T) {
+	// rs=overhead+1 case in calculateRecordPadSize
+	gcm, _ := createCipher(make([]byte, 16))
+	_ = AES128GCM.overhead(gcm)
+	// baseRecordSize = 1
+	pad := AES128GCM.calculateRecordPadSize(10, 1)
+	assert.Equal(t, 1, pad)
+}
+
+func TestContentEncoding_CalculateCipherBlockEnd_Truncated(t *testing.T) {
+	gcm, _ := createCipher(make([]byte, 16))
+	// end-start <= tagSize case
+	_, err := AES128GCM.calculateCipherBlockEnd(gcm, 0, 10, 4096)
+	assert.Equal(t, ErrTruncated, err)
 }

decrypt.go

@@ -76,13 +76,20 @@ func Decrypt(content []byte, opts ...Option) ([]byte, error) {
 }
 
 func readHeader(opt *options, content []byte) []byte {
+	contentLen := len(content)
 	if opt.encoding == AES128GCM {
-		baseOffset := uint32(keyLen + recodeSizeLen)
-		idLen := uint32(content[baseOffset])
+		baseOffset := keyLen + recodeSizeLen
+		if contentLen <= baseOffset {
+			return content
+		}
+		idLen := int(content[baseOffset])
 
 		opt.salt = content[0:keyLen]
 		opt.recordSize = binary.BigEndian.Uint32(content[keyLen:baseOffset])
 		baseOffset++
+		if contentLen < baseOffset+idLen {
+			return content
+		}
 		opt.keyID = content[baseOffset : baseOffset+idLen]
 
 		return content[baseOffset+idLen:]
@@ -91,7 +98,7 @@ func readHeader(opt *options, content []byte) []byte {
 }
 
 func decryptRecord(opt *options, gcm cipher.AEAD, nonce []byte, content []byte, last bool) ([]byte, error) {
-	result, err := gcm.Open(nil, nonce, content, nil)
+	result, err := gcm.Open(content[:0], nonce, content, nil)
 	if err != nil {
 		return nil, err
 	}

decrypt_test.go

@@ -32,6 +32,27 @@ func TestDecryptWithAESGCM(t *testing.T) {
 	assert.Equal(t, "hello world", string(plaintext))
 }
 
+func TestDecryptWithAESGCM_Truncated(t *testing.T) {
+	authSecret := d(t, "9HcXsQe3xLMG/w2HsYKrOA==")
+	salt := d(t, "mRGYnIzSJGeZnJ19lgQcfw==")
+	privateKey := d(t, "yfSYB+/vCEoWklHCG7F99cQ1vRwemFYn87jZc8PHBwU=")
+	senderPublicKey := d(t, "BGJXZ4zDA04RfSgTufdauZXcNYbe3oF/yEri5ETSuZLDx70gYi7w2ytak8U82H01P1HYnIvr2fEeX7NZpeHdnhM=")
+
+	// BaseRecordSize for AESGCM (rs=4096) is 4096 - 2 = 4094.
+	// Total block size including tag is 4096 + 16 = 4112.
+	content := make([]byte, 4112)
+	plaintext, err := Decrypt(content,
+		WithEncoding(AESGCM),
+		WithSalt(salt),
+		WithAuthSecret(authSecret),
+		WithPrivate(privateKey),
+		WithDh(senderPublicKey),
+	)
+
+	assert.Equal(t, ErrTruncated, err)
+	assert.Nil(t, plaintext)
+}
+
 func TestDecryptWithAESGCM_2record(t *testing.T) {
 	authSecret := d(t, "9HcXsQe3xLMG/w2HsYKrOA==")
 	salt := d(t, "mRGYnIzSJGeZnJ19lgQcfw==")

encrypt.go

@@ -9,6 +9,7 @@ package httpece
 
 import (
 	"crypto/cipher"
+	"encoding/binary"
 	"fmt"
 	"math"
 )
@@ -19,7 +20,8 @@ func Encrypt(plaintext []byte, opts ...Option) ([]byte, error) {
 	var err error
 
 	// Options
-	if opt, err = parseOptions(encrypt, opts); err != nil {
+	opt, err = parseOptions(encrypt, opts)
+	if err != nil {
 		return nil, err
 	}
 
@@ -103,7 +105,7 @@ func encryptRecord(opt *options, gcm cipher.AEAD, nonce, plaintext []byte, recor
 	if err != nil {
 		return nil, err
 	}
-	return gcm.Seal(nil, nonce, plaintextWithPadding, nil), nil
+	return gcm.Seal(plaintextWithPadding[:0], nonce, plaintextWithPadding, nil), nil
 }
 
 func writeHeader(opt *options, results [][]byte) ([][]byte, error) {
@@ -119,7 +121,7 @@ func writeHeader(opt *options, results [][]byte) ([][]byte, error) {
 		}
 		buffer := make([]byte, saltLen+4+1+keyIDLen)
 		copy(buffer, opt.salt)
-		copy(buffer[saltLen:], uint32ToBytes(opt.recordSize))
+		binary.BigEndian.PutUint32(buffer[saltLen:], opt.recordSize)
 		buffer[saltLen+4] = uint8(keyIDLen)
 		copy(buffer[saltLen+5:], opt.keyID)
 		return append(results, buffer), nil

encrypt_test.go

@@ -35,6 +35,38 @@ func TestEncryptWithAESGCM(t *testing.T) {
 	assert.Equal(t, "vOjpVgZE4IYn/uEJKk3DzZ4X+Qr1dgSSUkuIzQE=", e(content))
 }
 
+func TestEncryptDecrypt_Errors(t *testing.T) {
+	// Encrypt error: record size too small
+	_, err := Encrypt([]byte("test"), WithRecordSize(1))
+	assert.Error(t, err)
+
+	// Decrypt error: record size too small
+	_, err = Decrypt([]byte("test"), WithRecordSize(1))
+	assert.Error(t, err)
+
+	// Encrypt error: invalid salt length
+	_, err = Encrypt([]byte("test"), WithSalt(make([]byte, 5)))
+	assert.Error(t, err)
+}
+
+func TestEncrypt_HeaderErrors(t *testing.T) {
+	// writeHeader with too long keyID
+	opt := &options{
+		encoding: AES128GCM,
+		keyID:    make([]byte, 300),
+	}
+	_, err := writeHeader(opt, nil)
+	assert.Error(t, err)
+
+	// writeHeader with too long salt
+	opt = &options{
+		encoding: AES128GCM,
+		salt:     make([]byte, 300),
+	}
+	_, err = writeHeader(opt, nil)
+	assert.Error(t, err)
+}
+
 func TestEncryptWithAESGCM4093Byte(t *testing.T) {
 	authSecret := d(t, "9HcXsQe3xLMG/w2HsYKrOA==")
 	salt := d(t, "mRGYnIzSJGeZnJ19lgQcfw==")

keys.go

@@ -8,6 +8,7 @@
 package httpece
 
 import (
+	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/ecdh"
@@ -176,14 +177,7 @@ func extractSecret(opt *options) ([]byte, error) {
 }
 
 func buildInfo(base []byte, context []byte) string {
-	baseLen := len(base)
-	contextLen := len(context)
-	result := make([]byte, 0, baseLen+contextLen)
-	result = append(result, base...)
-	if contextLen > 0 {
-		result = append(result, context...)
-	}
-	return string(result)
+	return string(bytes.Join([][]byte{base, context}, nil))
 }
 
 func extractDH(opt *options) (secret []byte, context []byte, err error) {