xk6-frostfs/internal/datagen/payload.go

package datagen

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"hash"
	"io"
	"math/rand"

	"github.com/go-loremipsum/loremipsum"
)

// Payload represents arbitrary data to be packed into S3 or native object.
// Implementations could be thread-unsafe.
type Payload interface {
	// Reader returns io.Reader instance to read the payload.
	// Must not be called twice.
	Reader() io.Reader
	// Bytes is a helper which reads all data from Reader() into slice.
	// The sole purpose of this method is to simplify HTTP scenario,
	// where all payload needs to be read and wrapped.
	Bytes() []byte
	// Size returns payload size, which is equal to the total amount of data
	// that could be read from the Reader().
	Size() int
	// Hash returns payload sha256 hash. Must be called after all data is read from the reader.
	Hash() string
}

type bytesPayload struct {
	data []byte
}

func (p *bytesPayload) Reader() io.Reader {
	return bytes.NewReader(p.data)
}

func (p *bytesPayload) Size() int {
	return len(p.data)
}

func (p *bytesPayload) Hash() string {
	h := sha256.Sum256(p.data[:])
	return hex.EncodeToString(h[:])
}

func (p *bytesPayload) Bytes() []byte {
	return p.data
}

func NewFixedPayload(data []byte) Payload {
	return &bytesPayload{data: data}
}

type randomPayload struct {
	r    io.Reader
	s    hash.Hash
	h    string
	size int
}

func NewStreamPayload(size int, seed int64, typ string) Payload {
	var rr io.Reader
	switch typ {
	case "text":
		rr = &textReader{li: loremipsum.NewWithSeed(seed)}
	default:
		rr = rand.New(rand.NewSource(seed))
	}

	lr := io.LimitReader(rr, int64(size))
	// We need some buffering to write complete blocks in the TeeReader.
	// Streaming payload read is expected to be used for big objects, thus 4k seems like a good choice.
	br := bufio.NewReaderSize(lr, 4096)
	s := sha256.New()
	tr := io.TeeReader(br, s)
	return &randomPayload{
		r:    tr,
		s:    s,
		size: size,
	}
}

func (p *randomPayload) Reader() io.Reader {
	return p.r
}

func (p *randomPayload) Size() int {
	return p.size
}

func (p *randomPayload) Hash() string {
	if p.h == "" {
		p.h = hex.EncodeToString(p.s.Sum(nil))
		// Prevent possible misuse.
		p.r = nil
		p.s = nil
	}
	return p.h
}

func (p *randomPayload) Bytes() []byte {
	data, err := io.ReadAll(p.r)
	if err != nil {
		// We use only 2 readers, either `bytes.Reader` or `rand.Reader`.
		// None of them returns errors, thus encountering an error is a fatal error.
		panic(err)
	}
	return data
}

type textReader struct {
	li *loremipsum.LoremIpsum
}

func (r *textReader) Read(p []byte) (n int, err error) {
	paragraph := r.li.Paragraph()
	return copy(p, paragraph), nil
}