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  // AddUint48 appends a big-endian, 48-bit value to the byte string.
    99  func (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.
   104  func (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.
   109  func (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.
   135  type 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.
   140  type BuildError struct {
   141  	Err error
   142  }
   143  
   144  // AddUint8LengthPrefixed adds a 8-bit length-prefixed byte sequence.
   145  func (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.
   150  func (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.
   155  func (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.
   160  func (b *Builder) AddUint32LengthPrefixed(f BuilderContinuation) {
   161  	b.addLengthPrefixed(4, false, f)
   162  }
   163  
   164  func (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  
   187  func (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  
   216  func (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  
   294  func (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.
   314  func (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.
   335  type 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.
   345  func (b *Builder) AddValue(v MarshalingValue) {
   346  	err := v.Marshal(b)
   347  	if err != nil {
   348  		b.err = err
   349  	}
   350  }
   351
View as plain text