Source file src/pkg/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		"errors"
    11		"fmt"
    12		"io"
    13		"net/textproto"
    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, 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	var quoteEscaper = strings.NewReplacer("\\", "\\\\", `"`, "\\\"")
    89	
    90	func escapeQuotes(s string) string {
    91		return quoteEscaper.Replace(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, 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, 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) 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() 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     error // last error that occurred writing
   141	}
   142	
   143	func (p *part) close() error {
   144		p.closed = true
   145		return p.we
   146	}
   147	
   148	func (p *part) Write(d []byte) (n int, err error) {
   149		if p.closed {
   150			return 0, errors.New("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	}