Selectable Encryption Method

crypto.go

@@ -19,7 +19,14 @@ import (
 	"golang.org/x/crypto/pbkdf2"
 )
 
+type EncryptionMethod int
+
 const (
+	ZipStandardEncryption EncryptionMethod = 1
+	AES128Encryption EncryptionMethod = 2
+	AES192Encryption EncryptionMethod = 3
+	AES256Encryption EncryptionMethod = 4
+
 	// AES key lengths
 	aes128 = 16
 	aes192 = 24
@@ -379,13 +386,14 @@ func encryptStream(key []byte, w io.Writer) (io.Writer, error) {
 // newEncryptionWriter returns an io.Writer that when written to, 1. writes
 // out the salt, 2. writes out pwv, and 3. writes out authenticated, encrypted
 // data. The authcode will be written out in fileWriter.close().
-func newEncryptionWriter(w io.Writer, password passwordFn, fw *fileWriter) (io.Writer, error) {
-	var salt [16]byte
+func newEncryptionWriter(w io.Writer, password passwordFn, fw *fileWriter, aesstrength byte) (io.Writer, error) {
+	keysize := aesKeyLen(aesstrength)
+	salt := make([]byte, keysize / 2)
 	_, err := rand.Read(salt[:])
 	if err != nil {
 		return nil, errors.New("zip: unable to generate random salt")
 	}
-	ekey, akey, pwv := generateKeys(password(), salt[:], aes256)
+	ekey, akey, pwv := generateKeys(password(), salt[:], keysize)
 	fw.hmac = hmac.New(sha1.New, akey)
 	aw := &authWriter{
 		hmac: fw.hmac,
@@ -424,11 +432,23 @@ func (h *FileHeader) writeWinZipExtra() {
 	eb.uint16(7)                // following data size is 7
 	eb.uint16(2)                // ae 2
 	eb.uint16(0x4541)           // "AE"
-	eb.uint8(3)                 // aes256
+	eb.uint8(h.aesStrength)     // aes256
 	eb.uint16(h.Method)         // original compression method
 	h.Extra = append(h.Extra, buf[:]...)
 }
 
+func (h *FileHeader) setEncryptionMethod(enc EncryptionMethod) {
+	h.encryption = enc
+	switch enc {
+	case AES128Encryption:
+		h.aesStrength = 1
+	case AES192Encryption:
+		h.aesStrength = 2
+	case AES256Encryption:
+		h.aesStrength = 3
+	}
+}
+
 func (h *FileHeader) setEncryptionBit() {
 	h.Flags |= 0x1
 }
@@ -452,11 +472,12 @@ type passwordFn func() []byte
 // contents will be encrypted with AES-256 using the given password. The
 // file's contents must be written to the io.Writer before the next call
 // to Create, CreateHeader, or Close.
-func (w *Writer) Encrypt(name string, password string) (io.Writer, error) {
+func (w *Writer) Encrypt(name string, password string, enc EncryptionMethod) (io.Writer, error) {
 	fh := &FileHeader{
 		Name:   name,
 		Method: Deflate,
 	}
 	fh.SetPassword(password)
+	fh.setEncryptionMethod(enc)
 	return w.CreateHeader(fh)
 }

crypto_test.go

@@ -175,42 +175,44 @@ func TestPasswordWriteSimple(t *testing.T) {
 	contents := []byte("Hello World")
 	conLen := len(contents)
 
-	raw := new(bytes.Buffer)
-	zipw := NewWriter(raw)
-	w, err := zipw.Encrypt("hello.txt", "golang")
-	if err != nil {
-		t.Errorf("Expected to create a new FileHeader")
-	}
-	n, err := io.Copy(w, bytes.NewReader(contents))
-	if err != nil || n != int64(conLen) {
-		t.Errorf("Expected to write the full contents to the writer.")
-	}
-	zipw.Close()
+	for _, enc := range []EncryptionMethod{AES128Encryption, AES192Encryption, AES256Encryption} {
+		raw := new(bytes.Buffer)
+		zipw := NewWriter(raw)
+		w, err := zipw.Encrypt("hello.txt", "golang", enc)
+		if err != nil {
+			t.Errorf("Expected to create a new FileHeader")
+		}
+		n, err := io.Copy(w, bytes.NewReader(contents))
+		if err != nil || n != int64(conLen) {
+			t.Errorf("Expected to write the full contents to the writer.")
+		}
+		zipw.Close()
 
-	// Read the zip
-	buf := new(bytes.Buffer)
-	zipr, err := NewReader(bytes.NewReader(raw.Bytes()), int64(raw.Len()))
-	if err != nil {
-		t.Errorf("Expected to open a new zip reader: %v", err)
-	}
-	nn := len(zipr.File)
-	if nn != 1 {
-		t.Errorf("Expected to have one file in the zip archive, but has %d files", nn)
-	}
-	z := zipr.File[0]
-	z.SetPassword("golang")
-	rr, err := z.Open()
-	if err != nil {
-		t.Errorf("Expected to open the readcloser: %v", err)
-	}
-	n, err = io.Copy(buf, rr)
-	if err != nil {
-		t.Errorf("Expected to write to temporary buffer: %v", err)
-	}
-	if n != int64(conLen) {
-		t.Errorf("Expected to copy %d bytes to temp buffer, but copied %d bytes instead", conLen, n)
-	}
-	if !bytes.Equal(contents, buf.Bytes()) {
-		t.Errorf("Expected the unzipped contents to equal '%s', but was '%s' instead", contents, buf.Bytes())
+		// Read the zip
+		buf := new(bytes.Buffer)
+		zipr, err := NewReader(bytes.NewReader(raw.Bytes()), int64(raw.Len()))
+		if err != nil {
+			t.Errorf("Expected to open a new zip reader: %v", err)
+		}
+		nn := len(zipr.File)
+		if nn != 1 {
+			t.Errorf("Expected to have one file in the zip archive, but has %d files", nn)
+		}
+		z := zipr.File[0]
+		z.SetPassword("golang")
+		rr, err := z.Open()
+		if err != nil {
+			t.Errorf("Expected to open the readcloser: %v", err)
+		}
+		n, err = io.Copy(buf, rr)
+		if err != nil {
+			t.Errorf("Expected to write to temporary buffer: %v", err)
+		}
+		if n != int64(conLen) {
+			t.Errorf("Expected to copy %d bytes to temp buffer, but copied %d bytes instead", conLen, n)
+		}
+		if !bytes.Equal(contents, buf.Bytes()) {
+			t.Errorf("Expected the unzipped contents to equal '%s', but was '%s' instead", contents, buf.Bytes())
+		}
 	}
 }

struct.go

@@ -101,6 +101,7 @@ type FileHeader struct {
 	// https://www.imperialviolet.org/2015/05/16/aeads.html
 	DeferAuth bool
 
+	encryption  EncryptionMethod
 	password    passwordFn // Returns the password to use when reading/writing
 	ae          uint16
 	aesStrength byte

writer.go

@@ -232,7 +232,7 @@ func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
 		// we have a password and need to encrypt.
 		fh.writeWinZipExtra()
 		fh.Method = 99 // ok to change, we've gotten the comp and wrote extra
-		ew, err := newEncryptionWriter(sw, fh.password, fw)
+		ew, err := newEncryptionWriter(sw, fh.password, fw, fh.aesStrength)
 		if err != nil {
 			return nil, err
 		}