vendor/google.golang.org/api/gensupport/media.go

   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
   5 package gensupport
   6
   7 import (
   8         "bytes"
   9         "fmt"
  10         "io"
  11         "io/ioutil"
  12         "mime/multipart"
  13         "net/http"
  14         "net/textproto"
  15         "strings"
  16         "sync"
  17
  18         "google.golang.org/api/googleapi"
  19 )
  20
  21 const sniffBuffSize = 512
  22
  23 func newContentSniffer(r io.Reader) *contentSniffer {
  24         return &contentSniffer{r: r}
  25 }
  26
  27 // contentSniffer wraps a Reader, and reports the content type determined by sniffing up to 512 bytes from the Reader.
  28 type contentSniffer struct {
  29         r     io.Reader
  30         start []byte // buffer for the sniffed bytes.
  31         err   error  // set to any error encountered while reading bytes to be sniffed.
  32
  33         ctype   string // set on first sniff.
  34         sniffed bool   // set to true on first sniff.
  35 }
  36
  37 func (cs *contentSniffer) Read(p []byte) (n int, err error) {
  38         // Ensure that the content type is sniffed before any data is consumed from Reader.
  39         _, _ = cs.ContentType()
  40
  41         if len(cs.start) > 0 {
  42                 n := copy(p, cs.start)
  43                 cs.start = cs.start[n:]
  44                 return n, nil
  45         }
  46
  47         // We may have read some bytes into start while sniffing, even if the read ended in an error.
  48         // We should first return those bytes, then the error.
  49         if cs.err != nil {
  50                 return 0, cs.err
  51         }
  52
  53         // Now we have handled all bytes that were buffered while sniffing.  Now just delegate to the underlying reader.
  54         return cs.r.Read(p)
  55 }
  56
  57 // ContentType returns the sniffed content type, and whether the content type was succesfully sniffed.
  58 func (cs *contentSniffer) ContentType() (string, bool) {
  59         if cs.sniffed {
  60                 return cs.ctype, cs.ctype != ""
  61         }
  62         cs.sniffed = true
  63         // If ReadAll hits EOF, it returns err==nil.
  64         cs.start, cs.err = ioutil.ReadAll(io.LimitReader(cs.r, sniffBuffSize))
  65
  66         // Don't try to detect the content type based on possibly incomplete data.
  67         if cs.err != nil {
  68                 return "", false
  69         }
  70
  71         cs.ctype = http.DetectContentType(cs.start)
  72         return cs.ctype, true
  73 }
  74
  75 // DetermineContentType determines the content type of the supplied reader.
  76 // If the content type is already known, it can be specified via ctype.
  77 // Otherwise, the content of media will be sniffed to determine the content type.
  78 // If media implements googleapi.ContentTyper (deprecated), this will be used
  79 // instead of sniffing the content.
  80 // After calling DetectContentType the caller must not perform further reads on
  81 // media, but rather read from the Reader that is returned.
  82 func DetermineContentType(media io.Reader, ctype string) (io.Reader, string) {
  83         // Note: callers could avoid calling DetectContentType if ctype != "",
  84         // but doing the check inside this function reduces the amount of
  85         // generated code.
  86         if ctype != "" {
  87                 return media, ctype
  88         }
  89
  90         // For backwards compatability, allow clients to set content
  91         // type by providing a ContentTyper for media.
  92         if typer, ok := media.(googleapi.ContentTyper); ok {
  93                 return media, typer.ContentType()
  94         }
  95
  96         sniffer := newContentSniffer(media)
  97         if ctype, ok := sniffer.ContentType(); ok {
  98                 return sniffer, ctype
  99         }
 100         // If content type could not be sniffed, reads from sniffer will eventually fail with an error.
 101         return sniffer, ""
 102 }
 103
 104 type typeReader struct {
 105         io.Reader
 106         typ string
 107 }
 108
 109 // multipartReader combines the contents of multiple readers to create a multipart/related HTTP body.
 110 // Close must be called if reads from the multipartReader are abandoned before reaching EOF.
 111 type multipartReader struct {
 112         pr       *io.PipeReader
 113         ctype    string
 114         mu       sync.Mutex
 115         pipeOpen bool
 116 }
 117
 118 func newMultipartReader(parts []typeReader) *multipartReader {
 119         mp := &multipartReader{pipeOpen: true}
 120         var pw *io.PipeWriter
 121         mp.pr, pw = io.Pipe()
 122         mpw := multipart.NewWriter(pw)
 123         mp.ctype = "multipart/related; boundary=" + mpw.Boundary()
 124         go func() {
 125                 for _, part := range parts {
 126                         w, err := mpw.CreatePart(typeHeader(part.typ))
 127                         if err != nil {
 128                                 mpw.Close()
 129                                 pw.CloseWithError(fmt.Errorf("googleapi: CreatePart failed: %v", err))
 130                                 return
 131                         }
 132                         _, err = io.Copy(w, part.Reader)
 133                         if err != nil {
 134                                 mpw.Close()
 135                                 pw.CloseWithError(fmt.Errorf("googleapi: Copy failed: %v", err))
 136                                 return
 137                         }
 138                 }
 139
 140                 mpw.Close()
 141                 pw.Close()
 142         }()
 143         return mp
 144 }
 145
 146 func (mp *multipartReader) Read(data []byte) (n int, err error) {
 147         return mp.pr.Read(data)
 148 }
 149
 150 func (mp *multipartReader) Close() error {
 151         mp.mu.Lock()
 152         if !mp.pipeOpen {
 153                 mp.mu.Unlock()
 154                 return nil
 155         }
 156         mp.pipeOpen = false
 157         mp.mu.Unlock()
 158         return mp.pr.Close()
 159 }
 160
 161 // CombineBodyMedia combines a json body with media content to create a multipart/related HTTP body.
 162 // It returns a ReadCloser containing the combined body, and the overall "multipart/related" content type, with random boundary.
 163 //
 164 // The caller must call Close on the returned ReadCloser if reads are abandoned before reaching EOF.
 165 func CombineBodyMedia(body io.Reader, bodyContentType string, media io.Reader, mediaContentType string) (io.ReadCloser, string) {
 166         mp := newMultipartReader([]typeReader{
 167                 {body, bodyContentType},
 168                 {media, mediaContentType},
 169         })
 170         return mp, mp.ctype
 171 }
 172
 173 func typeHeader(contentType string) textproto.MIMEHeader {
 174         h := make(textproto.MIMEHeader)
 175         if contentType != "" {
 176                 h.Set("Content-Type", contentType)
 177         }
 178         return h
 179 }
 180
 181 // PrepareUpload determines whether the data in the supplied reader should be
 182 // uploaded in a single request, or in sequential chunks.
 183 // chunkSize is the size of the chunk that media should be split into.
 184 //
 185 // If chunkSize is zero, media is returned as the first value, and the other
 186 // two return values are nil, true.
 187 //
 188 // Otherwise, a MediaBuffer is returned, along with a bool indicating whether the
 189 // contents of media fit in a single chunk.
 190 //
 191 // After PrepareUpload has been called, media should no longer be used: the
 192 // media content should be accessed via one of the return values.
 193 func PrepareUpload(media io.Reader, chunkSize int) (r io.Reader, mb *MediaBuffer, singleChunk bool) {
 194         if chunkSize == 0 { // do not chunk
 195                 return media, nil, true
 196         }
 197         mb = NewMediaBuffer(media, chunkSize)
 198         _, _, _, err := mb.Chunk()
 199         // If err is io.EOF, we can upload this in a single request. Otherwise, err is
 200         // either nil or a non-EOF error. If it is the latter, then the next call to
 201         // mb.Chunk will return the same error. Returning a MediaBuffer ensures that this
 202         // error will be handled at some point.
 203         return nil, mb, err == io.EOF
 204 }
 205
 206 // MediaInfo holds information for media uploads. It is intended for use by generated
 207 // code only.
 208 type MediaInfo struct {
 209         // At most one of Media and MediaBuffer will be set.
 210         media           io.Reader
 211         buffer          *MediaBuffer
 212         singleChunk     bool
 213         mType           string
 214         size            int64 // mediaSize, if known.  Used only for calls to progressUpdater_.
 215         progressUpdater googleapi.ProgressUpdater
 216 }
 217
 218 // NewInfoFromMedia should be invoked from the Media method of a call. It returns a
 219 // MediaInfo populated with chunk size and content type, and a reader or MediaBuffer
 220 // if needed.
 221 func NewInfoFromMedia(r io.Reader, options []googleapi.MediaOption) *MediaInfo {
 222         mi := &MediaInfo{}
 223         opts := googleapi.ProcessMediaOptions(options)
 224         if !opts.ForceEmptyContentType {
 225                 r, mi.mType = DetermineContentType(r, opts.ContentType)
 226         }
 227         mi.media, mi.buffer, mi.singleChunk = PrepareUpload(r, opts.ChunkSize)
 228         return mi
 229 }
 230
 231 // NewInfoFromResumableMedia should be invoked from the ResumableMedia method of a
 232 // call. It returns a MediaInfo using the given reader, size and media type.
 233 func NewInfoFromResumableMedia(r io.ReaderAt, size int64, mediaType string) *MediaInfo {
 234         rdr := ReaderAtToReader(r, size)
 235         rdr, mType := DetermineContentType(rdr, mediaType)
 236         return &MediaInfo{
 237                 size:        size,
 238                 mType:       mType,
 239                 buffer:      NewMediaBuffer(rdr, googleapi.DefaultUploadChunkSize),
 240                 media:       nil,
 241                 singleChunk: false,
 242         }
 243 }
 244
 245 // SetProgressUpdater sets the progress updater for the media info.
 246 func (mi *MediaInfo) SetProgressUpdater(pu googleapi.ProgressUpdater) {
 247         if mi != nil {
 248                 mi.progressUpdater = pu
 249         }
 250 }
 251
 252 // UploadType determines the type of upload: a single request, or a resumable
 253 // series of requests.
 254 func (mi *MediaInfo) UploadType() string {
 255         if mi.singleChunk {
 256                 return "multipart"
 257         }
 258         return "resumable"
 259 }
 260
 261 // UploadRequest sets up an HTTP request for media upload. It adds headers
 262 // as necessary, and returns a replacement for the body and a function for http.Request.GetBody.
 263 func (mi *MediaInfo) UploadRequest(reqHeaders http.Header, body io.Reader) (newBody io.Reader, getBody func() (io.ReadCloser, error), cleanup func()) {
 264         cleanup = func() {}
 265         if mi == nil {
 266                 return body, nil, cleanup
 267         }
 268         var media io.Reader
 269         if mi.media != nil {
 270                 // This only happens when the caller has turned off chunking. In that
 271                 // case, we write all of media in a single non-retryable request.
 272                 media = mi.media
 273         } else if mi.singleChunk {
 274                 // The data fits in a single chunk, which has now been read into the MediaBuffer.
 275                 // We obtain that chunk so we can write it in a single request. The request can
 276                 // be retried because the data is stored in the MediaBuffer.
 277                 media, _, _, _ = mi.buffer.Chunk()
 278         }
 279         if media != nil {
 280                 fb := readerFunc(body)
 281                 fm := readerFunc(media)
 282                 combined, ctype := CombineBodyMedia(body, "application/json", media, mi.mType)
 283                 if fb != nil && fm != nil {
 284                         getBody = func() (io.ReadCloser, error) {
 285                                 rb := ioutil.NopCloser(fb())
 286                                 rm := ioutil.NopCloser(fm())
 287                                 r, _ := CombineBodyMedia(rb, "application/json", rm, mi.mType)
 288                                 return r, nil
 289                         }
 290                 }
 291                 cleanup = func() { combined.Close() }
 292                 reqHeaders.Set("Content-Type", ctype)
 293                 body = combined
 294         }
 295         if mi.buffer != nil && mi.mType != "" && !mi.singleChunk {
 296                 reqHeaders.Set("X-Upload-Content-Type", mi.mType)
 297         }
 298         return body, getBody, cleanup
 299 }
 300
 301 // readerFunc returns a function that always returns an io.Reader that has the same
 302 // contents as r, provided that can be done without consuming r. Otherwise, it
 303 // returns nil.
 304 // See http.NewRequest (in net/http/request.go).
 305 func readerFunc(r io.Reader) func() io.Reader {
 306         switch r := r.(type) {
 307         case *bytes.Buffer:
 308                 buf := r.Bytes()
 309                 return func() io.Reader { return bytes.NewReader(buf) }
 310         case *bytes.Reader:
 311                 snapshot := *r
 312                 return func() io.Reader { r := snapshot; return &r }
 313         case *strings.Reader:
 314                 snapshot := *r
 315                 return func() io.Reader { r := snapshot; return &r }
 316         default:
 317                 return nil
 318         }
 319 }
 320
 321 // ResumableUpload returns an appropriately configured ResumableUpload value if the
 322 // upload is resumable, or nil otherwise.
 323 func (mi *MediaInfo) ResumableUpload(locURI string) *ResumableUpload {
 324         if mi == nil || mi.singleChunk {
 325                 return nil
 326         }
 327         return &ResumableUpload{
 328                 URI:       locURI,
 329                 Media:     mi.buffer,
 330                 MediaType: mi.mType,
 331                 Callback: func(curr int64) {
 332                         if mi.progressUpdater != nil {
 333                                 mi.progressUpdater(curr, mi.size)
 334                         }
 335                 },
 336         }
 337 }
 338
 339 // SetGetBody sets the GetBody field of req to f.
 340 func SetGetBody(req *http.Request, f func() (io.ReadCloser, error)) {
 341         req.GetBody = f
 342 }