diff --git a/cryptobyte/builder.go b/cryptobyte/builder.go index 765bf8e..d0c6c3f 100644 --- a/cryptobyte/builder.go +++ b/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 diff --git a/cryptobyte/cryptobyte_test.go b/cryptobyte/cryptobyte_test.go index f294dd5..b859cc9 100644 --- a/cryptobyte/cryptobyte_test.go +++ b/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) {