internal/age,internal/format: implement armored file generation

cmd/age/age.go

@@ -24,15 +24,16 @@ func main() {
 	decryptFlag := flag.Bool("d", false, "decrypt the input")
 	outFlag := flag.String("o", "", "output to `FILE` (default stdout)")
 	inFlag := flag.String("i", "", "read from `FILE` (default stdin)")
+	armorFlag := flag.Bool("a", false, "generate an armored file")
 	flag.Parse()
 
 	switch {
 	case *generateFlag:
-		if *decryptFlag || *inFlag != "" {
+		if *decryptFlag || *inFlag != "" || *armorFlag {
 			log.Fatalf("Invalid flag combination")
 		}
 	case *decryptFlag:
-		if *generateFlag {
+		if *generateFlag || *armorFlag {
 			log.Fatalf("Invalid flag combination")
 		}
 	default: // encrypt
@@ -62,7 +63,7 @@ func main() {
 	case *decryptFlag:
 		decrypt(in, out)
 	default:
-		encrypt(in, out)
+		encrypt(in, out, *armorFlag)
 	}
 }
 
@@ -81,7 +82,7 @@ func generate(out io.Writer) {
 	fmt.Fprintf(out, "%s\n", k)
 }
 
-func encrypt(in io.Reader, out io.Writer) {
+func encrypt(in io.Reader, out io.Writer, armor bool) {
 	var recipients []age.Recipient
 	for _, arg := range flag.Args() {
 		r, err := parseRecipient(arg)
@@ -94,7 +95,11 @@ func encrypt(in io.Reader, out io.Writer) {
 		log.Fatalf("Missing recipients!")
 	}
 
-	w, err := age.Encrypt(out, recipients...)
+	ageEncrypt := age.Encrypt
+	if armor {
+		ageEncrypt = age.EncryptWithArmor
+	}
+	w, err := ageEncrypt(out, recipients...)
 	if err != nil {
 		log.Fatalf("Error initializing encryption: %v", err)
 	}

internal/age/age.go

@@ -36,6 +36,14 @@ type Recipient interface {
 }
 
 func Encrypt(dst io.Writer, recipients ...Recipient) (io.WriteCloser, error) {
+	return encrypt(dst, false, recipients...)
+}
+
+func EncryptWithArmor(dst io.Writer, recipients ...Recipient) (io.WriteCloser, error) {
+	return encrypt(dst, true, recipients...)
+}
+
+func encrypt(dst io.Writer, armor bool, recipients ...Recipient) (io.WriteCloser, error) {
 	if len(recipients) == 0 {
 		return nil, errors.New("no recipients specified")
 	}
@@ -45,7 +53,7 @@ func Encrypt(dst io.Writer, recipients ...Recipient) (io.WriteCloser, error) {
 		return nil, err
 	}
 
-	hdr := &format.Header{}
+	hdr := &format.Header{Armor: armor}
 	for i, r := range recipients {
 		if r.Type() == "scrypt" && len(recipients) != 1 {
 			return nil, errors.New("an scrypt recipient must be the only one")
@@ -66,15 +74,25 @@ func Encrypt(dst io.Writer, recipients ...Recipient) (io.WriteCloser, error) {
 		return nil, fmt.Errorf("failed to write header: %v", err)
 	}
 
+	var finalDst io.WriteCloser
+	if armor {
+		finalDst = format.ArmoredWriter(dst)
+	} else {
+		// stream.Writer takes a WriteCloser, and will propagate Close calls (so
+		// that the ArmoredWriter will get closed), but we don't want to expose
+		// that behavior to our caller.
+		finalDst = format.NopCloser(dst)
+	}
+
 	nonce := make([]byte, 16)
 	if _, err := rand.Read(nonce); err != nil {
 		return nil, err
 	}
-	if _, err := dst.Write(nonce); err != nil {
+	if _, err := finalDst.Write(nonce); err != nil {
 		return nil, fmt.Errorf("failed to write nonce: %v", err)
 	}
 
-	return stream.NewWriter(streamKey(fileKey, nonce), dst)
+	return stream.NewWriter(streamKey(fileKey, nonce), finalDst)
 }
 
 func Decrypt(src io.Reader, identities ...Identity) (io.Reader, error) {

internal/format/armor.go

@@ -0,0 +1,69 @@
+// Copyright 2019 Google LLC
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+package format
+
+import "io"
+
+import "encoding/base64"
+
+type newlineWriter struct {
+	dst     io.Writer
+	written int
+}
+
+func (w *newlineWriter) Write(p []byte) (n int, err error) {
+	for len(p) > 0 {
+		remainingInLine := columnsPerLine - (w.written % columnsPerLine)
+		if remainingInLine == columnsPerLine && w.written != 0 {
+			if _, err := w.dst.Write([]byte("\n")); err != nil {
+				return n, err
+			}
+		}
+		toWrite := remainingInLine
+		if toWrite > len(p) {
+			toWrite = len(p)
+		}
+		nn, err := w.dst.Write(p[:toWrite])
+		n += nn
+		w.written += nn
+		p = p[nn:]
+		if err != nil {
+			return n, err
+		}
+	}
+	return n, nil
+}
+
+type CloserFunc func() error
+
+func (f CloserFunc) Close() error { return f() }
+
+type nopCloser struct {
+	io.Writer
+}
+
+func (nopCloser) Close() error { return nil }
+
+func NopCloser(w io.Writer) io.WriteCloser { return nopCloser{w} }
+
+func ArmoredWriter(dst io.Writer) io.WriteCloser {
+	// TODO: write a test with aligned and misaligned sizes, and 8 and 10 steps.
+	w := base64.NewEncoder(b64, &newlineWriter{dst: dst})
+	return struct {
+		io.Writer
+		io.Closer
+	}{
+		Writer: w,
+		Closer: CloserFunc(func() error {
+			if err := w.Close(); err != nil {
+				return err
+			}
+			_, err := dst.Write([]byte("\n--- end of file ---\n"))
+			return err
+		}),
+	}
+}

internal/format/format.go

@@ -18,6 +18,7 @@ import (
 )
 
 type Header struct {
+	Armor      bool
 	Recipients []*Recipient
 	MAC        []byte
 }
@@ -40,9 +41,11 @@ func DecodeString(s string) ([]byte, error) {
 
 var EncodeToString = b64.EncodeToString
 
-const bytesPerLine = 56 / 4 * 3 // 56 columns of Base64
+const columnsPerLine = 56
+const bytesPerLine = columnsPerLine / 4 * 3
 
 const intro = "This is a file encrypted with age-tool.com, version 1\n"
+const introWithArmor = "This is an armored file encrypted with age-tool.com, version 1\n"
 
 var recipientPrefix = []byte("->")
 var footerPrefix = []byte("---")
@@ -59,22 +62,26 @@ func (r *Recipient) Marshal(w io.Writer) error {
 	if _, err := io.WriteString(w, "\n"); err != nil {
 		return err
 	}
-	for i := 0; i < len(r.Body); i += bytesPerLine {
+	ww := base64.NewEncoder(b64, &newlineWriter{dst: w})
-		n := bytesPerLine
+	if _, err := ww.Write(r.Body); err != nil {
-		if n > len(r.Body)-i {
+		return err
-			n = len(r.Body) - i
-		}
-		s := EncodeToString(r.Body[i : i+n])
-		if _, err := io.WriteString(w, s+"\n"); err != nil {
-			return err
-		}
 	}
-	return nil
+	if err := ww.Close(); err != nil {
+		return err
+	}
+	_, err := io.WriteString(w, "\n")
+	return err
 }
 
 func (h *Header) MarshalWithoutMAC(w io.Writer) error {
-	if _, err := io.WriteString(w, intro); err != nil {
+	if h.Armor {
-		return err
+		if _, err := io.WriteString(w, introWithArmor); err != nil {
+			return err
+		}
+	} else {
+		if _, err := io.WriteString(w, intro); err != nil {
+			return err
+		}
 	}
 	for _, r := range h.Recipients {
 		if err := r.Marshal(w); err != nil {

internal/stream/stream.go

@@ -4,6 +4,7 @@
 // license that can be found in the LICENSE file or at
 // https://developers.google.com/open-source/licenses/bsd
 
+// Package stream implements a variant of the STREAM chunked encryption scheme.
 package stream
 
 import (
@@ -131,14 +132,14 @@ func setLastChunkFlag(nonce *[chacha20poly1305.NonceSize]byte) {
 
 type Writer struct {
 	a         cipher.AEAD
-	dst       io.Writer
+	dst       io.WriteCloser
 	unwritten []byte // backed by buf
 	buf       [encChunkSize]byte
 	nonce     [chacha20poly1305.NonceSize]byte
 	err       error
 }
 
-func NewWriter(key []byte, dst io.Writer) (*Writer, error) {
+func NewWriter(key []byte, dst io.WriteCloser) (*Writer, error) {
 	aead, err := chacha20poly1305.New(key)
 	if err != nil {
 		return nil, err
@@ -177,6 +178,8 @@ func (w *Writer) Write(p []byte) (n int, err error) {
 	return total, nil
 }
 
+// Close will flush the last chunk and call the underlying
+// WriteCloser's Close method.
 func (w *Writer) Close() error {
 	if w.err != nil {
 		return w.err
@@ -185,10 +188,11 @@ func (w *Writer) Close() error {
 	err := w.flushChunk(lastChunk)
 	if err != nil {
 		w.err = err
-	} else {
+		return err
-		w.err = errors.New("stream.Writer is already closed")
 	}
-	return err
+	w.err = errors.New("stream.Writer is already closed")
+
+	return w.dst.Close()
 }
 
 const (