storj/pkg/eestream/transform.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()
}