122 lines
2.6 KiB
Go
122 lines
2.6 KiB
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
|
|
}
|