"context"

"errors"

"fmt"

"hash"

"hash/crc32"

"io"

"io/fs"

"math"

"google.golang.org/api/iterator"

)

// maxChecksumZeroArraySize is the maximum amount of memory to allocate for

// updating the checksum. A larger size will occupy more memory but will require

// fewer updates when computing the crc32c of a full object.

// TODO: test the performance of smaller values for this.

const maxChecksumZeroArraySize = 4 * 1024 * 1024

// Downloader manages a set of parallelized downloads.

type Downloader struct {

client *storage.Client

}

func (d *Downloader) addResult(input *DownloadObjectInput, result *DownloadOutput) {

copiedResult := *result // make a copy so that callbacks do not affect the result

copiedResult := *result // make a copy so that callbacks do not affect the result

if input.directory {

f := input.Destination.(*os.File)

input.directoryObjectOutputs <- copiedResult

}

}

// TODO: check checksum if full object

if d.config.asynchronous || input.directory {

input.Callback(result)

break // no more work; exit

}

out := input.downloadShard(d.client, d.config.perOperationTimeout, d.config.partSize)

if input.shard == 0 {

if out.Err != nil {

// Don't queue more shards if the first failed.

d.addResult(input, out)

} else {

numShards := numShards(out.Attrs, input.Range, d.config.partSize)

if numShards <= 1 {

// Download completed with a single shard.

d.addResult(input, out)

} else {

// Queue more shards.

outs := d.queueShards(input, out.Attrs.Generation, numShards)

// Start a goroutine that gathers shards sent to the output

// channel and adds the result once it has received all shards.

go d.gatherShards(input, outs, numShards)

}

}

d.startDownload(input)

} else {

out := input.downloadShard(d.client, d.config.perOperationTimeout, d.config.partSize)

// If this isn't the first shard, send to the output channel specific to the object.

// This should never block since the channel is buffered to exactly the number of shards.

input.shardOutputs <- out

d.workers.Done()

}

// startDownload downloads the first shard and schedules subsequent shards

// if necessary.

func (d *Downloader) startDownload(input *DownloadObjectInput) {

var out *DownloadOutput

// Full object read. Request the full object and only read partSize bytes

// (or the full object, if smaller than partSize), so that we can avoid a

// metadata call to grab the CRC32C for JSON downloads.

if fullObjectRead(input.Range) {

input.checkCRC = true

out = input.downloadFirstShard(d.client, d.config.perOperationTimeout, d.config.partSize)

} else {

out = input.downloadShard(d.client, d.config.perOperationTimeout, d.config.partSize)

}

if out.Err != nil {

// Don't queue more shards if the first failed.

d.addResult(input, out)

return

}

numShards := numShards(out.Attrs, input.Range, d.config.partSize)

input.checkCRC = input.checkCRC && !out.Attrs.Decompressed // do not checksum if the object was decompressed

if numShards > 1 {

outs := d.queueShards(input, out.Attrs.Generation, numShards)

// Start a goroutine that gathers shards sent to the output

// channel and adds the result once it has received all shards.

go d.gatherShards(input, out, outs, numShards, out.crc32c)

} else {

// Download completed with a single shard.

if input.checkCRC {

if err := checksumObject(out.crc32c, out.Attrs.CRC32C); err != nil {

out.Err = err

}

}

d.addResult(input, out)

}

}

// queueShards queues all subsequent shards of an object after the first.

// The results should be forwarded to the returned channel.

func (d *Downloader) queueShards(in *DownloadObjectInput, gen int64, shards int) <-chan *DownloadOutput {

// It will add the result to the Downloader once it has received all shards.

// gatherShards cancels remaining shards if any shard errored.

// It does not do any checking to verify that shards are for the same object.

func (d *Downloader) gatherShards(in *DownloadObjectInput, outs <-chan *DownloadOutput, shards int) {

func (d *Downloader) gatherShards(in *DownloadObjectInput, out *DownloadOutput, outs <-chan *DownloadOutput, shards int, firstPieceCRC uint32) {

errs := []error{}

var shardOut *DownloadOutput

orderedChecksums := make([]crc32cPiece, shards-1)

for i := 1; i < shards; i++ {

// Add monitoring here? This could hang if any individual piece does.

shardOut = <-outs

shardOut := <-outs

// We can ignore errors that resulted from a previous error.

// Note that we may still get some cancel errors if they

errs = append(errs, shardOut.Err)

in.cancelCtx(errCancelAllShards)

}

orderedChecksums[shardOut.shard-1] = crc32cPiece{sum: shardOut.crc32c, length: shardOut.shardLength}

}

// All pieces gathered.

if len(errs) == 0 && in.checkCRC && out.Attrs != nil {

fullCrc := joinCRC32C(firstPieceCRC, orderedChecksums)

if err := checksumObject(fullCrc, out.Attrs.CRC32C); err != nil {

errs = append(errs, err)

}

}

// All pieces gathered; return output. Any shard output will do.

shardOut.Range = in.Range

// Prepare output.

out.Range = in.Range

if len(errs) != 0 {

shardOut.Err = fmt.Errorf("download shard errors:\n%w", errors.Join(errs...))

out.Err = fmt.Errorf("download shard errors:\n%w", errors.Join(errs...))

}

if shardOut.Attrs != nil {

shardOut.Attrs.StartOffset = 0

if out.Attrs != nil {

out.Attrs.StartOffset = 0

if in.Range != nil {

shardOut.Attrs.StartOffset = in.Range.Offset

out.Attrs.StartOffset = in.Range.Offset

}

}

d.addResult(in, shardOut)

d.addResult(in, out)

}

// gatherObjectOutputs receives from the given channel exactly numObjects times.

shardOutputs chan<- *DownloadOutput

directory bool // input was queued by calling DownloadDirectory

directoryObjectOutputs chan<- DownloadOutput

checkCRC bool

}

// downloadShard will read a specific object piece into in.Destination.

// If timeout is less than 0, no timeout is set.

func (in *DownloadObjectInput) downloadShard(client *storage.Client, timeout time.Duration, partSize int64) (out *DownloadOutput) {

out = &DownloadOutput{Bucket: in.Bucket, Object: in.Object, Range: in.Range}

// Set timeout.

ctx := in.ctx

if timeout > 0 {

c, cancel := context.WithTimeout(ctx, timeout)

defer cancel()

ctx = c

}

// The first shard will be sent as download many, since we do not know yet

// if it will be sharded.

method := downloadMany

if in.shard != 0 {

method = downloadSharded

}

ctx = setUsageMetricHeader(ctx, method)

// Set options on the object.

o := client.Bucket(in.Bucket).Object(in.Object)

if in.Conditions != nil {

o = o.If(*in.Conditions)

}

if in.Generation != nil {

o = o.Generation(*in.Generation)

}

if len(in.EncryptionKey) > 0 {

o = o.Key(in.EncryptionKey)

}

objRange := shardRange(in.Range, partSize, in.shard)

ctx := in.setOptionsOnContext(timeout)

o := in.setOptionsOnObject(client)

// Read.

r, err := o.NewRangeReader(ctx, objRange.Offset, objRange.Length)

if err != nil {

out.Err = err

}

}

w := io.NewOffsetWriter(in.Destination, offset)

_, err = io.Copy(w, r)

var w io.Writer

w = io.NewOffsetWriter(in.Destination, offset)

var crcHash hash.Hash32

if in.checkCRC {

crcHash = crc32.New(crc32.MakeTable(crc32.Castagnoli))

w = io.MultiWriter(w, crcHash)

}

n, err := io.Copy(w, r)

if err != nil {

out.Err = err

r.Close()

return

}

if err = r.Close(); err != nil {

out.Err = err

return

}

out.Attrs = &r.Attrs

out.shard = in.shard

out.shardLength = n

if in.checkCRC {

out.crc32c = crcHash.Sum32()

}

return

}

// downloadFirstShard will read the first object piece into in.Destination.

// If timeout is less than 0, no timeout is set.

func (in *DownloadObjectInput) downloadFirstShard(client *storage.Client, timeout time.Duration, partSize int64) (out *DownloadOutput) {

out = &DownloadOutput{Bucket: in.Bucket, Object: in.Object, Range: in.Range}

ctx := in.setOptionsOnContext(timeout)

o := in.setOptionsOnObject(client)

r, err := o.NewReader(ctx)

if err != nil {

out.Err = err

return

}

var w io.Writer

w = io.NewOffsetWriter(in.Destination, 0)

var crcHash hash.Hash32

if in.checkCRC {

crcHash = crc32.New(crc32.MakeTable(crc32.Castagnoli))

w = io.MultiWriter(w, crcHash)

}

// Copy only the first partSize bytes before closing the reader.

// If we encounter an EOF, the file was smaller than partSize.

n, err := io.CopyN(w, r, partSize)

if err != nil && err != io.EOF {

out.Err = err

r.Close()

return

}

out.Attrs = &r.Attrs

out.shard = in.shard

out.shardLength = n

if in.checkCRC {

out.crc32c = crcHash.Sum32()

}

return

}

func (in *DownloadObjectInput) setOptionsOnContext(timeout time.Duration) context.Context {

ctx := in.ctx

if timeout > 0 {

c, cancel := context.WithTimeout(ctx, timeout)

defer cancel()

ctx = c

}

// The first shard will be sent as download many, since we do not know yet

// if it will be sharded.

method := downloadMany

if in.shard != 0 {

method = downloadSharded

}

return setUsageMetricHeader(ctx, method)

}

func (in *DownloadObjectInput) setOptionsOnObject(client *storage.Client) *storage.ObjectHandle {

o := client.Bucket(in.Bucket).Object(in.Object)

if in.Conditions != nil {

o = o.If(*in.Conditions)

}

if in.Generation != nil {

o = o.Generation(*in.Generation)

}

if len(in.EncryptionKey) > 0 {

o = o.Key(in.EncryptionKey)

}

return o

}

// DownloadDirectoryInput is the input for a directory to download.

type DownloadDirectoryInput struct {

// Bucket is the bucket in GCS to download from. Required.

Range *DownloadRange // requested range, if it was specified

Err error // error occurring during download

Attrs *storage.ReaderObjectAttrs // attributes of downloaded object, if successful

shard int

shardLength int64

crc32c uint32

}

// TODO: use built-in after go < 1.21 is dropped.

header := fmt.Sprintf("%s/%s", usageMetricKey, method)

return callctx.SetHeaders(ctx, xGoogHeaderKey, header)

}

type crc32cPiece struct {

sum uint32 // crc32c checksum of the piece

length int64 // number of bytes in this piece

}

// joinCRC32C pieces together the initial checksum with the orderedChecksums

// provided to calculate the checksum of the whole.

func joinCRC32C(initialChecksum uint32, orderedChecksums []crc32cPiece) uint32 {

base := initialChecksum

zeroes := make([]byte, maxChecksumZeroArraySize)

for _, part := range orderedChecksums {

// Precondition Base (flip every bit)

base ^= 0xFFFFFFFF

// Zero pad base crc32c. To conserve memory, do so with only maxChecksumZeroArraySize

// at a time. Reuse the zeroes array where possible.

var padded int64 = 0

for padded < part.length {

desiredZeroes := min(part.length-padded, maxChecksumZeroArraySize)

base = crc32.Update(base, crc32.MakeTable(crc32.Castagnoli), zeroes[:desiredZeroes])

padded += desiredZeroes

}

// Postcondition Base (same as precondition, this switches the bits back)

base ^= 0xFFFFFFFF

// Bitwise OR between Base and Part to produce a new Base

base ^= part.sum

}

return base

}

func fullObjectRead(r *DownloadRange) bool {

return r == nil || (r.Offset == 0 && r.Length < 0)

}

func checksumObject(got, want uint32) error {

// Only checksum the object if we have a valid CRC32C.

if want != 0 && want != got {

return fmt.Errorf("bad CRC on read: got %d, want %d", got, want)

}

return nil

}