diff --git a/cryptobyte/builder.go b/cryptobyte/builder.go index d0c6c3f..ca7b1db 100644 --- a/cryptobyte/builder.go +++ b/cryptobyte/builder.go @@ -274,9 +274,11 @@ func (b *Builder) flushChild() { return } - if !b.fixedSize { - b.result = child.result // In case child reallocated result. + if b.fixedSize && &b.result[0] != &child.result[0] { + panic("cryptobyte: BuilderContinuation reallocated a fixed-size buffer") } + + b.result = child.result } func (b *Builder) add(bytes ...byte) { @@ -284,7 +286,7 @@ func (b *Builder) add(bytes ...byte) { return } if b.child != nil { - panic("attempted write while child is pending") + panic("cryptobyte: attempted write while child is pending") } if len(b.result)+len(bytes) < len(bytes) { b.err = errors.New("cryptobyte: length overflow") @@ -304,7 +306,7 @@ func (b *Builder) Unwrite(n int) { return } if b.child != nil { - panic("attempted unwrite while child is pending") + panic("cryptobyte: attempted unwrite while child is pending") } length := len(b.result) - b.pendingLenLen - b.offset if length < 0 { diff --git a/cryptobyte/cryptobyte_test.go b/cryptobyte/cryptobyte_test.go index b859cc9..fe658ef 100644 --- a/cryptobyte/cryptobyte_test.go +++ b/cryptobyte/cryptobyte_test.go @@ -412,6 +412,33 @@ func TestUnwrite(t *testing.T) { }) } +func TestFixedBuilderLengthPrefixed(t *testing.T) { + bufCap := 10 + inner := bytes.Repeat([]byte{0xff}, bufCap-2) + buf := make([]byte, 0, bufCap) + b := NewFixedBuilder(buf) + b.AddUint16LengthPrefixed(func(b *Builder) { + b.AddBytes(inner) + }) + if got := b.BytesOrPanic(); len(got) != bufCap { + t.Errorf("Expected output lenght to be %d, got %d", bufCap, len(got)) + } +} + +func TestFixedBuilderPanicReallocate(t *testing.T) { + defer func() { + recover() + }() + + b := NewFixedBuilder(make([]byte, 0, 10)) + b1 := NewFixedBuilder(make([]byte, 0, 10)) + b.AddUint16LengthPrefixed(func(b *Builder) { + *b = *b1 + }) + + t.Error("Builder did not panic") +} + // ASN.1 func TestASN1Int64(t *testing.T) {