writer.go 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185
  1. // Copyright 2011 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 multipart
  5. import (
  6. "bytes"
  7. "crypto/rand"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "net/textproto"
  12. "strings"
  13. )
  14. // A Writer generates multipart messages.
  15. type Writer struct {
  16. w io.Writer
  17. boundary string
  18. lastpart *part
  19. }
  20. // NewWriter returns a new multipart Writer with a random boundary,
  21. // writing to w.
  22. func NewWriter(w io.Writer) *Writer {
  23. return &Writer{
  24. w: w,
  25. boundary: randomBoundary(),
  26. }
  27. }
  28. // Boundary returns the Writer's boundary.
  29. func (w *Writer) Boundary() string {
  30. return w.boundary
  31. }
  32. // SetBoundary overrides the Writer's default randomly-generated
  33. // boundary separator with an explicit value.
  34. //
  35. // SetBoundary must be called before any parts are created, may only
  36. // contain certain ASCII characters, and must be 1-69 bytes long.
  37. func (w *Writer) SetBoundary(boundary string) error {
  38. if w.lastpart != nil {
  39. return errors.New("mime: SetBoundary called after write")
  40. }
  41. // rfc2046#section-5.1.1
  42. if len(boundary) < 1 || len(boundary) > 69 {
  43. return errors.New("mime: invalid boundary length")
  44. }
  45. for _, b := range boundary {
  46. if 'A' <= b && b <= 'Z' || 'a' <= b && b <= 'z' || '0' <= b && b <= '9' {
  47. continue
  48. }
  49. switch b {
  50. case '\'', '(', ')', '+', '_', ',', '-', '.', '/', ':', '=', '?':
  51. continue
  52. }
  53. return errors.New("mime: invalid boundary character")
  54. }
  55. w.boundary = boundary
  56. return nil
  57. }
  58. // FormDataContentType returns the Content-Type for an HTTP
  59. // multipart/form-data with this Writer's Boundary.
  60. func (w *Writer) FormDataContentType() string {
  61. return "multipart/form-data; boundary=" + w.boundary
  62. }
  63. func randomBoundary() string {
  64. var buf [30]byte
  65. _, err := io.ReadFull(rand.Reader, buf[:])
  66. if err != nil {
  67. panic(err)
  68. }
  69. return fmt.Sprintf("%x", buf[:])
  70. }
  71. // CreatePart creates a new multipart section with the provided
  72. // header. The body of the part should be written to the returned
  73. // Writer. After calling CreatePart, any previous part may no longer
  74. // be written to.
  75. func (w *Writer) CreatePart(header textproto.MIMEHeader) (io.Writer, error) {
  76. if w.lastpart != nil {
  77. if err := w.lastpart.close(); err != nil {
  78. return nil, err
  79. }
  80. }
  81. var b bytes.Buffer
  82. if w.lastpart != nil {
  83. fmt.Fprintf(&b, "\r\n--%s\r\n", w.boundary)
  84. } else {
  85. fmt.Fprintf(&b, "--%s\r\n", w.boundary)
  86. }
  87. // TODO(bradfitz): move this to textproto.MimeHeader.Write(w), have it sort
  88. // and clean, like http.Header.Write(w) does.
  89. for k, vv := range header {
  90. for _, v := range vv {
  91. fmt.Fprintf(&b, "%s: %s\r\n", k, v)
  92. }
  93. }
  94. fmt.Fprintf(&b, "\r\n")
  95. _, err := io.Copy(w.w, &b)
  96. if err != nil {
  97. return nil, err
  98. }
  99. p := &part{
  100. mw: w,
  101. }
  102. w.lastpart = p
  103. return p, nil
  104. }
  105. var quoteEscaper = strings.NewReplacer("\\", "\\\\", `"`, "\\\"")
  106. func escapeQuotes(s string) string {
  107. return quoteEscaper.Replace(s)
  108. }
  109. // CreateFormFile is a convenience wrapper around CreatePart. It creates
  110. // a new form-data header with the provided field name and file name.
  111. func (w *Writer) CreateFormFile(fieldname, filename string) (io.Writer, error) {
  112. h := make(textproto.MIMEHeader)
  113. h.Set("Content-Disposition",
  114. fmt.Sprintf(`form-data; name="%s"; filename="%s"`,
  115. escapeQuotes(fieldname), escapeQuotes(filename)))
  116. h.Set("Content-Type", "application/octet-stream")
  117. return w.CreatePart(h)
  118. }
  119. // CreateFormField calls CreatePart with a header using the
  120. // given field name.
  121. func (w *Writer) CreateFormField(fieldname string) (io.Writer, error) {
  122. h := make(textproto.MIMEHeader)
  123. h.Set("Content-Disposition",
  124. fmt.Sprintf(`form-data; name="%s"`, escapeQuotes(fieldname)))
  125. return w.CreatePart(h)
  126. }
  127. // WriteField calls CreateFormField and then writes the given value.
  128. func (w *Writer) WriteField(fieldname, value string) error {
  129. p, err := w.CreateFormField(fieldname)
  130. if err != nil {
  131. return err
  132. }
  133. _, err = p.Write([]byte(value))
  134. return err
  135. }
  136. // Close finishes the multipart message and writes the trailing
  137. // boundary end line to the output.
  138. func (w *Writer) Close() error {
  139. if w.lastpart != nil {
  140. if err := w.lastpart.close(); err != nil {
  141. return err
  142. }
  143. w.lastpart = nil
  144. }
  145. _, err := fmt.Fprintf(w.w, "\r\n--%s--\r\n", w.boundary)
  146. return err
  147. }
  148. type part struct {
  149. mw *Writer
  150. closed bool
  151. we error // last error that occurred writing
  152. }
  153. func (p *part) close() error {
  154. p.closed = true
  155. return p.we
  156. }
  157. func (p *part) Write(d []byte) (n int, err error) {
  158. if p.closed {
  159. return 0, errors.New("multipart: can't write to finished part")
  160. }
  161. n, err = p.mw.w.Write(d)
  162. if err != nil {
  163. p.we = err
  164. }
  165. return
  166. }