storj/pkg/eestream/transform.go
2018-08-06 17:24:30 +03:00

166 lines
5.0 KiB
Go

// Copyright (C) 2018 Storj Labs, Inc.
// See LICENSE for copying information.
package eestream
import (
"bytes"
"context"
"io"
"io/ioutil"
"storj.io/storj/internal/pkg/readcloser"
"storj.io/storj/pkg/ranger"
)
// A Transformer is a data transformation that may change the size of the blocks
// of data it operates on in a deterministic fashion.
type Transformer interface {
InBlockSize() int // The block size prior to transformation
OutBlockSize() int // The block size after transformation
Transform(out, in []byte, blockNum int64) ([]byte, error)
}
type transformedReader struct {
r io.ReadCloser
t Transformer
blockNum int64
inbuf []byte
outbuf []byte
expectedSize int64
bytesRead int
}
// TransformReader applies a Transformer to a Reader. startingBlockNum should
// probably be 0 unless you know you're already starting at a block offset.
func TransformReader(r io.ReadCloser, t Transformer,
startingBlockNum int64) io.ReadCloser {
return &transformedReader{
r: r,
t: t,
blockNum: startingBlockNum,
inbuf: make([]byte, t.InBlockSize()),
outbuf: make([]byte, 0, t.OutBlockSize()),
}
}
// TransformReaderSize creates a TransformReader with expected size,
// i.e. the number of bytes that is expected to be read from this reader.
// If less than the expected bytes are read, the reader will return
// io.ErrUnexpectedEOF instead of io.EOF.
func TransformReaderSize(r io.ReadCloser, t Transformer,
startingBlockNum int64, expectedSize int64) io.ReadCloser {
return &transformedReader{
r: r,
t: t,
blockNum: startingBlockNum,
inbuf: make([]byte, t.InBlockSize()),
outbuf: make([]byte, 0, t.OutBlockSize()),
expectedSize: expectedSize,
}
}
func (t *transformedReader) Read(p []byte) (n int, err error) {
if len(t.outbuf) <= 0 {
// If there's no more buffered data left, let's fill the buffer with
// the next block
b, err := io.ReadFull(t.r, t.inbuf)
t.bytesRead += b
if err == io.EOF && int64(t.bytesRead) < t.expectedSize {
return 0, io.ErrUnexpectedEOF
} else if err != nil {
return 0, err
}
t.outbuf, err = t.t.Transform(t.outbuf, t.inbuf, t.blockNum)
if err != nil {
return 0, Error.Wrap(err)
}
t.blockNum++
}
// return as much as we can from the current buffered block
n = copy(p, t.outbuf)
// slide the uncopied data to the beginning of the buffer
copy(t.outbuf, t.outbuf[n:])
// resize the buffer
t.outbuf = t.outbuf[:len(t.outbuf)-n]
return n, nil
}
func (t *transformedReader) Close() error {
return t.r.Close()
}
type transformedRanger struct {
rr ranger.RangeCloser
t Transformer
}
// Transform will apply a Transformer to a Ranger.
func Transform(rr ranger.RangeCloser, t Transformer) (ranger.RangeCloser, error) {
if rr.Size()%int64(t.InBlockSize()) != 0 {
return nil, Error.New("invalid transformer and range reader combination." +
"the range reader size is not a multiple of the block size")
}
return &transformedRanger{rr: rr, t: t}, nil
}
func (t *transformedRanger) Size() int64 {
blocks := t.rr.Size() / int64(t.t.InBlockSize())
return blocks * int64(t.t.OutBlockSize())
}
// calcEncompassingBlocks is a useful helper function that, given an offset,
// length, and blockSize, will tell you which blocks contain the requested
// offset and length
func calcEncompassingBlocks(offset, length int64, blockSize int) (
firstBlock, blockCount int64) {
firstBlock = offset / int64(blockSize)
if length <= 0 {
return firstBlock, 0
}
lastBlock := (offset + length) / int64(blockSize)
if (offset+length)%int64(blockSize) == 0 {
return firstBlock, lastBlock - firstBlock
}
return firstBlock, 1 + lastBlock - firstBlock
}
func (t *transformedRanger) Range(ctx context.Context, offset, length int64) (io.ReadCloser, error) {
// Range may not have been called for block-aligned offsets and lengths, so
// let's figure out which blocks encompass the request
firstBlock, blockCount := calcEncompassingBlocks(
offset, length, t.t.OutBlockSize())
// If block count is 0, there is nothing to transform, so return a dumb
// reader that will just return io.EOF on read
if blockCount == 0 {
return ioutil.NopCloser(bytes.NewReader(nil)), nil
}
// okay, now let's get the range on the underlying ranger for those blocks
// and then Transform it.
r, err := t.rr.Range(ctx,
firstBlock*int64(t.t.InBlockSize()),
blockCount*int64(t.t.InBlockSize()))
if err != nil {
return nil, err
}
tr := TransformReaderSize(r, t.t, firstBlock, blockCount*int64(t.t.InBlockSize()))
// the range we got potentially includes more than we wanted. if the
// offset started past the beginning of the first block, we need to
// swallow the first few bytes
_, err = io.CopyN(ioutil.Discard, tr,
offset-firstBlock*int64(t.t.OutBlockSize()))
if err != nil {
if err == io.EOF {
return nil, io.ErrUnexpectedEOF
}
return nil, Error.Wrap(err)
}
// the range might have been too long. only return what was requested
return readcloser.LimitReadCloser(tr, length), nil
}
func (t *transformedRanger) Close() error {
return t.rr.Close()
}