github.com/ava-labs/avalanchego@v1.11.11/codec/manager.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package codec
     5  
     6  import (
     7  	"errors"
     8  	"fmt"
     9  	"sync"
    10  
    11  	"github.com/ava-labs/avalanchego/utils/units"
    12  	"github.com/ava-labs/avalanchego/utils/wrappers"
    13  )
    14  
    15  const (
    16  	VersionSize = wrappers.ShortLen
    17  
    18  	// default max size, in bytes, of something being marshaled by Marshal()
    19  	defaultMaxSize = 256 * units.KiB
    20  
    21  	// initial capacity of byte slice that values are marshaled into.
    22  	// Larger value --> need less memory allocations but possibly have allocated but unused memory
    23  	// Smaller value --> need more memory allocations but more efficient use of allocated memory
    24  	initialSliceCap = 128
    25  )
    26  
    27  var (
    28  	ErrUnknownVersion    = errors.New("unknown codec version")
    29  	ErrMarshalNil        = errors.New("can't marshal nil pointer or interface")
    30  	ErrUnmarshalNil      = errors.New("can't unmarshal nil")
    31  	ErrUnmarshalTooBig   = errors.New("byte array exceeds maximum length")
    32  	ErrCantPackVersion   = errors.New("couldn't pack codec version")
    33  	ErrCantUnpackVersion = errors.New("couldn't unpack codec version")
    34  	ErrDuplicatedVersion = errors.New("duplicated codec version")
    35  )
    36  
    37  var _ Manager = (*manager)(nil)
    38  
    39  // Manager describes the functionality for managing codec versions.
    40  type Manager interface {
    41  	// Associate the given codec with the given version ID
    42  	RegisterCodec(version uint16, codec Codec) error
    43  
    44  	// Size returns the size, in bytes, of [value] when it's marshaled
    45  	// using the codec with the given version.
    46  	// RegisterCodec must have been called with that version.
    47  	// If [value] is nil, returns [ErrMarshalNil]
    48  	Size(version uint16, value interface{}) (int, error)
    49  
    50  	// Marshal the given value using the codec with the given version.
    51  	// RegisterCodec must have been called with that version.
    52  	Marshal(version uint16, source interface{}) (destination []byte, err error)
    53  
    54  	// Unmarshal the given bytes into the given destination. [destination] must
    55  	// be a pointer or an interface. Returns the version of the codec that
    56  	// produces the given bytes.
    57  	Unmarshal(source []byte, destination interface{}) (version uint16, err error)
    58  }
    59  
    60  // NewManager returns a new codec manager.
    61  func NewManager(maxSize int) Manager {
    62  	return &manager{
    63  		maxSize: maxSize,
    64  		codecs:  map[uint16]Codec{},
    65  	}
    66  }
    67  
    68  // NewDefaultManager returns a new codec manager.
    69  func NewDefaultManager() Manager {
    70  	return NewManager(defaultMaxSize)
    71  }
    72  
    73  type manager struct {
    74  	lock    sync.RWMutex
    75  	maxSize int
    76  	codecs  map[uint16]Codec
    77  }
    78  
    79  // RegisterCodec is used to register a new codec version that can be used to
    80  // (un)marshal with.
    81  func (m *manager) RegisterCodec(version uint16, codec Codec) error {
    82  	m.lock.Lock()
    83  	defer m.lock.Unlock()
    84  
    85  	if _, exists := m.codecs[version]; exists {
    86  		return ErrDuplicatedVersion
    87  	}
    88  	m.codecs[version] = codec
    89  	return nil
    90  }
    91  
    92  func (m *manager) Size(version uint16, value interface{}) (int, error) {
    93  	if value == nil {
    94  		return 0, ErrMarshalNil // can't marshal nil
    95  	}
    96  
    97  	m.lock.RLock()
    98  	c, exists := m.codecs[version]
    99  	m.lock.RUnlock()
   100  
   101  	if !exists {
   102  		return 0, ErrUnknownVersion
   103  	}
   104  
   105  	res, err := c.Size(value)
   106  
   107  	// Add [VersionSize] for the codec version
   108  	return VersionSize + res, err
   109  }
   110  
   111  // To marshal an interface, [value] must be a pointer to the interface.
   112  func (m *manager) Marshal(version uint16, value interface{}) ([]byte, error) {
   113  	if value == nil {
   114  		return nil, ErrMarshalNil // can't marshal nil
   115  	}
   116  
   117  	m.lock.RLock()
   118  	c, exists := m.codecs[version]
   119  	m.lock.RUnlock()
   120  	if !exists {
   121  		return nil, ErrUnknownVersion
   122  	}
   123  
   124  	p := wrappers.Packer{
   125  		MaxSize: m.maxSize,
   126  		Bytes:   make([]byte, 0, initialSliceCap),
   127  	}
   128  	p.PackShort(version)
   129  	if p.Errored() {
   130  		return nil, ErrCantPackVersion // Should never happen
   131  	}
   132  	return p.Bytes, c.MarshalInto(value, &p)
   133  }
   134  
   135  // Unmarshal unmarshals [bytes] into [dest], where [dest] must be a pointer or
   136  // interface.
   137  func (m *manager) Unmarshal(bytes []byte, dest interface{}) (uint16, error) {
   138  	if dest == nil {
   139  		return 0, ErrUnmarshalNil
   140  	}
   141  
   142  	if byteLen := len(bytes); byteLen > m.maxSize {
   143  		return 0, fmt.Errorf("%w: %d > %d", ErrUnmarshalTooBig, byteLen, m.maxSize)
   144  	}
   145  
   146  	p := wrappers.Packer{
   147  		Bytes: bytes,
   148  	}
   149  	version := p.UnpackShort()
   150  	if p.Errored() { // Make sure the codec version is correct
   151  		return 0, ErrCantUnpackVersion
   152  	}
   153  
   154  	m.lock.RLock()
   155  	c, exists := m.codecs[version]
   156  	m.lock.RUnlock()
   157  	if !exists {
   158  		return version, ErrUnknownVersion
   159  	}
   160  	return version, c.Unmarshal(p.Bytes[p.Offset:], dest)
   161  }