builder.mx raw

   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([]byte{:lenLen}...)
 195  
 196  	if b.inContinuation == nil {
 197  		v := false
 198  		b.inContinuation = &v
 199  	}
 200  
 201  	b.child = &Builder{
 202  		result:         b.result,
 203  		fixedSize:      b.fixedSize,
 204  		offset:         offset,
 205  		pendingLenLen:  lenLen,
 206  		pendingIsASN1:  isASN1,
 207  		inContinuation: b.inContinuation,
 208  	}
 209  
 210  	b.callContinuation(f, b.child)
 211  	b.flushChild()
 212  	if b.child != nil {
 213  		panic("cryptobyte: internal error")
 214  	}
 215  }
 216  
 217  func (b *Builder) flushChild() {
 218  	if b.child == nil {
 219  		return
 220  	}
 221  	b.child.flushChild()
 222  	child := b.child
 223  	b.child = nil
 224  
 225  	if child.err != nil {
 226  		b.err = child.err
 227  		return
 228  	}
 229  
 230  	length := len(child.result) - child.pendingLenLen - child.offset
 231  
 232  	if length < 0 {
 233  		panic("cryptobyte: internal error") // result unexpectedly shrunk
 234  	}
 235  
 236  	if child.pendingIsASN1 {
 237  		// For ASN.1, we reserved a single byte for the length. If that turned out
 238  		// to be incorrect, we have to move the contents along in order to make
 239  		// space.
 240  		if child.pendingLenLen != 1 {
 241  			panic("cryptobyte: internal error")
 242  		}
 243  		var lenLen, lenByte uint8
 244  		if int64(length) > 0xfffffffe {
 245  			b.err = errors.New("pending ASN.1 child too long")
 246  			return
 247  		} else if length > 0xffffff {
 248  			lenLen = 5
 249  			lenByte = 0x80 | 4
 250  		} else if length > 0xffff {
 251  			lenLen = 4
 252  			lenByte = 0x80 | 3
 253  		} else if length > 0xff {
 254  			lenLen = 3
 255  			lenByte = 0x80 | 2
 256  		} else if length > 0x7f {
 257  			lenLen = 2
 258  			lenByte = 0x80 | 1
 259  		} else {
 260  			lenLen = 1
 261  			lenByte = uint8(length)
 262  			length = 0
 263  		}
 264  
 265  		// Insert the initial length byte, make space for successive length bytes,
 266  		// and adjust the offset.
 267  		child.result[child.offset] = lenByte
 268  		extraBytes := int(lenLen - 1)
 269  		if extraBytes != 0 {
 270  			child.add([]byte{:extraBytes}...)
 271  			childStart := child.offset + child.pendingLenLen
 272  			copy(child.result[childStart+extraBytes:], child.result[childStart:])
 273  		}
 274  		child.offset++
 275  		child.pendingLenLen = extraBytes
 276  	}
 277  
 278  	l := length
 279  	for i := child.pendingLenLen - 1; i >= 0; i-- {
 280  		child.result[child.offset+i] = uint8(l)
 281  		l >>= 8
 282  	}
 283  	if l != 0 {
 284  		b.err = fmt.Errorf("cryptobyte: pending child length %d exceeds %d-byte length prefix", length, child.pendingLenLen)
 285  		return
 286  	}
 287  
 288  	if b.fixedSize && &b.result[0] != &child.result[0] {
 289  		panic("cryptobyte: BuilderContinuation reallocated a fixed-size buffer")
 290  	}
 291  
 292  	b.result = child.result
 293  }
 294  
 295  func (b *Builder) add(bytes ...byte) {
 296  	if b.err != nil {
 297  		return
 298  	}
 299  	if b.child != nil {
 300  		panic("cryptobyte: attempted write while child is pending")
 301  	}
 302  	if len(b.result)+len(bytes) < len(bytes) {
 303  		b.err = errors.New("cryptobyte: length overflow")
 304  	}
 305  	if b.fixedSize && len(b.result)+len(bytes) > cap(b.result) {
 306  		b.err = errors.New("cryptobyte: Builder is exceeding its fixed-size buffer")
 307  		return
 308  	}
 309  	b.result = append(b.result, bytes...)
 310  }
 311  
 312  // Unwrite rolls back non-negative n bytes written directly to the Builder.
 313  // An attempt by a child builder passed to a continuation to unwrite bytes
 314  // from its parent will panic.
 315  func (b *Builder) Unwrite(n int) {
 316  	if b.err != nil {
 317  		return
 318  	}
 319  	if b.child != nil {
 320  		panic("cryptobyte: attempted unwrite while child is pending")
 321  	}
 322  	length := len(b.result) - b.pendingLenLen - b.offset
 323  	if length < 0 {
 324  		panic("cryptobyte: internal error")
 325  	}
 326  	if n < 0 {
 327  		panic("cryptobyte: attempted to unwrite negative number of bytes")
 328  	}
 329  	if n > length {
 330  		panic("cryptobyte: attempted to unwrite more than was written")
 331  	}
 332  	b.result = b.result[:len(b.result)-n]
 333  }
 334  
 335  // A MarshalingValue marshals itself into a Builder.
 336  type MarshalingValue interface {
 337  	// Marshal is called by Builder.AddValue. It receives a pointer to a builder
 338  	// to marshal itself into. It may return an error that occurred during
 339  	// marshaling, such as unset or invalid values.
 340  	Marshal(b *Builder) error
 341  }
 342  
 343  // AddValue calls Marshal on v, passing a pointer to the builder to append to.
 344  // If Marshal returns an error, it is set on the Builder so that subsequent
 345  // appends don't have an effect.
 346  func (b *Builder) AddValue(v MarshalingValue) {
 347  	err := v.Marshal(b)
 348  	if err != nil {
 349  		b.err = err
 350  	}
 351  }
 352