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