forked from emersion/go-message
-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathwriter.go
144 lines (121 loc) · 3.98 KB
/
writer.go
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
136
137
138
139
140
141
142
143
144
package message
import (
"bytes"
"errors"
"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"] != "" {
err := ww.mw.SetBoundary(mediaParams["boundary"])
// If the original boundary is malformed or does not adhere to the RFC
// then we end up in a mismatch between the generated boundary
// -the one we use to delinate parts of the body-
// and the original boundary present in the header.
// This results in an invalid message strucutre.
if err != nil && err != textproto.SetBoundaryCalledAfterWriteErr {
mediaParams["boundary"] = ww.mw.Boundary()
header.SetContentType(mediaType, mediaParams)
}
} 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
}
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) {
var b bytes.Buffer
return w.CreatePartWithBuffer(header, &b)
}
// CreatePartWithBuffer behaves the same as CreatPart but allows the user to specify a target byte buffer used by the
// multipart writer.
func (w *Writer) CreatePartWithBuffer(header Header, buffer *bytes.Buffer) (*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.CreatePartWithBuffer(header.Header, buffer)
if err != nil {
return nil, err
}
ww.Writer = pw
return cw, nil
}