1// Copyright 2017 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package cryptobyte
6
7import (
8	"errors"
9	"fmt"
10)
11
12// A Builder builds byte strings from fixed-length and length-prefixed values.
13// Builders either allocate space as needed, or are ‘fixed’, which means that
14// they write into a given buffer and produce an error if it's exhausted.
15//
16// The zero value is a usable Builder that allocates space as needed.
17//
18// Simple values are marshaled and appended to a Builder using methods on the
19// Builder. Length-prefixed values are marshaled by providing a
20// BuilderContinuation, which is a function that writes the inner contents of
21// the value to a given Builder. See the documentation for BuilderContinuation
22// for details.
23type Builder struct {
24	err            error
25	result         []byte
26	fixedSize      bool
27	child          *Builder
28	offset         int
29	pendingLenLen  int
30	pendingIsASN1  bool
31	inContinuation *bool
32}
33
34// NewBuilder creates a Builder that appends its output to the given buffer.
35// Like append(), the slice will be reallocated if its capacity is exceeded.
36// Use Bytes to get the final buffer.
37func NewBuilder(buffer []byte) *Builder {
38	return &Builder{
39		result: buffer,
40	}
41}
42
43// NewFixedBuilder creates a Builder that appends its output into the given
44// buffer. This builder does not reallocate the output buffer. Writes that
45// would exceed the buffer's capacity are treated as an error.
46func NewFixedBuilder(buffer []byte) *Builder {
47	return &Builder{
48		result:    buffer,
49		fixedSize: true,
50	}
51}
52
53// SetError sets the value to be returned as the error from Bytes. Writes
54// performed after calling SetError are ignored.
55func (b *Builder) SetError(err error) {
56	b.err = err
57}
58
59// Bytes returns the bytes written by the builder or an error if one has
60// occurred during building.
61func (b *Builder) Bytes() ([]byte, error) {
62	if b.err != nil {
63		return nil, b.err
64	}
65	return b.result[b.offset:], nil
66}
67
68// BytesOrPanic returns the bytes written by the builder or panics if an error
69// has occurred during building.
70func (b *Builder) BytesOrPanic() []byte {
71	if b.err != nil {
72		panic(b.err)
73	}
74	return b.result[b.offset:]
75}
76
77// AddUint8 appends an 8-bit value to the byte string.
78func (b *Builder) AddUint8(v uint8) {
79	b.add(byte(v))
80}
81
82// AddUint16 appends a big-endian, 16-bit value to the byte string.
83func (b *Builder) AddUint16(v uint16) {
84	b.add(byte(v>>8), byte(v))
85}
86
87// AddUint24 appends a big-endian, 24-bit value to the byte string. The highest
88// byte of the 32-bit input value is silently truncated.
89func (b *Builder) AddUint24(v uint32) {
90	b.add(byte(v>>16), byte(v>>8), byte(v))
91}
92
93// AddUint32 appends a big-endian, 32-bit value to the byte string.
94func (b *Builder) AddUint32(v uint32) {
95	b.add(byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
96}
97
98// AddUint48 appends a big-endian, 48-bit value to the byte string.
99func (b *Builder) AddUint48(v uint64) {
100	b.add(byte(v>>40), byte(v>>32), byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
101}
102
103// AddUint64 appends a big-endian, 64-bit value to the byte string.
104func (b *Builder) AddUint64(v uint64) {
105	b.add(byte(v>>56), byte(v>>48), byte(v>>40), byte(v>>32), byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
106}
107
108// AddBytes appends a sequence of bytes to the byte string.
109func (b *Builder) AddBytes(v []byte) {
110	b.add(v...)
111}
112
113// BuilderContinuation is a continuation-passing interface for building
114// length-prefixed byte sequences. Builder methods for length-prefixed
115// sequences (AddUint8LengthPrefixed etc) will invoke the BuilderContinuation
116// supplied to them. The child builder passed to the continuation can be used
117// to build the content of the length-prefixed sequence. For example:
118//
119//	parent := cryptobyte.NewBuilder()
120//	parent.AddUint8LengthPrefixed(func (child *Builder) {
121//	  child.AddUint8(42)
122//	  child.AddUint8LengthPrefixed(func (grandchild *Builder) {
123//	    grandchild.AddUint8(5)
124//	  })
125//	})
126//
127// It is an error to write more bytes to the child than allowed by the reserved
128// length prefix. After the continuation returns, the child must be considered
129// invalid, i.e. users must not store any copies or references of the child
130// that outlive the continuation.
131//
132// If the continuation panics with a value of type BuildError then the inner
133// error will be returned as the error from Bytes. If the child panics
134// otherwise then Bytes will repanic with the same value.
135type BuilderContinuation func(child *Builder)
136
137// BuildError wraps an error. If a BuilderContinuation panics with this value,
138// the panic will be recovered and the inner error will be returned from
139// Builder.Bytes.
140type BuildError struct {
141	Err error
142}
143
144// AddUint8LengthPrefixed adds a 8-bit length-prefixed byte sequence.
145func (b *Builder) AddUint8LengthPrefixed(f BuilderContinuation) {
146	b.addLengthPrefixed(1, false, f)
147}
148
149// AddUint16LengthPrefixed adds a big-endian, 16-bit length-prefixed byte sequence.
150func (b *Builder) AddUint16LengthPrefixed(f BuilderContinuation) {
151	b.addLengthPrefixed(2, false, f)
152}
153
154// AddUint24LengthPrefixed adds a big-endian, 24-bit length-prefixed byte sequence.
155func (b *Builder) AddUint24LengthPrefixed(f BuilderContinuation) {
156	b.addLengthPrefixed(3, false, f)
157}
158
159// AddUint32LengthPrefixed adds a big-endian, 32-bit length-prefixed byte sequence.
160func (b *Builder) AddUint32LengthPrefixed(f BuilderContinuation) {
161	b.addLengthPrefixed(4, false, f)
162}
163
164func (b *Builder) callContinuation(f BuilderContinuation, arg *Builder) {
165	if !*b.inContinuation {
166		*b.inContinuation = true
167
168		defer func() {
169			*b.inContinuation = false
170
171			r := recover()
172			if r == nil {
173				return
174			}
175
176			if buildError, ok := r.(BuildError); ok {
177				b.err = buildError.Err
178			} else {
179				panic(r)
180			}
181		}()
182	}
183
184	f(arg)
185}
186
187func (b *Builder) addLengthPrefixed(lenLen int, isASN1 bool, f BuilderContinuation) {
188	// Subsequent writes can be ignored if the builder has encountered an error.
189	if b.err != nil {
190		return
191	}
192
193	offset := len(b.result)
194	b.add(make([]byte, lenLen)...)
195
196	if b.inContinuation == nil {
197		b.inContinuation = new(bool)
198	}
199
200	b.child = &Builder{
201		result:         b.result,
202		fixedSize:      b.fixedSize,
203		offset:         offset,
204		pendingLenLen:  lenLen,
205		pendingIsASN1:  isASN1,
206		inContinuation: b.inContinuation,
207	}
208
209	b.callContinuation(f, b.child)
210	b.flushChild()
211	if b.child != nil {
212		panic("cryptobyte: internal error")
213	}
214}
215
216func (b *Builder) flushChild() {
217	if b.child == nil {
218		return
219	}
220	b.child.flushChild()
221	child := b.child
222	b.child = nil
223
224	if child.err != nil {
225		b.err = child.err
226		return
227	}
228
229	length := len(child.result) - child.pendingLenLen - child.offset
230
231	if length < 0 {
232		panic("cryptobyte: internal error") // result unexpectedly shrunk
233	}
234
235	if child.pendingIsASN1 {
236		// For ASN.1, we reserved a single byte for the length. If that turned out
237		// to be incorrect, we have to move the contents along in order to make
238		// space.
239		if child.pendingLenLen != 1 {
240			panic("cryptobyte: internal error")
241		}
242		var lenLen, lenByte uint8
243		if int64(length) > 0xfffffffe {
244			b.err = errors.New("pending ASN.1 child too long")
245			return
246		} else if length > 0xffffff {
247			lenLen = 5
248			lenByte = 0x80 | 4
249		} else if length > 0xffff {
250			lenLen = 4
251			lenByte = 0x80 | 3
252		} else if length > 0xff {
253			lenLen = 3
254			lenByte = 0x80 | 2
255		} else if length > 0x7f {
256			lenLen = 2
257			lenByte = 0x80 | 1
258		} else {
259			lenLen = 1
260			lenByte = uint8(length)
261			length = 0
262		}
263
264		// Insert the initial length byte, make space for successive length bytes,
265		// and adjust the offset.
266		child.result[child.offset] = lenByte
267		extraBytes := int(lenLen - 1)
268		if extraBytes != 0 {
269			child.add(make([]byte, extraBytes)...)
270			childStart := child.offset + child.pendingLenLen
271			copy(child.result[childStart+extraBytes:], child.result[childStart:])
272		}
273		child.offset++
274		child.pendingLenLen = extraBytes
275	}
276
277	l := length
278	for i := child.pendingLenLen - 1; i >= 0; i-- {
279		child.result[child.offset+i] = uint8(l)
280		l >>= 8
281	}
282	if l != 0 {
283		b.err = fmt.Errorf("cryptobyte: pending child length %d exceeds %d-byte length prefix", length, child.pendingLenLen)
284		return
285	}
286
287	if b.fixedSize && &b.result[0] != &child.result[0] {
288		panic("cryptobyte: BuilderContinuation reallocated a fixed-size buffer")
289	}
290
291	b.result = child.result
292}
293
294func (b *Builder) add(bytes ...byte) {
295	if b.err != nil {
296		return
297	}
298	if b.child != nil {
299		panic("cryptobyte: attempted write while child is pending")
300	}
301	if len(b.result)+len(bytes) < len(bytes) {
302		b.err = errors.New("cryptobyte: length overflow")
303	}
304	if b.fixedSize && len(b.result)+len(bytes) > cap(b.result) {
305		b.err = errors.New("cryptobyte: Builder is exceeding its fixed-size buffer")
306		return
307	}
308	b.result = append(b.result, bytes...)
309}
310
311// Unwrite rolls back non-negative n bytes written directly to the Builder.
312// An attempt by a child builder passed to a continuation to unwrite bytes
313// from its parent will panic.
314func (b *Builder) Unwrite(n int) {
315	if b.err != nil {
316		return
317	}
318	if b.child != nil {
319		panic("cryptobyte: attempted unwrite while child is pending")
320	}
321	length := len(b.result) - b.pendingLenLen - b.offset
322	if length < 0 {
323		panic("cryptobyte: internal error")
324	}
325	if n < 0 {
326		panic("cryptobyte: attempted to unwrite negative number of bytes")
327	}
328	if n > length {
329		panic("cryptobyte: attempted to unwrite more than was written")
330	}
331	b.result = b.result[:len(b.result)-n]
332}
333
334// A MarshalingValue marshals itself into a Builder.
335type MarshalingValue interface {
336	// Marshal is called by Builder.AddValue. It receives a pointer to a builder
337	// to marshal itself into. It may return an error that occurred during
338	// marshaling, such as unset or invalid values.
339	Marshal(b *Builder) error
340}
341
342// AddValue calls Marshal on v, passing a pointer to the builder to append to.
343// If Marshal returns an error, it is set on the Builder so that subsequent
344// appends don't have an effect.
345func (b *Builder) AddValue(v MarshalingValue) {
346	err := v.Marshal(b)
347	if err != nil {
348		b.err = err
349	}
350}
351