libgo/go/mime/multipart/writer.go

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