Clean up the API and writer code.

Added an Example in documentation for code usage.
Cleaned up documentation and README.txt.

README.txt

@@ -6,17 +6,19 @@ This package DOES NOT intend to implement the encryption methods
 mentioned in the original PKWARE spec (sections 6.0 and 7.0):
 https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT
 
-The process
+Status - Alpha. More tests and code clean up next.
-============
-1. hello.txt -> compressed -> encrypted -> .zip
-2. .zip -> decrypted -> decompressed -> hello.txt
 
 Roadmap
 ========
 Reading - Done.
-Writing - Starting.
+Writing - Done.
 Testing - Needs more.
 
+The process
+============
+1. hello.txt -> compressed -> encrypted -> .zip
+2. .zip -> decrypted -> decompressed -> hello.txt
+
 WinZip AES specifies
 =====================
 1. Encryption-Decryption w/ AES-CTR (128, 192, or 256 bits)

crypto.go

@@ -26,8 +26,22 @@ const (
 	aes256 = 32
 )
 
+func aesKeyLen(strength byte) int {
+	switch strength {
+	case 1:
+		return aes128
+	case 2:
+		return aes192
+	case 3:
+		return aes256
+	default:
+		return 0
+	}
+}
+
 // Encryption/Decryption Errors
 var (
+	ErrEncryption     = errors.New("zip: encryption error")
 	ErrDecryption     = errors.New("zip: decryption error")
 	ErrPassword       = errors.New("zip: invalid password")
 	ErrAuthentication = errors.New("zip: authentication failed")
@@ -38,12 +52,11 @@ var (
 // repeatedly encrypting an incrementing counter and
 // xoring the resulting stream of data with the input.
 
-// This is a reimplementation of Go's CTR mode to allow
+// This is a re-implementation of Go's CTR mode to allow
-// for little-endian, left-aligned uint32 counter. Go's
+// for a little-endian, left-aligned uint32 counter, which
-// cipher.NewCTR follows the NIST Standard SP 800-38A, pp 13-15
+// is required for WinZip AES encryption. Go's cipher.NewCTR
-// which has a big-endian, right-aligned counter. WinZip
+// follows the NIST Standard SP 800-38A, pp 13-15
-// AES requires the CTR mode to have a little-endian,
+// which has a big-endian, right-aligned counter.
-// left-aligned counter.
 
 type ctr struct {
 	b       cipher.Block
@@ -55,7 +68,7 @@ type ctr struct {
 const streamBufferSize = 512
 
 // NewWinZipCTR returns a Stream which encrypts/decrypts using the given Block in
-// counter mode. The counter is initially set to 1.
+// counter mode. The counter is initially set to 1 per WinZip AES.
 func newWinZipCTR(block cipher.Block) cipher.Stream {
 	bufSize := streamBufferSize
 	if bufSize < block.BlockSize() {
@@ -128,6 +141,29 @@ func xorBytes(dst, a, b []byte) int {
 	return n
 }
 
+// newAuthReader returns either a buffered or streaming authentication reader.
+// Buffered authentication is recommended. Streaming authentication is only
+// recommended if: 1. you buffer the data yourself and wait for authentication
+// before streaming to another source such as the network, or 2. you just don't
+// care about authenticating unknown ciphertext before use :).
+func newAuthReader(akey []byte, data, adata io.Reader, streaming bool) io.Reader {
+	ar := authReader{
+		data:  data,
+		adata: adata,
+		mac:   hmac.New(sha1.New, akey),
+		err:   nil,
+		auth:  false,
+	}
+	if streaming {
+		return &ar
+	}
+	return &bufferedAuthReader{
+		ar,
+		new(bytes.Buffer),
+	}
+}
+
+// Streaming authentication
 type authReader struct {
 	data  io.Reader // data to be authenticated
 	adata io.Reader // the authentication code to read
@@ -136,7 +172,6 @@ type authReader struct {
 	auth  bool
 }
 
-// Streaming authentication
 func (a *authReader) Read(p []byte) (int, error) {
 	if a.err != nil {
 		return 0, a.err
@@ -173,34 +208,12 @@ func (a *authReader) Read(p []byte) (int, error) {
 	return n, a.err
 }
 
-// newAuthReader returns either a buffered or streaming authentication reader.
+// buffered authentication
-// Buffered authentication is recommended. Streaming authentication is only
-// recommended if: 1. you buffer the data yourself and wait for authentication
-// before streaming to another source such as the network, or 2. you just don't
-// care about authenticating unknown ciphertext before use :).
-func newAuthReader(akey []byte, data, adata io.Reader, streaming bool) io.Reader {
-	ar := authReader{
-		data:  data,
-		adata: adata,
-		mac:   hmac.New(sha1.New, akey),
-		err:   nil,
-		auth:  false,
-	}
-	if streaming {
-		return &ar
-	}
-	return &bufferedAuthReader{
-		ar,
-		new(bytes.Buffer),
-	}
-}
-
 type bufferedAuthReader struct {
 	authReader
 	buf *bytes.Buffer // buffer to store data to authenticate
 }
 
-// buffered authentication
 func (a *bufferedAuthReader) Read(b []byte) (int, error) {
 	// check for sticky error
 	if a.err != nil {
@@ -275,10 +288,10 @@ func newDecryptionReader(r *io.SectionReader, f *File) (io.Reader, error) {
 	salt := saltpwvv[:saltLen]
 	pwvv := saltpwvv[saltLen : saltLen+2]
 	// generate keys
-	if f.Password == nil {
+	if f.password == nil {
 		return nil, ErrPassword
 	}
-	decKey, authKey, pwv := generateKeys(f.Password(), salt, keyLen)
+	decKey, authKey, pwv := generateKeys(f.password(), salt, keyLen)
 	// check password verifier (pwv)
 	// Change to use crypto/subtle for constant time comparison
 	if !checkPasswordVerification(pwv, pwvv) {
@@ -294,7 +307,7 @@ func newDecryptionReader(r *io.SectionReader, f *File) (io.Reader, error) {
 	authOff := dataOff + dataLen
 	authcode := io.NewSectionReader(r, authOff, 10)
 	// setup auth reader, (buffered)/streaming
-	ar := newAuthReader(authKey, data, authcode, false)
+	ar := newAuthReader(authKey, data, authcode, f.DeferAuth)
 	// return decryption reader
 	dr := decryptStream(decKey, ar)
 	if dr == nil {
@@ -313,19 +326,7 @@ func decryptStream(key []byte, ciphertext io.Reader) io.Reader {
 	return reader
 }
 
-func aesKeyLen(strength byte) int {
+// writes encrypted data to hmac as it passes through
-	switch strength {
-	case 1:
-		return aes128
-	case 2:
-		return aes192
-	case 3:
-		return aes256
-	default:
-		return 0
-	}
-}
-
 type authWriter struct {
 	hmac hash.Hash // from fw.hmac
 	w    io.Writer // this will be the compCount writer
@@ -344,8 +345,8 @@ 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
+	es    io.Writer // where to write plaintext
-	first bool      // first write?
+	first bool      // first write
 	err   error     // last error
 }
 
@@ -368,16 +369,26 @@ func (ew *encryptionWriter) Write(p []byte) (int, error) {
 	return ew.es.Write(p)
 }
 
-// newEncryptionWriter returns a io.Writer that when written to, 1. writes
+func encryptStream(key []byte, w io.Writer) (io.Writer, error) {
-// out the salt, 2. writes out pwv, 3. writes out encrypted the data, and finally
+	block, err := aes.NewCipher(key)
-// 4. will write to hmac.
+	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
+}
+
+// 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, 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)
+	ekey, akey, pwv := generateKeys(fh.password(), salt[:], aes256)
 	fw.hmac = hmac.New(sha1.New, akey)
 	aw := &authWriter{
 		hmac: fw.hmac,
@@ -397,29 +408,56 @@ func newEncryptionWriter(w io.Writer, fh *FileHeader, fw *fileWriter) (io.Writer
 	return ew, nil
 }
 
-func encryptStream(key []byte, w io.Writer) (io.Writer, error) {
+// IsEncrypted indicates whether this file's data is encrypted.
-	block, err := aes.NewCipher(key)
+func (h *FileHeader) IsEncrypted() bool {
-	if err != nil {
+	return h.Flags&0x1 == 1
-		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() {
+// WinZip AE-2 specifies that no CRC value is written and
+// should be skipped when reading.
+func (h *FileHeader) isAE2() bool {
+	return h.ae == 2
+}
+
+func (h *FileHeader) writeWinZipExtra() {
 	// total size is 11 bytes
 	var buf [11]byte
 	eb := writeBuf(buf[:])
-	eb.uint16(winzipAesExtraId)
+	eb.uint16(winzipAesExtraId) // 0x9901
-	eb.uint16(7)         // following data size is 7
+	eb.uint16(7)                // following data size is 7
-	eb.uint16(2)         // ae 2
+	eb.uint16(2)                // ae 2
-	eb.uint16(0x4541)    // "AE"
+	eb.uint16(0x4541)           // "AE"
-	eb.uint8(3)          // aes256
+	eb.uint8(3)                 // aes256
-	eb.uint16(fh.Method) // original compression method
+	eb.uint16(h.Method)         // original compression method
-	fh.Extra = append(fh.Extra, buf[:]...)
+	h.Extra = append(h.Extra, buf[:]...)
 }
 
-func (fh *FileHeader) setEncryptionBit() {
+func (h *FileHeader) setEncryptionBit() {
-	fh.Flags |= 0x1
+	h.Flags |= 0x1
+}
+
+// SetPassword sets the password used for encryption/decryption.
+func (h *FileHeader) SetPassword(password string) {
+	if !h.IsEncrypted() {
+		h.setEncryptionBit()
+	}
+	h.password = func() []byte {
+		return []byte(password)
+	}
+}
+
+// PasswordFn is a function that returns the password
+// as a byte slice
+type passwordFn func() []byte
+
+// Encrypt is similar to Create except that it will encrypt the file contents
+// using AES-256 with the given password. Must follow all the same constraints
+// as Create.
+func (w *Writer) Encrypt(name string, password string) (io.Writer, error) {
+	fh := &FileHeader{
+		Name:   name,
+		Method: Deflate,
+	}
+	fh.SetPassword(password)
+	return w.CreateHeader(fh)
 }

crypto_test.go

@@ -7,10 +7,6 @@ import (
 	"testing"
 )
 
-func pwFn() []byte {
-	return []byte("golang")
-}
-
 // Test simple password reading.
 func TestPasswordReadSimple(t *testing.T) {
 	file := "hello-aes.zip"
@@ -30,7 +26,7 @@ func TestPasswordReadSimple(t *testing.T) {
 	if f.Method != 0 {
 		t.Errorf("Expected %s to have its Method set to 0.", file)
 	}
-	f.Password = pwFn
+	f.SetPassword("golang")
 	rc, err := f.Open()
 	if err != nil {
 		t.Errorf("Expected to open the readcloser: %v.", err)
@@ -45,6 +41,7 @@ func TestPasswordReadSimple(t *testing.T) {
 }
 
 // Test for multi-file password protected zip.
+// Each file can be protected with a different password.
 func TestPasswordHelloWorldAes(t *testing.T) {
 	file := "world-aes.zip"
 	expecting := "helloworld"
@@ -61,7 +58,7 @@ func TestPasswordHelloWorldAes(t *testing.T) {
 		if !f.IsEncrypted() {
 			t.Errorf("Expected %s to be encrypted.", f.FileInfo().Name)
 		}
-		f.Password = pwFn
+		f.SetPassword("golang")
 		rc, err := f.Open()
 		if err != nil {
 			t.Errorf("Expected to open readcloser: %v", err)
@@ -91,7 +88,7 @@ func TestPasswordMacbethAct1(t *testing.T) {
 		if !f.IsEncrypted() {
 			t.Errorf("Expected %s to be encrypted.", f.Name)
 		}
-		f.Password = pwFn
+		f.SetPassword("golang")
 		rc, err := f.Open()
 		if err != nil {
 			t.Errorf("Expected to open readcloser: %v", err)
@@ -131,7 +128,7 @@ func TestPasswordAE1BadCRC(t *testing.T) {
 		if !f.IsEncrypted() {
 			t.Errorf("Expected zip to be encrypted")
 		}
-		f.Password = pwFn
+		f.SetPassword("golang")
 		rc, err := f.Open()
 		if err != nil {
 			t.Errorf("Expected the readcloser to open.")
@@ -162,7 +159,7 @@ func TestPasswordTamperedData(t *testing.T) {
 		if !f.IsEncrypted() {
 			t.Errorf("Expected zip to be encrypted")
 		}
-		f.Password = pwFn
+		f.SetPassword("golang")
 		rc, err := f.Open()
 		if err != nil {
 			t.Errorf("Expected the readcloser to open.")
@@ -178,14 +175,9 @@ 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)
+	w, err := zipw.Encrypt("hello.txt", "golang")
 	if err != nil {
 		t.Errorf("Expected to create a new FileHeader")
 	}
@@ -206,7 +198,7 @@ func TestPasswordWriteSimple(t *testing.T) {
 		t.Errorf("Expected to have one file in the zip archive, but has %d files", nn)
 	}
 	z := zipr.File[0]
-	z.Password = pwFn
+	z.SetPassword("golang")
 	rr, err := z.Open()
 	if err != nil {
 		t.Errorf("Expected to open the readcloser: %v", err)

example_test.go

@@ -58,9 +58,6 @@ func ExampleReader() {
 	// Iterate through the files in the archive,
 	// printing some of their contents.
 	for _, f := range r.File {
-		// if f.IsEncrypted() {
-		// 	f.SetPassword([]byte("password"))
-		// }
 		fmt.Printf("Contents of %s:\n", f.Name)
 		rc, err := f.Open()
 		if err != nil {
@@ -77,3 +74,40 @@ func ExampleReader() {
 	// Contents of README:
 	// This is the source code repository for the Go programming language.
 }
+
+func ExampleWriter_Encrypt() {
+	contents := []byte("Hello World")
+
+	// write a password zip
+	raw := new(bytes.Buffer)
+	zipw := zip.NewWriter(raw)
+	w, err := zipw.Encrypt("hello.txt", "golang")
+	if err != nil {
+		log.Fatal(err)
+	}
+	_, err = io.Copy(w, bytes.NewReader(contents))
+	if err != nil {
+		log.Fatal(err)
+	}
+	zipw.Close()
+
+	// read the password zip
+	zipr, err := zip.NewReader(bytes.NewReader(raw.Bytes()), int64(raw.Len()))
+	if err != nil {
+		log.Fatal(err)
+	}
+	for _, z := range zipr.File {
+		z.SetPassword("golang")
+		rr, err := z.Open()
+		if err != nil {
+			log.Fatal(err)
+		}
+		_, err = io.Copy(os.Stdout, rr)
+		if err != nil {
+			log.Fatal(err)
+		}
+		rr.Close()
+	}
+	// Output:
+	// Hello World
+}

struct.go

@@ -17,7 +17,7 @@ for normal archives both fields will be the same. For files requiring
 the ZIP64 format the 32 bit fields will be 0xffffffff and the 64 bit
 fields must be used instead.
 
-Can read/write AES encrypted files that use Winzip's AES encryption method.
+Can read/write password protected files that use Winzip's AES encryption method.
 See: http://www.winzip.com/aes_info.htm
 */
 package zip
@@ -93,31 +93,19 @@ type FileHeader struct {
 	ExternalAttrs      uint32 // Meaning depends on CreatorVersion
 	Comment            string
 
-	// DeferAuth determines whether hmac checks happen before
+	// DeferAuth being set to true will delay hmac auth/integrity
-	// any ciphertext is decrypted. It is recommended to leave this
+	// checks when decrypting a file meaning the reader will be
-	// set to false. For more detail:
+	// getting unauthenticated plaintext. 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
+	password    passwordFn // Returns the password to use when reading/writing
 	ae          uint16
 	aesStrength byte
 }
 
-// PasswordFn is a function that returns the password
-// as a byte slice
-type PasswordFn func() []byte
-
-// IsEncrypted indicates whether this file's data is encrypted.
-func (f *FileHeader) IsEncrypted() bool {
-	return f.Flags&0x1 == 1
-}
-
-func (f *FileHeader) isAE2() bool {
-	return f.ae == 2
-}
-
 // FileInfo returns an os.FileInfo for the FileHeader.
 func (h *FileHeader) FileInfo() os.FileInfo {
 	return headerFileInfo{h}

writer.go

@@ -228,10 +228,8 @@ func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
 	}
 	// check for password
 	var sw io.Writer = fw.compCount
-	if fh.Password != nil {
+	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)