builder.go

     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  
     5  package cryptobyte
     6  
     7  import (
     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.
    23  type 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.
    37  func 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.
    46  func 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.
    55  func (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.
    61  func (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.
    70  func (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.
    78  func (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.
    83  func (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.
    89  func (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.
    94  func (b *Builder) AddUint32(v uint32) {
    95  	b.add(byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
    96  }
    97  
    98  // AddBytes appends a sequence of bytes to the byte string.
    99  func (b *Builder) AddBytes(v []byte) {
   100  	b.add(v...)
   101  }
   102  
   103  // BuilderContinuation is a continuation-passing interface for building
   104  // length-prefixed byte sequences. Builder methods for length-prefixed
   105  // sequences (AddUint8LengthPrefixed etc) will invoke the BuilderContinuation
   106  // supplied to them. The child builder passed to the continuation can be used
   107  // to build the content of the length-prefixed sequence. For example:
   108  //
   109  //   parent := cryptobyte.NewBuilder()
   110  //   parent.AddUint8LengthPrefixed(func (child *Builder) {
   111  //     child.AddUint8(42)
   112  //     child.AddUint8LengthPrefixed(func (grandchild *Builder) {
   113  //       grandchild.AddUint8(5)
   114  //     })
   115  //   })
   116  //
   117  // It is an error to write more bytes to the child than allowed by the reserved
   118  // length prefix. After the continuation returns, the child must be considered
   119  // invalid, i.e. users must not store any copies or references of the child
   120  // that outlive the continuation.
   121  //
   122  // If the continuation panics with a value of type BuildError then the inner
   123  // error will be returned as the error from Bytes. If the child panics
   124  // otherwise then Bytes will repanic with the same value.
   125  type BuilderContinuation func(child *Builder)
   126  
   127  // BuildError wraps an error. If a BuilderContinuation panics with this value,
   128  // the panic will be recovered and the inner error will be returned from
   129  // Builder.Bytes.
   130  type BuildError struct {
   131  	Err error
   132  }
   133  
   134  // AddUint8LengthPrefixed adds a 8-bit length-prefixed byte sequence.
   135  func (b *Builder) AddUint8LengthPrefixed(f BuilderContinuation) {
   136  	b.addLengthPrefixed(1, false, f)
   137  }
   138  
   139  // AddUint16LengthPrefixed adds a big-endian, 16-bit length-prefixed byte sequence.
   140  func (b *Builder) AddUint16LengthPrefixed(f BuilderContinuation) {
   141  	b.addLengthPrefixed(2, false, f)
   142  }
   143  
   144  // AddUint24LengthPrefixed adds a big-endian, 24-bit length-prefixed byte sequence.
   145  func (b *Builder) AddUint24LengthPrefixed(f BuilderContinuation) {
   146  	b.addLengthPrefixed(3, false, f)
   147  }
   148  
   149  // AddUint32LengthPrefixed adds a big-endian, 32-bit length-prefixed byte sequence.
   150  func (b *Builder) AddUint32LengthPrefixed(f BuilderContinuation) {
   151  	b.addLengthPrefixed(4, false, f)
   152  }
   153  
   154  func (b *Builder) callContinuation(f BuilderContinuation, arg *Builder) {
   155  	if !*b.inContinuation {
   156  		*b.inContinuation = true
   157  
   158  		defer func() {
   159  			*b.inContinuation = false
   160  
   161  			r := recover()
   162  			if r == nil {
   163  				return
   164  			}
   165  
   166  			if buildError, ok := r.(BuildError); ok {
   167  				b.err = buildError.Err
   168  			} else {
   169  				panic(r)
   170  			}
   171  		}()
   172  	}
   173  
   174  	f(arg)
   175  }
   176  
   177  func (b *Builder) addLengthPrefixed(lenLen int, isASN1 bool, f BuilderContinuation) {
   178  	// Subsequent writes can be ignored if the builder has encountered an error.
   179  	if b.err != nil {
   180  		return
   181  	}
   182  
   183  	offset := len(b.result)
   184  	b.add(make([]byte, lenLen)...)
   185  
   186  	if b.inContinuation == nil {
   187  		b.inContinuation = new(bool)
   188  	}
   189  
   190  	b.child = &Builder{
   191  		result:         b.result,
   192  		fixedSize:      b.fixedSize,
   193  		offset:         offset,
   194  		pendingLenLen:  lenLen,
   195  		pendingIsASN1:  isASN1,
   196  		inContinuation: b.inContinuation,
   197  	}
   198  
   199  	b.callContinuation(f, b.child)
   200  	b.flushChild()
   201  	if b.child != nil {
   202  		panic("cryptobyte: internal error")
   203  	}
   204  }
   205  
   206  func (b *Builder) flushChild() {
   207  	if b.child == nil {
   208  		return
   209  	}
   210  	b.child.flushChild()
   211  	child := b.child
   212  	b.child = nil
   213  
   214  	if child.err != nil {
   215  		b.err = child.err
   216  		return
   217  	}
   218  
   219  	length := len(child.result) - child.pendingLenLen - child.offset
   220  
   221  	if length < 0 {
   222  		panic("cryptobyte: internal error") // result unexpectedly shrunk
   223  	}
   224  
   225  	if child.pendingIsASN1 {
   226  		// For ASN.1, we reserved a single byte for the length. If that turned out
   227  		// to be incorrect, we have to move the contents along in order to make
   228  		// space.
   229  		if child.pendingLenLen != 1 {
   230  			panic("cryptobyte: internal error")
   231  		}
   232  		var lenLen, lenByte uint8
   233  		if int64(length) > 0xfffffffe {
   234  			b.err = errors.New("pending ASN.1 child too long")
   235  			return
   236  		} else if length > 0xffffff {
   237  			lenLen = 5
   238  			lenByte = 0x80 | 4
   239  		} else if length > 0xffff {
   240  			lenLen = 4
   241  			lenByte = 0x80 | 3
   242  		} else if length > 0xff {
   243  			lenLen = 3
   244  			lenByte = 0x80 | 2
   245  		} else if length > 0x7f {
   246  			lenLen = 2
   247  			lenByte = 0x80 | 1
   248  		} else {
   249  			lenLen = 1
   250  			lenByte = uint8(length)
   251  			length = 0
   252  		}
   253  
   254  		// Insert the initial length byte, make space for successive length bytes,
   255  		// and adjust the offset.
   256  		child.result[child.offset] = lenByte
   257  		extraBytes := int(lenLen - 1)
   258  		if extraBytes != 0 {
   259  			child.add(make([]byte, extraBytes)...)
   260  			childStart := child.offset + child.pendingLenLen
   261  			copy(child.result[childStart+extraBytes:], child.result[childStart:])
   262  		}
   263  		child.offset++
   264  		child.pendingLenLen = extraBytes
   265  	}
   266  
   267  	l := length
   268  	for i := child.pendingLenLen - 1; i >= 0; i-- {
   269  		child.result[child.offset+i] = uint8(l)
   270  		l >>= 8
   271  	}
   272  	if l != 0 {
   273  		b.err = fmt.Errorf("cryptobyte: pending child length %d exceeds %d-byte length prefix", length, child.pendingLenLen)
   274  		return
   275  	}
   276  
   277  	if b.fixedSize && &b.result[0] != &child.result[0] {
   278  		panic("cryptobyte: BuilderContinuation reallocated a fixed-size buffer")
   279  	}
   280  
   281  	b.result = child.result
   282  }
   283  
   284  func (b *Builder) add(bytes ...byte) {
   285  	if b.err != nil {
   286  		return
   287  	}
   288  	if b.child != nil {
   289  		panic("cryptobyte: attempted write while child is pending")
   290  	}
   291  	if len(b.result)+len(bytes) < len(bytes) {
   292  		b.err = errors.New("cryptobyte: length overflow")
   293  	}
   294  	if b.fixedSize && len(b.result)+len(bytes) > cap(b.result) {
   295  		b.err = errors.New("cryptobyte: Builder is exceeding its fixed-size buffer")
   296  		return
   297  	}
   298  	b.result = append(b.result, bytes...)
   299  }
   300  
   301  // Unwrite rolls back n bytes written directly to the Builder. An attempt by a
   302  // child builder passed to a continuation to unwrite bytes from its parent will
   303  // panic.
   304  func (b *Builder) Unwrite(n int) {
   305  	if b.err != nil {
   306  		return
   307  	}
   308  	if b.child != nil {
   309  		panic("cryptobyte: attempted unwrite while child is pending")
   310  	}
   311  	length := len(b.result) - b.pendingLenLen - b.offset
   312  	if length < 0 {
   313  		panic("cryptobyte: internal error")
   314  	}
   315  	if n > length {
   316  		panic("cryptobyte: attempted to unwrite more than was written")
   317  	}
   318  	b.result = b.result[:len(b.result)-n]
   319  }
   320  
   321  // A MarshalingValue marshals itself into a Builder.
   322  type MarshalingValue interface {
   323  	// Marshal is called by Builder.AddValue. It receives a pointer to a builder
   324  	// to marshal itself into. It may return an error that occurred during
   325  	// marshaling, such as unset or invalid values.
   326  	Marshal(b *Builder) error
   327  }
   328  
   329  // AddValue calls Marshal on v, passing a pointer to the builder to append to.
   330  // If Marshal returns an error, it is set on the Builder so that subsequent
   331  // appends don't have an effect.
   332  func (b *Builder) AddValue(v MarshalingValue) {
   333  	err := v.Marshal(b)
   334  	if err != nil {
   335  		b.err = err
   336  	}
   337  }
   338
View as plain text