117 lines
2.8 KiB
Go
117 lines
2.8 KiB
Go
package message
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/textproto"
|
|
"sort"
|
|
"strings"
|
|
)
|
|
|
|
// From https://golang.org/src/mime/multipart/writer.go?s=2140:2215#L76
|
|
func writeHeader(w io.Writer, header Header) error {
|
|
keys := make([]string, 0, len(header))
|
|
for k := range header {
|
|
keys = append(keys, k)
|
|
}
|
|
sort.Strings(keys)
|
|
for _, k := range keys {
|
|
for _, v := range header[k] {
|
|
if _, err := io.WriteString(w, formatHeaderField(k, v)); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
_, err := io.WriteString(w, "\r\n")
|
|
return err
|
|
}
|
|
|
|
// A Writer formats entities.
|
|
type Writer struct {
|
|
w io.Writer
|
|
c io.Closer
|
|
mw *multipart.Writer
|
|
}
|
|
|
|
// newWriter creates a new Writer writing to w with the provided header. Nothing
|
|
// is written to w when it is called. header is modified in-place.
|
|
func newWriter(w io.Writer, header Header) *Writer {
|
|
ww := &Writer{w: w}
|
|
|
|
mediaType, mediaParams, _ := header.ContentType()
|
|
if strings.HasPrefix(mediaType, "multipart/") {
|
|
ww.mw = multipart.NewWriter(ww.w)
|
|
|
|
// Do not set ww's io.Closer for now: if this is a multipart entity but
|
|
// CreatePart is not used (only Write is used), then the final boundary
|
|
// is expected to be written by the user too. In this case, ww.Close
|
|
// shouldn't write the final boundary.
|
|
|
|
if mediaParams["boundary"] != "" {
|
|
ww.mw.SetBoundary(mediaParams["boundary"])
|
|
} else {
|
|
mediaParams["boundary"] = ww.mw.Boundary()
|
|
header.SetContentType(mediaType, mediaParams)
|
|
}
|
|
|
|
header.Del("Content-Transfer-Encoding")
|
|
} else {
|
|
wc := encodingWriter(header.Get("Content-Transfer-Encoding"), ww.w)
|
|
ww.w = wc
|
|
ww.c = wc
|
|
}
|
|
|
|
return ww
|
|
}
|
|
|
|
// CreateWriter creates a new Writer writing to w. If header contains an
|
|
// encoding, data written to the Writer will automatically be encoded with it.
|
|
func CreateWriter(w io.Writer, header Header) (*Writer, error) {
|
|
ww := newWriter(w, header)
|
|
if err := writeHeader(w, header); err != nil {
|
|
return nil, err
|
|
}
|
|
return ww, nil
|
|
}
|
|
|
|
// Write implements io.Writer.
|
|
func (w *Writer) Write(b []byte) (int, error) {
|
|
return w.w.Write(b)
|
|
}
|
|
|
|
// Close implements io.Closer.
|
|
func (w *Writer) Close() error {
|
|
if w.c != nil {
|
|
return w.c.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// CreatePart returns a Writer to a new part in this multipart entity. If this
|
|
// entity is not multipart, it fails. The body of the part should be written to
|
|
// the returned io.WriteCloser.
|
|
func (w *Writer) CreatePart(header Header) (*Writer, error) {
|
|
if w.mw == nil {
|
|
return nil, errors.New("cannot create a part in a non-multipart message")
|
|
}
|
|
|
|
if w.c == nil {
|
|
// We know that the user calls CreatePart so Close should write the final
|
|
// boundary
|
|
w.c = w.mw
|
|
}
|
|
|
|
// cw -> ww -> pw -> w.mw -> w.w
|
|
|
|
ww := &struct{ io.Writer }{nil}
|
|
cw := newWriter(ww, header)
|
|
pw, err := w.mw.CreatePart(textproto.MIMEHeader(header))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
ww.Writer = pw
|
|
return cw, nil
|
|
}
|