cryptobyte: add (*Builder).Unwrite and (*Builder).SetError

Unwrite allows programs to rollback builders more reliably and
efficiently than by copying a Builder (which might waste an allocation
and depends on internal behavior). This is useful for example to remove
a length-prefixed field if it ends up being empty.

SetError allows simple Builder extensions to set errors without making
MarshalingValue wrappers.

Based on the experience of CL 144115.

Change-Id: I9a785b81b51b15af49418b5bdb71c4ef222ccc46
Reviewed-on: https://go-review.googlesource.com/c/145317
Reviewed-by: Adam Langley <agl@golang.org>

Author: FiloSottile

cryptobyte/builder.go

@@ -50,6 +50,12 @@ func NewFixedBuilder(buffer []byte) *Builder {
 	}
 }
 
+// SetError sets the value to be returned as the error from Bytes. Writes
+// performed after calling SetError are ignored.
+func (b *Builder) SetError(err error) {
+	b.err = err
+}
+
 // Bytes returns the bytes written by the builder or an error if one has
 // occurred during building.
 func (b *Builder) Bytes() ([]byte, error) {
@@ -94,7 +100,7 @@ func (b *Builder) AddBytes(v []byte) {
 	b.add(v...)
 }
 
-// BuilderContinuation is continuation-passing interface for building
+// BuilderContinuation is a continuation-passing interface for building
 // length-prefixed byte sequences. Builder methods for length-prefixed
 // sequences (AddUint8LengthPrefixed etc) will invoke the BuilderContinuation
 // supplied to them. The child builder passed to the continuation can be used
@@ -290,6 +296,26 @@ func (b *Builder) add(bytes ...byte) {
 	b.result = append(b.result, bytes...)
 }
 
+// Unwrite rolls back n bytes written directly to the Builder. An attempt by a
+// child builder passed to a continuation to unwrite bytes from its parent will
+// panic.
+func (b *Builder) Unwrite(n int) {
+	if b.err != nil {
+		return
+	}
+	if b.child != nil {
+		panic("attempted unwrite while child is pending")
+	}
+	length := len(b.result) - b.pendingLenLen - b.offset
+	if length < 0 {
+		panic("cryptobyte: internal error")
+	}
+	if n > length {
+		panic("cryptobyte: attempted to unwrite more than was written")
+	}
+	b.result = b.result[:len(b.result)-n]
+}
+
 // A MarshalingValue marshals itself into a Builder.
 type MarshalingValue interface {
 	// Marshal is called by Builder.AddValue. It receives a pointer to a builder

cryptobyte/cryptobyte_test.go

@@ -327,12 +327,14 @@ func TestWriteWithPendingChild(t *testing.T) {
 	var b Builder
 	b.AddUint8LengthPrefixed(func(c *Builder) {
 		c.AddUint8LengthPrefixed(func(d *Builder) {
-			defer func() {
-				if recover() == nil {
-					t.Errorf("recover() = nil, want error; c.AddUint8() did not panic")
-				}
+			func() {
+				defer func() {
+					if recover() == nil {
+						t.Errorf("recover() = nil, want error; c.AddUint8() did not panic")
+					}
+				}()
+				c.AddUint8(2) // panics
 			}()
-			c.AddUint8(2) // panics
 
 			defer func() {
 				if recover() == nil {
@@ -351,6 +353,65 @@ func TestWriteWithPendingChild(t *testing.T) {
 	})
 }
 
+func TestSetError(t *testing.T) {
+	const errorStr = "TestSetError"
+	var b Builder
+	b.SetError(errors.New(errorStr))
+
+	ret, err := b.Bytes()
+	if ret != nil {
+		t.Error("expected nil result")
+	}
+	if err == nil {
+		t.Fatal("unexpected nil error")
+	}
+	if s := err.Error(); s != errorStr {
+		t.Errorf("expected error %q, got %v", errorStr, s)
+	}
+}
+
+func TestUnwrite(t *testing.T) {
+	var b Builder
+	b.AddBytes([]byte{1, 2, 3, 4, 5})
+	b.Unwrite(2)
+	if err := builderBytesEq(&b, 1, 2, 3); err != nil {
+		t.Error(err)
+	}
+
+	func() {
+		defer func() {
+			if recover() == nil {
+				t.Errorf("recover() = nil, want error; b.Unwrite() did not panic")
+			}
+		}()
+		b.Unwrite(4) // panics
+	}()
+
+	b = Builder{}
+	b.AddBytes([]byte{1, 2, 3, 4, 5})
+	b.AddUint8LengthPrefixed(func(b *Builder) {
+		b.AddBytes([]byte{1, 2, 3, 4, 5})
+
+		defer func() {
+			if recover() == nil {
+				t.Errorf("recover() = nil, want error; b.Unwrite() did not panic")
+			}
+		}()
+		b.Unwrite(6) // panics
+	})
+
+	b = Builder{}
+	b.AddBytes([]byte{1, 2, 3, 4, 5})
+	b.AddUint8LengthPrefixed(func(c *Builder) {
+		defer func() {
+			if recover() == nil {
+				t.Errorf("recover() = nil, want error; b.Unwrite() did not panic")
+			}
+		}()
+		b.Unwrite(2) // panics (attempted unwrite while child is pending)
+	})
+}
+
 // ASN.1
 
 func TestASN1Int64(t *testing.T) {