This closes qax-os#2237, fix panic when open encrypted workbook in some cases (qax-os#2238)

- Add more unit test for crypt, improve test code coverage

Co-authored-by: shuotan <[email protected]>

Author: sqdtss

crypt.go

@@ -145,7 +145,10 @@ func Decrypt(raw []byte, opts *Options) (packageBuf []byte, err error) {
 	if err != nil {
 		return
 	}
-	encryptionInfoBuf, encryptedPackageBuf := extractPart(doc)
+	var encryptionInfoBuf, encryptedPackageBuf []byte
+	if encryptionInfoBuf, encryptedPackageBuf, err = extractPart(doc); err != nil {
+		return
+	}
 	mechanism, err := encryptionMechanism(encryptionInfoBuf)
 	if err != nil || mechanism == "extensible" {
 		return
@@ -186,24 +189,25 @@ func Encrypt(raw []byte, opts *Options) ([]byte, error) {
 }
 
 // extractPart extract data from storage by specified part name.
-func extractPart(doc *mscfb.Reader) (encryptionInfoBuf, encryptedPackageBuf []byte) {
+func extractPart(doc *mscfb.Reader) ([]byte, []byte, error) {
+	var encryptionInfoBuf, encryptedPackageBuf []byte
 	for entry, err := doc.Next(); err == nil; entry, err = doc.Next() {
 		switch entry.Name {
 		case "EncryptionInfo":
 			buf := make([]byte, entry.Size)
-			i, _ := doc.Read(buf)
-			if i > 0 {
-				encryptionInfoBuf = buf
+			if _, err := doc.Read(buf); err != nil {
+				return encryptionInfoBuf, encryptedPackageBuf, err
 			}
+			encryptionInfoBuf = buf
 		case "EncryptedPackage":
 			buf := make([]byte, entry.Size)
-			i, _ := doc.Read(buf)
-			if i > 0 {
-				encryptedPackageBuf = buf
+			if _, err := doc.Read(buf); err != nil {
+				return encryptionInfoBuf, encryptedPackageBuf, err
 			}
+			encryptedPackageBuf = buf
 		}
 	}
-	return
+	return encryptionInfoBuf, encryptedPackageBuf, nil
 }
 
 // encryptionMechanism parse password-protected documents created mechanism.

crypt_test.go

@@ -13,6 +13,7 @@ package excelize
 
 import (
 	"bytes"
+	"encoding/base64"
 	"encoding/binary"
 	"os"
 	"path/filepath"
@@ -57,7 +58,8 @@ func TestEncrypt(t *testing.T) {
 
 	doc, err := mscfb.New(bytes.NewReader(raw))
 	assert.NoError(t, err)
-	encryptionInfoBuf, encryptedPackageBuf := extractPart(doc)
+	encryptionInfoBuf, encryptedPackageBuf, err := extractPart(doc)
+	assert.NoError(t, err)
 	binary.LittleEndian.PutUint64(encryptionInfoBuf[20:32], uint64(0))
 	_, err = standardDecrypt(encryptionInfoBuf, encryptedPackageBuf, &Options{Password: "password"})
 	assert.NoError(t, err)
@@ -73,6 +75,121 @@ func TestEncrypt(t *testing.T) {
 	assert.EqualError(t, err, "illegal base64 data at input byte 0")
 	_, err = createIV([]byte{0}, Encryption{KeyData: KeyData{SaltValue: "=="}})
 	assert.EqualError(t, err, "illegal base64 data at input byte 0")
+	// Test error handling for EncryptionInfo parse failure
+	compoundFile := &cfb{
+		paths:   []string{"Root Entry/"},
+		sectors: []sector{{name: "Root Entry", typeID: 5}},
+	}
+	compoundFile.put("EncryptionInfo", []byte{})
+	_, err = OpenReader(bytes.NewReader(compoundFile.write()))
+	assert.Equal(t, ErrWorkbookFileFormat, err)
+	// Test error handling for EncryptedPackage parse failure
+	compoundFile = &cfb{
+		paths:   []string{"Root Entry/"},
+		sectors: []sector{{name: "Root Entry", typeID: 5}},
+	}
+	encryptionInfo := make([]byte, 100)
+	binary.LittleEndian.PutUint16(encryptionInfo[:2], 4)
+	binary.LittleEndian.PutUint16(encryptionInfo[2:4], 4)
+	compoundFile.put("EncryptionInfo", encryptionInfo)
+	compoundFile.put("EncryptedPackage", []byte{})
+	_, err = OpenReader(bytes.NewReader(compoundFile.write()))
+	assert.Equal(t, ErrWorkbookFileFormat, err)
+	// Test createIV when iv length is less than block size
+	_, err = createIV(0, Encryption{
+		KeyData: KeyData{
+			HashAlgorithm: "md5",
+			BlockSize:     32,
+			SaltValue:     base64.StdEncoding.EncodeToString([]byte("")),
+		},
+	})
+	assert.NoError(t, err)
+	// Test decryptPackage error with padding
+	input := make([]byte, 18)
+	binary.LittleEndian.PutUint64(input[:8], 10)
+	for i := 8; i < 18; i++ {
+		input[i] = byte(i)
+	}
+	_, err = decryptPackage(make([]byte, 32), input, Encryption{
+		KeyData: KeyData{
+			HashAlgorithm: "sha256",
+			BlockSize:     16,
+			SaltValue:     base64.StdEncoding.EncodeToString([]byte("")),
+		},
+	})
+	assert.NoError(t, err)
+	// Test IV creation error with invalid salt
+	input = make([]byte, 4104)
+	binary.LittleEndian.PutUint64(input[:8], 4096)
+	_, err = decryptPackage(make([]byte, 32), input, Encryption{
+		KeyData: KeyData{
+			HashAlgorithm: "sha256",
+			BlockSize:     16,
+			SaltValue:     "==",
+		},
+	})
+	assert.Error(t, err)
+	// Test decrypt error with invalid key
+	_, err = decryptPackage([]byte{}, input, Encryption{
+		KeyData: KeyData{
+			HashAlgorithm: "sha256",
+			BlockSize:     16,
+			SaltValue:     base64.StdEncoding.EncodeToString([]byte("")),
+		},
+	})
+	assert.Error(t, err)
+	// Test put with path that is a prefix of name
+	compoundFile = &cfb{
+		paths:   []string{"Root Entry/"},
+		sectors: []sector{{name: "Root Entry", typeID: 5}},
+	}
+	compoundFile.put("Root Entry/Test", []byte(""))
+	compoundFile = &cfb{
+		paths:   []string{"Root"},
+		sectors: []sector{{name: "Root", typeID: 5}},
+	}
+	compoundFile.put("Test", []byte(""))
+	// Test compare function with different scenarios
+	compoundFile = &cfb{}
+	assert.Equal(t, -1, compoundFile.compare("Root/A", "Root/B"))
+	assert.Equal(t, 1, compoundFile.compare("Root/B", "Root/A"))
+	assert.Equal(t, -1, compoundFile.compare("Root", "Root/Child"))
+	// Test prepare with typeID == 0 sector
+	compoundFile = &cfb{
+		paths: []string{"Root Entry/", "Skip/", "Valid/"},
+		sectors: []sector{
+			{name: "Root Entry", typeID: 5},
+			{name: "Skip", typeID: 0},
+			{name: "Valid", typeID: 2},
+		},
+	}
+	compoundFile.prepare()
+	// Test locate with FATSectors incrementing but staying <= 109
+	compoundFile = &cfb{
+		paths:   []string{"Root Entry/"},
+		sectors: []sector{{name: "Root Entry", typeID: 5, content: []byte{}}},
+	}
+	numFiles := 1533
+	for i := range numFiles {
+		name := strings.Repeat("F", i%50+1)
+		compoundFile.paths = append(compoundFile.paths, name+"/")
+		compoundFile.sectors = append(compoundFile.sectors, sector{
+			name:    name,
+			typeID:  2,
+			content: make([]byte, 4096),
+		})
+	}
+	compoundFile.locate()
+	// Test writeDirectoryEntry with empty clsID
+	compoundFile = &cfb{
+		paths: []string{"Root Entry/", "File1/"},
+		sectors: []sector{
+			{name: "Root Entry", typeID: 5, content: []byte{}, clsID: headerCLSID},
+			{name: "File1", typeID: 2, content: []byte("test"), clsID: []byte{}},
+		},
+	}
+	compoundFile.stream = make([]byte, 10000)
+	compoundFile.writeDirectoryEntry([]int{1, 0, 1, 0, 1, 0, 0, 0})
 }
 
 func TestEncryptionMechanism(t *testing.T) {