You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

152 lines
4.5 KiB

6 years ago
12 years ago
  1. package operation
  2. import (
  3. "bytes"
  4. "compress/flate"
  5. "compress/gzip"
  6. "encoding/json"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "io/ioutil"
  11. "mime"
  12. "mime/multipart"
  13. "net/http"
  14. "net/textproto"
  15. "path/filepath"
  16. "strings"
  17. "github.com/joeslay/seaweedfs/weed/glog"
  18. "github.com/joeslay/seaweedfs/weed/security"
  19. "github.com/joeslay/seaweedfs/weed/util"
  20. )
  21. type UploadResult struct {
  22. Name string `json:"name,omitempty"`
  23. Size uint32 `json:"size,omitempty"`
  24. Error string `json:"error,omitempty"`
  25. ETag string `json:"eTag,omitempty"`
  26. }
  27. var (
  28. client *http.Client
  29. )
  30. func init() {
  31. client = &http.Client{Transport: &http.Transport{
  32. MaxIdleConnsPerHost: 1024,
  33. }}
  34. }
  35. var fileNameEscaper = strings.NewReplacer("\\", "\\\\", "\"", "\\\"")
  36. // Upload sends a POST request to a volume server to upload the content with adjustable compression level
  37. func UploadWithLocalCompressionLevel(uploadUrl string, filename string, reader io.Reader, isGzipped bool, mtype string, pairMap map[string]string, jwt security.EncodedJwt, compressionLevel int) (*UploadResult, error) {
  38. if compressionLevel < 1 {
  39. compressionLevel = 1
  40. }
  41. if compressionLevel > 9 {
  42. compressionLevel = 9
  43. }
  44. return doUpload(uploadUrl, filename, reader, isGzipped, mtype, pairMap, compressionLevel, jwt)
  45. }
  46. // Upload sends a POST request to a volume server to upload the content with fast compression
  47. func Upload(uploadUrl string, filename string, reader io.Reader, isGzipped bool, mtype string, pairMap map[string]string, jwt security.EncodedJwt) (*UploadResult, error) {
  48. return doUpload(uploadUrl, filename, reader, isGzipped, mtype, pairMap, flate.BestSpeed, jwt)
  49. }
  50. func doUpload(uploadUrl string, filename string, reader io.Reader, isGzipped bool, mtype string, pairMap map[string]string, compression int, jwt security.EncodedJwt) (*UploadResult, error) {
  51. contentIsGzipped := isGzipped
  52. shouldGzipNow := false
  53. if !isGzipped {
  54. if shouldBeZipped, iAmSure := util.IsGzippableFileType(filepath.Base(filename), mtype); iAmSure && shouldBeZipped {
  55. shouldGzipNow = true
  56. contentIsGzipped = true
  57. }
  58. }
  59. return upload_content(uploadUrl, func(w io.Writer) (err error) {
  60. if shouldGzipNow {
  61. gzWriter, _ := gzip.NewWriterLevel(w, compression)
  62. _, err = io.Copy(gzWriter, reader)
  63. gzWriter.Close()
  64. } else {
  65. _, err = io.Copy(w, reader)
  66. }
  67. return
  68. }, filename, contentIsGzipped, mtype, pairMap, jwt)
  69. }
  70. func upload_content(uploadUrl string, fillBufferFunction func(w io.Writer) error, filename string, isGzipped bool, mtype string, pairMap map[string]string, jwt security.EncodedJwt) (*UploadResult, error) {
  71. body_buf := bytes.NewBufferString("")
  72. body_writer := multipart.NewWriter(body_buf)
  73. h := make(textproto.MIMEHeader)
  74. h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, fileNameEscaper.Replace(filename)))
  75. if mtype == "" {
  76. mtype = mime.TypeByExtension(strings.ToLower(filepath.Ext(filename)))
  77. }
  78. if mtype != "" {
  79. h.Set("Content-Type", mtype)
  80. }
  81. if isGzipped {
  82. h.Set("Content-Encoding", "gzip")
  83. }
  84. file_writer, cp_err := body_writer.CreatePart(h)
  85. if cp_err != nil {
  86. glog.V(0).Infoln("error creating form file", cp_err.Error())
  87. return nil, cp_err
  88. }
  89. if err := fillBufferFunction(file_writer); err != nil {
  90. glog.V(0).Infoln("error copying data", err)
  91. return nil, err
  92. }
  93. content_type := body_writer.FormDataContentType()
  94. if err := body_writer.Close(); err != nil {
  95. glog.V(0).Infoln("error closing body", err)
  96. return nil, err
  97. }
  98. req, postErr := http.NewRequest("POST", uploadUrl, body_buf)
  99. if postErr != nil {
  100. glog.V(0).Infoln("failing to upload to", uploadUrl, postErr.Error())
  101. return nil, postErr
  102. }
  103. req.Header.Set("Content-Type", content_type)
  104. for k, v := range pairMap {
  105. req.Header.Set(k, v)
  106. }
  107. if jwt != "" {
  108. req.Header.Set("Authorization", "BEARER "+string(jwt))
  109. }
  110. resp, post_err := client.Do(req)
  111. if post_err != nil {
  112. glog.V(0).Infoln("failing to upload to", uploadUrl, post_err.Error())
  113. return nil, post_err
  114. }
  115. defer resp.Body.Close()
  116. etag := getEtag(resp)
  117. resp_body, ra_err := ioutil.ReadAll(resp.Body)
  118. if ra_err != nil {
  119. return nil, ra_err
  120. }
  121. var ret UploadResult
  122. unmarshal_err := json.Unmarshal(resp_body, &ret)
  123. if unmarshal_err != nil {
  124. glog.V(0).Infoln("failing to read upload response", uploadUrl, string(resp_body))
  125. return nil, unmarshal_err
  126. }
  127. if ret.Error != "" {
  128. return nil, errors.New(ret.Error)
  129. }
  130. ret.ETag = etag
  131. return &ret, nil
  132. }
  133. func getEtag(r *http.Response) (etag string) {
  134. etag = r.Header.Get("ETag")
  135. if strings.HasPrefix(etag, "\"") && strings.HasSuffix(etag, "\"") {
  136. etag = etag[1 : len(etag)-1]
  137. }
  138. return
  139. }