1// Copyright 2016 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package gensupport
6
7import (
8	"bytes"
9	"io"
10
11	"google.golang.org/api/googleapi"
12)
13
14// MediaBuffer buffers data from an io.Reader to support uploading media in retryable chunks.
15type MediaBuffer struct {
16	media io.Reader
17
18	chunk []byte // The current chunk which is pending upload.  The capacity is the chunk size.
19	err   error  // Any error generated when populating chunk by reading media.
20
21	// The absolute position of chunk in the underlying media.
22	off int64
23}
24
25func NewMediaBuffer(media io.Reader, chunkSize int) *MediaBuffer {
26	return &MediaBuffer{media: media, chunk: make([]byte, 0, chunkSize)}
27}
28
29// Chunk returns the current buffered chunk, the offset in the underlying media
30// from which the chunk is drawn, and the size of the chunk.
31// Successive calls to Chunk return the same chunk between calls to Next.
32func (mb *MediaBuffer) Chunk() (chunk io.Reader, off int64, size int, err error) {
33	// There may already be data in chunk if Next has not been called since the previous call to Chunk.
34	if mb.err == nil && len(mb.chunk) == 0 {
35		mb.err = mb.loadChunk()
36	}
37	return bytes.NewReader(mb.chunk), mb.off, len(mb.chunk), mb.err
38}
39
40// loadChunk will read from media into chunk, up to the capacity of chunk.
41func (mb *MediaBuffer) loadChunk() error {
42	bufSize := cap(mb.chunk)
43	mb.chunk = mb.chunk[:bufSize]
44
45	read := 0
46	var err error
47	for err == nil && read < bufSize {
48		var n int
49		n, err = mb.media.Read(mb.chunk[read:])
50		read += n
51	}
52	mb.chunk = mb.chunk[:read]
53	return err
54}
55
56// Next advances to the next chunk, which will be returned by the next call to Chunk.
57// Calls to Next without a corresponding prior call to Chunk will have no effect.
58func (mb *MediaBuffer) Next() {
59	mb.off += int64(len(mb.chunk))
60	mb.chunk = mb.chunk[0:0]
61}
62
63type readerTyper struct {
64	io.Reader
65	googleapi.ContentTyper
66}
67
68// ReaderAtToReader adapts a ReaderAt to be used as a Reader.
69// If ra implements googleapi.ContentTyper, then the returned reader
70// will also implement googleapi.ContentTyper, delegating to ra.
71func ReaderAtToReader(ra io.ReaderAt, size int64) io.Reader {
72	r := io.NewSectionReader(ra, 0, size)
73	if typer, ok := ra.(googleapi.ContentTyper); ok {
74		return readerTyper{r, typer}
75	}
76	return r
77}
78