media.go 12 KB

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