media.go 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
  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. "fmt"
  7. "io"
  8. "io/ioutil"
  9. "mime/multipart"
  10. "net/http"
  11. "net/textproto"
  12. "google.golang.org/api/googleapi"
  13. )
  14. const sniffBuffSize = 512
  15. func newContentSniffer(r io.Reader) *contentSniffer {
  16. return &contentSniffer{r: r}
  17. }
  18. // contentSniffer wraps a Reader, and reports the content type determined by sniffing up to 512 bytes from the Reader.
  19. type contentSniffer struct {
  20. r io.Reader
  21. start []byte // buffer for the sniffed bytes.
  22. err error // set to any error encountered while reading bytes to be sniffed.
  23. ctype string // set on first sniff.
  24. sniffed bool // set to true on first sniff.
  25. }
  26. func (cs *contentSniffer) Read(p []byte) (n int, err error) {
  27. // Ensure that the content type is sniffed before any data is consumed from Reader.
  28. _, _ = cs.ContentType()
  29. if len(cs.start) > 0 {
  30. n := copy(p, cs.start)
  31. cs.start = cs.start[n:]
  32. return n, nil
  33. }
  34. // We may have read some bytes into start while sniffing, even if the read ended in an error.
  35. // We should first return those bytes, then the error.
  36. if cs.err != nil {
  37. return 0, cs.err
  38. }
  39. // Now we have handled all bytes that were buffered while sniffing. Now just delegate to the underlying reader.
  40. return cs.r.Read(p)
  41. }
  42. // ContentType returns the sniffed content type, and whether the content type was succesfully sniffed.
  43. func (cs *contentSniffer) ContentType() (string, bool) {
  44. if cs.sniffed {
  45. return cs.ctype, cs.ctype != ""
  46. }
  47. cs.sniffed = true
  48. // If ReadAll hits EOF, it returns err==nil.
  49. cs.start, cs.err = ioutil.ReadAll(io.LimitReader(cs.r, sniffBuffSize))
  50. // Don't try to detect the content type based on possibly incomplete data.
  51. if cs.err != nil {
  52. return "", false
  53. }
  54. cs.ctype = http.DetectContentType(cs.start)
  55. return cs.ctype, true
  56. }
  57. // DetermineContentType determines the content type of the supplied reader.
  58. // If the content type is already known, it can be specified via ctype.
  59. // Otherwise, the content of media will be sniffed to determine the content type.
  60. // If media implements googleapi.ContentTyper (deprecated), this will be used
  61. // instead of sniffing the content.
  62. // After calling DetectContentType the caller must not perform further reads on
  63. // media, but rather read from the Reader that is returned.
  64. func DetermineContentType(media io.Reader, ctype string) (io.Reader, string) {
  65. // Note: callers could avoid calling DetectContentType if ctype != "",
  66. // but doing the check inside this function reduces the amount of
  67. // generated code.
  68. if ctype != "" {
  69. return media, ctype
  70. }
  71. // For backwards compatability, allow clients to set content
  72. // type by providing a ContentTyper for media.
  73. if typer, ok := media.(googleapi.ContentTyper); ok {
  74. return media, typer.ContentType()
  75. }
  76. sniffer := newContentSniffer(media)
  77. if ctype, ok := sniffer.ContentType(); ok {
  78. return sniffer, ctype
  79. }
  80. // If content type could not be sniffed, reads from sniffer will eventually fail with an error.
  81. return sniffer, ""
  82. }
  83. type typeReader struct {
  84. io.Reader
  85. typ string
  86. }
  87. // multipartReader combines the contents of multiple readers to creat a multipart/related HTTP body.
  88. // Close must be called if reads from the multipartReader are abandoned before reaching EOF.
  89. type multipartReader struct {
  90. pr *io.PipeReader
  91. pipeOpen bool
  92. ctype string
  93. }
  94. func newMultipartReader(parts []typeReader) *multipartReader {
  95. mp := &multipartReader{pipeOpen: true}
  96. var pw *io.PipeWriter
  97. mp.pr, pw = io.Pipe()
  98. mpw := multipart.NewWriter(pw)
  99. mp.ctype = "multipart/related; boundary=" + mpw.Boundary()
  100. go func() {
  101. for _, part := range parts {
  102. w, err := mpw.CreatePart(typeHeader(part.typ))
  103. if err != nil {
  104. mpw.Close()
  105. pw.CloseWithError(fmt.Errorf("googleapi: CreatePart failed: %v", err))
  106. return
  107. }
  108. _, err = io.Copy(w, part.Reader)
  109. if err != nil {
  110. mpw.Close()
  111. pw.CloseWithError(fmt.Errorf("googleapi: Copy failed: %v", err))
  112. return
  113. }
  114. }
  115. mpw.Close()
  116. pw.Close()
  117. }()
  118. return mp
  119. }
  120. func (mp *multipartReader) Read(data []byte) (n int, err error) {
  121. return mp.pr.Read(data)
  122. }
  123. func (mp *multipartReader) Close() error {
  124. if !mp.pipeOpen {
  125. return nil
  126. }
  127. mp.pipeOpen = false
  128. return mp.pr.Close()
  129. }
  130. // CombineBodyMedia combines a json body with media content to create a multipart/related HTTP body.
  131. // It returns a ReadCloser containing the combined body, and the overall "multipart/related" content type, with random boundary.
  132. //
  133. // The caller must call Close on the returned ReadCloser if reads are abandoned before reaching EOF.
  134. func CombineBodyMedia(body io.Reader, bodyContentType string, media io.Reader, mediaContentType string) (io.ReadCloser, string) {
  135. mp := newMultipartReader([]typeReader{
  136. {body, bodyContentType},
  137. {media, mediaContentType},
  138. })
  139. return mp, mp.ctype
  140. }
  141. func typeHeader(contentType string) textproto.MIMEHeader {
  142. h := make(textproto.MIMEHeader)
  143. h.Set("Content-Type", contentType)
  144. return h
  145. }
  146. // PrepareUpload determines whether the data in the supplied reader should be
  147. // uploaded in a single request, or in sequential chunks.
  148. // chunkSize is the size of the chunk that media should be split into.
  149. // If chunkSize is non-zero and the contents of media do not fit in a single
  150. // chunk (or there is an error reading media), then media will be returned as a
  151. // ResumableBuffer. Otherwise, media will be returned as a Reader.
  152. //
  153. // After PrepareUpload has been called, media should no longer be used: the
  154. // media content should be accessed via one of the return values.
  155. func PrepareUpload(media io.Reader, chunkSize int) (io.Reader,
  156. *ResumableBuffer) {
  157. if chunkSize == 0 { // do not chunk
  158. return media, nil
  159. }
  160. rb := NewResumableBuffer(media, chunkSize)
  161. rdr, _, _, err := rb.Chunk()
  162. if err == io.EOF { // we can upload this in a single request
  163. return rdr, nil
  164. }
  165. // err might be a non-EOF error. If it is, the next call to rb.Chunk will
  166. // return the same error. Returning a ResumableBuffer ensures that this error
  167. // will be handled at some point.
  168. return nil, rb
  169. }