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