1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135
|
package message
import (
"errors"
"fmt"
"io"
"strings"
"github.com/emersion/go-message/textproto"
)
// Writer writes message entities.
//
// If the message is not multipart, it should be used as a WriteCloser. Don't
// forget to call Close.
//
// If the message is multipart, users can either use CreatePart to write child
// parts or Write to directly pipe a multipart message. In any case, Close must
// be called at the end.
type Writer struct {
w io.Writer
c io.Closer
mw *textproto.MultipartWriter
}
// createWriter 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 createWriter(w io.Writer, header *Header) (*Writer, error) {
ww := &Writer{w: w}
mediaType, mediaParams, _ := header.ContentType()
if strings.HasPrefix(mediaType, "multipart/") {
ww.mw = textproto.NewMultipartWriter(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, err := encodingWriter(header.Get("Content-Transfer-Encoding"), ww.w)
if err != nil {
return nil, err
}
ww.w = wc
ww.c = wc
}
switch strings.ToLower(mediaParams["charset"]) {
case "", "us-ascii", "utf-8":
// This is OK
default:
// Anything else is invalid
return nil, fmt.Errorf("unhandled charset %q", mediaParams["charset"])
}
return ww, nil
}
// CreateWriter creates a new message writer to w. If header contains an
// encoding, data written to the Writer will automatically be encoded with it.
// The charset needs to be utf-8 or us-ascii.
func CreateWriter(w io.Writer, header Header) (*Writer, error) {
// ensure that modifications are invisible to the caller
header = header.Copy()
// If the message uses MIME, it has to include MIME-Version
if !header.Has("Mime-Version") {
header.Set("MIME-Version", "1.0")
}
ww, err := createWriter(w, &header)
if err != nil {
return nil, err
}
if err := textproto.WriteHeader(w, header.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}
// ensure that modifications are invisible to the caller
header = header.Copy()
cw, err := createWriter(ww, &header)
if err != nil {
return nil, err
}
pw, err := w.mw.CreatePart(header.Header)
if err != nil {
return nil, err
}
ww.Writer = pw
return cw, nil
}
|