github.com/vipernet-xyz/tm@v0.34.24/test/loadtime/payload/payload.go (about)

     1  package payload
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/rand"
     6  	"encoding/hex"
     7  	"fmt"
     8  	"math"
     9  
    10  	"google.golang.org/protobuf/proto"
    11  	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
    12  )
    13  
    14  const keyPrefix = "a="
    15  const maxPayloadSize = 4 * 1024 * 1024
    16  
    17  // NewBytes generates a new payload and returns the encoded representation of
    18  // the payload as a slice of bytes. NewBytes uses the fields on the Options
    19  // to create the payload.
    20  func NewBytes(p *Payload) ([]byte, error) {
    21  	p.Padding = make([]byte, 1)
    22  	if p.Time == nil {
    23  		p.Time = timestamppb.Now()
    24  	}
    25  	us, err := CalculateUnpaddedSize(p)
    26  	if err != nil {
    27  		return nil, err
    28  	}
    29  	if p.Size > maxPayloadSize {
    30  		return nil, fmt.Errorf("configured size %d is too large (>%d)", p.Size, maxPayloadSize)
    31  	}
    32  	pSize := int(p.Size) // #nosec -- The "if" above makes this cast safe
    33  	if pSize < us {
    34  		return nil, fmt.Errorf("configured size %d not large enough to fit unpadded transaction of size %d", pSize, us)
    35  	}
    36  
    37  	// We halve the padding size because we transform the TX to hex
    38  	p.Padding = make([]byte, (pSize-us)/2)
    39  	_, err = rand.Read(p.Padding)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  	b, err := proto.Marshal(p)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  	h := []byte(hex.EncodeToString(b))
    48  
    49  	// prepend a single key so that the kv store only ever stores a single
    50  	// transaction instead of storing all tx and ballooning in size.
    51  	return append([]byte(keyPrefix), h...), nil
    52  }
    53  
    54  // FromBytes extracts a paylod from the byte representation of the payload.
    55  // FromBytes leaves the padding untouched, returning it to the caller to handle
    56  // or discard per their preference.
    57  func FromBytes(b []byte) (*Payload, error) {
    58  	trH := bytes.TrimPrefix(b, []byte(keyPrefix))
    59  	if bytes.Equal(b, trH) {
    60  		return nil, fmt.Errorf("payload bytes missing key prefix '%s'", keyPrefix)
    61  	}
    62  	trB, err := hex.DecodeString(string(trH))
    63  	if err != nil {
    64  		return nil, err
    65  	}
    66  
    67  	p := &Payload{}
    68  	err = proto.Unmarshal(trB, p)
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  	return p, nil
    73  }
    74  
    75  // MaxUnpaddedSize returns the maximum size that a payload may be if no padding
    76  // is included.
    77  func MaxUnpaddedSize() (int, error) {
    78  	p := &Payload{
    79  		Time:        timestamppb.Now(),
    80  		Connections: math.MaxUint64,
    81  		Rate:        math.MaxUint64,
    82  		Size:        math.MaxUint64,
    83  		Padding:     make([]byte, 1),
    84  	}
    85  	return CalculateUnpaddedSize(p)
    86  }
    87  
    88  // CalculateUnpaddedSize calculates the size of the passed in payload for the
    89  // purpose of determining how much padding to add to add to reach the target size.
    90  // CalculateUnpaddedSize returns an error if the payload Padding field is longer than 1.
    91  func CalculateUnpaddedSize(p *Payload) (int, error) {
    92  	if len(p.Padding) != 1 {
    93  		return 0, fmt.Errorf("expected length of padding to be 1, received %d", len(p.Padding))
    94  	}
    95  	b, err := proto.Marshal(p)
    96  	if err != nil {
    97  		return 0, err
    98  	}
    99  	h := []byte(hex.EncodeToString(b))
   100  	return len(h) + len(keyPrefix), nil
   101  }