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