Add password protected writing

crypto.go

@@ -9,6 +9,7 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/hmac"
+	"crypto/rand"
 	"crypto/sha1"
 	"crypto/subtle"
 	"errors"
@@ -212,7 +213,7 @@ func (a *bufferedAuthReader) Read(b []byte) (int, error) {
 			a.err = io.ErrUnexpectedEOF
 			return 0, a.err
 		}
-		ab := new(bytes.Buffer)
+		ab := new(bytes.Buffer) // remove this buffer and io.Copy to mac
 		nn, err := io.Copy(ab, a.adata)
 		if err != nil || nn != 10 {
 			a.err = io.ErrUnexpectedEOF
@@ -266,10 +267,6 @@ func newDecryptionReader(r *io.SectionReader, f *File) (io.Reader, error) {
 	if saltLen == 0 {
 		return nil, ErrDecryption
 	}
-	// Change to a streaming implementation
-	// Maybe not such a good idea after all. See:
-	// https://www.imperialviolet.org/2014/06/27/streamingencryption.html
-	// https://www.imperialviolet.org/2015/05/16/aeads.html
 	// grab the salt, pwvv, data, and authcode
 	saltpwvv := make([]byte, saltLen+2)
 	if _, err := r.Read(saltpwvv); err != nil {
@@ -328,3 +325,101 @@ func aesKeyLen(strength byte) int {
 		return 0
 	}
 }
+
+type authWriter struct {
+	hmac hash.Hash // from fw.hmac
+	w    io.Writer // this will be the compCount writer
+}
+
+func (aw *authWriter) Write(p []byte) (int, error) {
+	_, err := aw.hmac.Write(p)
+	if err != nil {
+		return 0, err
+	}
+	return aw.w.Write(p)
+}
+
+// writes out the salt, pwv, and then the encrypted file data
+type encryptionWriter struct {
+	pwv   []byte    // password verification code to be written
+	salt  []byte    // salt to be written
+	w     io.Writer // where to write the salt + pwv
+	es    io.Writer // where to write encrypted file data
+	first bool      // first write?
+	err   error     // last error
+}
+
+func (ew *encryptionWriter) Write(p []byte) (int, error) {
+	if ew.err != nil {
+		return 0, ew.err
+	}
+	if ew.first {
+		// if our first time writing
+		// must write out the salt and pwv first unencrypted
+		_, err1 := ew.w.Write(ew.salt)
+		_, err2 := ew.w.Write(ew.pwv)
+		if err1 != nil || err2 != nil {
+			ew.err = errors.New("zip: error writing salt or pwv")
+			return 0, ew.err
+		}
+		ew.first = false
+	}
+	// now just pass on to the encryption stream
+	return ew.es.Write(p)
+}
+
+// newEncryptionWriter returns a io.Writer that when written to, 1. writes
+// out the salt, 2. writes out pwv, 3. writes out encrypted the data, and finally
+// 4. will write to hmac.
+func newEncryptionWriter(w io.Writer, fh *FileHeader, fw *fileWriter) (io.Writer, error) {
+	var salt [16]byte
+	_, err := rand.Read(salt[:])
+	if err != nil {
+		return nil, errors.New("zip: unable to generate random salt")
+	}
+	ekey, akey, pwv := generateKeys(fh.Password(), salt[:], aes256)
+	fw.hmac = hmac.New(sha1.New, akey)
+	aw := &authWriter{
+		hmac: fw.hmac,
+		w:    w,
+	}
+	es, err := encryptStream(ekey, aw)
+	if err != nil {
+		return nil, err
+	}
+	ew := &encryptionWriter{
+		pwv:   pwv,
+		salt:  salt[:],
+		w:     w,
+		es:    es,
+		first: true,
+	}
+	return ew, nil
+}
+
+func encryptStream(key []byte, w io.Writer) (io.Writer, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, errors.New("zip: couldn't create AES cipher")
+	}
+	stream := newWinZipCTR(block)
+	writer := &cipher.StreamWriter{S: stream, W: w}
+	return writer, nil
+}
+
+func (fh *FileHeader) writeWinZipExtra() {
+	// total size is 11 bytes
+	var buf [11]byte
+	eb := writeBuf(buf[:])
+	eb.uint16(winzipAesExtraId)
+	eb.uint16(7)         // following data size is 7
+	eb.uint16(2)         // ae 2
+	eb.uint16(0x4541)    // "AE"
+	eb.uint8(3)          // aes256
+	eb.uint16(fh.Method) // original compression method
+	fh.Extra = append(fh.Extra, buf[:]...)
+}
+
+func (fh *FileHeader) setEncryptionBit() {
+	fh.Flags |= 0x1
+}

crypto_test.go

@@ -12,7 +12,7 @@ func pwFn() []byte {
 }
 
 // Test simple password reading.
-func TestPasswordSimple(t *testing.T) {
+func TestPasswordReadSimple(t *testing.T) {
 	file := "hello-aes.zip"
 	var buf bytes.Buffer
 	r, err := OpenReader(filepath.Join("testdata", file))
@@ -173,3 +173,52 @@ func TestPasswordTamperedData(t *testing.T) {
 		}
 	}
 }
+
+func TestPasswordWriteSimple(t *testing.T) {
+	contents := []byte("Hello World")
+	conLen := len(contents)
+
+	// Write a zip
+	fh := &FileHeader{
+		Name:     "hello.txt",
+		Password: pwFn,
+	}
+	raw := new(bytes.Buffer)
+	zipw := NewWriter(raw)
+	w, err := zipw.CreateHeader(fh)
+	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.Password = pwFn
+	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

@@ -93,8 +93,14 @@ type FileHeader struct {
 	ExternalAttrs      uint32 // Meaning depends on CreatorVersion
 	Comment            string
 
-	// encryption fields
+	// DeferAuth determines whether hmac checks happen before
-	Password    PasswordFn // The password to use when reading/writing
+	// any ciphertext is decrypted. It is recommended to leave this
+	// set to false. For more detail:
+	// https://www.imperialviolet.org/2014/06/27/streamingencryption.html
+	// https://www.imperialviolet.org/2015/05/16/aeads.html
+	DeferAuth bool
+
+	Password    PasswordFn // Returns the password to use when reading/writing
 	ae          uint16
 	aesStrength byte
 }

writer.go

@@ -211,7 +211,8 @@ func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
 	}
 
 	fh.Flags |= 0x8 // we will write a data descriptor
-
+	// TODO(alex): Look at spec and see if these need to be changed
+	// when using encryption.
 	fh.CreatorVersion = fh.CreatorVersion&0xff00 | zipVersion20 // preserve compatibility byte
 	fh.ReaderVersion = zipVersion20
 
@@ -220,12 +221,27 @@ func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
 		compCount: &countWriter{w: w.cw},
 		crc32:     crc32.NewIEEE(),
 	}
+	// Get the compressor before possibly changing Method to 99 due to password
 	comp := compressor(fh.Method)
 	if comp == nil {
 		return nil, ErrAlgorithm
 	}
+	// check for password
+	var sw io.Writer = fw.compCount
+	if fh.Password != nil {
+		// we have a password and need to encrypt.
+		// 1. Set encryption bit in fh.Flags
+		fh.setEncryptionBit()
+		fh.writeWinZipExtra()
+		fh.Method = 99 // ok to change, we've gotten the comp and wrote extra
+		ew, err := newEncryptionWriter(sw, fh, fw)
+		if err != nil {
+			return nil, errors.New("zip: unable to create an encryption writer")
+		}
+		sw = ew
+	}
 	var err error
-	fw.comp, err = comp(fw.compCount)
+	fw.comp, err = comp(sw)
 	if err != nil {
 		return nil, err
 	}
@@ -278,6 +294,8 @@ type fileWriter struct {
 	compCount *countWriter
 	crc32     hash.Hash32
 	closed    bool
+
+	hmac hash.Hash // possible hmac used for authentication when encrypting
 }
 
 func (w *fileWriter) Write(p []byte) (int, error) {
@@ -296,10 +314,21 @@ func (w *fileWriter) close() error {
 	if err := w.comp.Close(); err != nil {
 		return err
 	}
-
+	// if encrypted grab the hmac and write it out
+	if w.header.IsEncrypted() {
+		authCode := w.hmac.Sum(nil)
+		authCode = authCode[:10]
+		_, err := w.compCount.Write(authCode)
+		if err != nil {
+			return errors.New("zip: error writing authcode")
+		}
+	}
 	// update FileHeader
 	fh := w.header.FileHeader
+	// ae-2 we don't write out CRC
+	if !fh.IsEncrypted() {
 		fh.CRC32 = w.crc32.Sum32()
+	}
 	fh.CompressedSize64 = uint64(w.compCount.count)
 	fh.UncompressedSize64 = uint64(w.rawCount.count)
 
@@ -358,6 +387,11 @@ func (w nopCloser) Close() error {
 
 type writeBuf []byte
 
+func (b *writeBuf) uint8(v uint8) {
+	(*b)[0] = v
+	*b = (*b)[1:]
+}
+
 func (b *writeBuf) uint16(v uint16) {
 	binary.LittleEndian.PutUint16(*b, v)
 	*b = (*b)[2:]