encoding/pem: change Encode, EncodeToMemory not to generate partial PEM blocks

Originally these routines could not fail except by
returning errors from the underlying writer.

Then we realized that header keys containing colons
needed to be rejected, and we started returning an error
from Encode. But that only happens after writing a
partial PEM block to the underlying writer, which is
unfortunate, but at least it was undocumented.

CL 77790 then documented this unfortunate behavior.

Instead of documenting unfortunate behavior, fix it.

Change-Id: Ic7467a576c4cecd16a99138571a1269cc4f96204
Reviewed-on: https://go-review.googlesource.com/82076
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

src/encoding/pem/pem.go

@@ -252,8 +252,18 @@ func writeHeader(out io.Writer, k, v string) error {
 	return err
 }
 
-// Encode writes the Block b to out.
+// Encode writes the PEM encoding of b to out.
 func Encode(out io.Writer, b *Block) error {
+	// Check for invalid block before writing any output.
+	for k := range b.Headers {
+		if strings.Contains(k, ":") {
+			return errors.New("pem: cannot encode a header key that contains a colon")
+		}
+	}
+
+	// All errors below are relayed from underlying io.Writer,
+	// so it is now safe to write data.
+
 	if _, err := out.Write(pemStart[1:]); err != nil {
 		return err
 	}
@@ -282,9 +292,6 @@ func Encode(out io.Writer, b *Block) error {
 		// For consistency of output, write other headers sorted by key.
 		sort.Strings(h)
 		for _, k := range h {
-			if strings.Contains(k, ":") {
-				return errors.New("pem: cannot encode a header key that contains a colon")
-			}
 			if err := writeHeader(out, k, b.Headers[k]); err != nil {
 				return err
 			}
@@ -311,12 +318,15 @@ func Encode(out io.Writer, b *Block) error {
 	return err
 }
 
-// EncodeToMemory returns the Block b.
+// EncodeToMemory returns the PEM encoding of b.
 //
-// EncodeToMemory will return an incomplete PEM encoded structure if an invalid block is given.
-// To catch errors, Blocks with user-supplied headers should use Encode.
+// If b has invalid headers and cannot be encoded,
+// EncodeToMemory returns nil. If it is important to
+// report details about this error case, use Encode instead.
 func EncodeToMemory(b *Block) []byte {
 	var buf bytes.Buffer
-	Encode(&buf, b)
+	if err := Encode(&buf, b); err != nil {
+		return nil
+	}
 	return buf.Bytes()
 }

src/encoding/pem/pem_test.go

@@ -590,3 +590,17 @@ N4XPksobn/NO2IDvPM7N9ZCe+aeyDEkE8QmP6mPScLuGvzSrsgOxWTMWF7Dbdzj0
 tJQLJRZ+ItT5Irl4owSEBNLahC1j3fhQavbj9WVAfKk=
 -----END RSA PRIVATE KEY-----
 `
+
+func TestBadEncode(t *testing.T) {
+	b := &Block{Type: "BAD", Headers: map[string]string{"X:Y": "Z"}}
+	var buf bytes.Buffer
+	if err := Encode(&buf, b); err == nil {
+		t.Fatalf("Encode did not report invalid header")
+	}
+	if buf.Len() != 0 {
+		t.Fatalf("Encode wrote data before reporting invalid header")
+	}
+	if data := EncodeToMemory(b); data != nil {
+		t.Fatalf("EncodeToMemory returned non-nil data")
+	}
+}