1 package textproto
2 3 import (
4 "bufio"
5 "bytes"
6 "fmt"
7 "io"
8 "net/textproto"
9 "sort"
10 "strings"
11 )
12 13 type headerField struct {
14 b []byte // Raw header field, including whitespace
15 16 k string
17 v string
18 }
19 20 func newHeaderField(k, v string, b []byte) *headerField {
21 return &headerField{k: textproto.CanonicalMIMEHeaderKey(k), v: v, b: b}
22 }
23 24 func (f *headerField) raw() ([]byte, error) {
25 if f.b != nil {
26 return f.b, nil
27 } else {
28 for pos, ch := range f.k {
29 // check if character is a printable US-ASCII except ':'
30 if !(ch >= '!' && ch < ':' || ch > ':' && ch <= '~') {
31 return nil, fmt.Errorf("field name contains incorrect symbols (\\x%x at %v)", ch, pos)
32 }
33 }
34 35 if pos := strings.IndexAny(f.v, "\r\n"); pos != -1 {
36 return nil, fmt.Errorf("field value contains \\r\\n (at %v)", pos)
37 }
38 39 return []byte(formatHeaderField(f.k, f.v)), nil
40 }
41 }
42 43 // A Header represents the key-value pairs in a message header.
44 //
45 // The header representation is idempotent: if the header can be read and
46 // written, the result will be exactly the same as the original (including
47 // whitespace and header field ordering). This is required for e.g. DKIM.
48 //
49 // Mutating the header is restricted: the only two allowed operations are
50 // inserting a new header field at the top and deleting a header field. This is
51 // again necessary for DKIM.
52 type Header struct {
53 // Fields are in reverse order so that inserting a new field at the top is
54 // cheap.
55 l []*headerField
56 m map[string][]*headerField
57 }
58 59 func makeHeaderMap(fs []*headerField) map[string][]*headerField {
60 if len(fs) == 0 {
61 return nil
62 }
63 64 m := make(map[string][]*headerField, len(fs))
65 for i, f := range fs {
66 m[f.k] = append(m[f.k], fs[i])
67 }
68 return m
69 }
70 71 func newHeader(fs []*headerField) Header {
72 // Reverse order
73 for i := len(fs)/2 - 1; i >= 0; i-- {
74 opp := len(fs) - 1 - i
75 fs[i], fs[opp] = fs[opp], fs[i]
76 }
77 78 return Header{l: fs, m: makeHeaderMap(fs)}
79 }
80 81 // HeaderFromMap creates a header from a map of header fields.
82 //
83 // This function is provided for interoperability with the standard library.
84 // If possible, ReadHeader should be used instead to avoid loosing information.
85 // The map representation looses the ordering of the fields, the capitalization
86 // of the header keys, and the whitespace of the original header.
87 func HeaderFromMap(m map[string][]string) Header {
88 fs := make([]*headerField, 0, len(m))
89 for k, vs := range m {
90 for _, v := range vs {
91 fs = append(fs, newHeaderField(k, v, nil))
92 }
93 }
94 95 sort.SliceStable(fs, func(i, j int) bool {
96 return fs[i].k < fs[j].k
97 })
98 99 return newHeader(fs)
100 }
101 102 // AddRaw adds the raw key, value pair to the header.
103 //
104 // The supplied byte slice should be a complete field in the "Key: Value" form
105 // including trailing CRLF. If there is no comma in the input - AddRaw panics.
106 // No changes are made to kv contents and it will be copied into WriteHeader
107 // output as is.
108 //
109 // kv is directly added to the underlying structure and therefore should not be
110 // modified after the AddRaw call.
111 func (h *Header) AddRaw(kv []byte) {
112 colon := bytes.IndexByte(kv, ':')
113 if colon == -1 {
114 panic("textproto: Header.AddRaw: missing colon")
115 }
116 k := textproto.CanonicalMIMEHeaderKey(string(trim(kv[:colon])))
117 v := trimAroundNewlines(kv[colon+1:])
118 119 if h.m == nil {
120 h.m = make(map[string][]*headerField)
121 }
122 123 f := newHeaderField(k, v, kv)
124 h.l = append(h.l, f)
125 h.m[k] = append(h.m[k], f)
126 }
127 128 // Add adds the key, value pair to the header. It prepends to any existing
129 // fields associated with key.
130 //
131 // Key and value should obey character requirements of RFC 6532.
132 // If you need to format or fold lines manually, use AddRaw.
133 func (h *Header) Add(k, v string) {
134 k = textproto.CanonicalMIMEHeaderKey(k)
135 136 if h.m == nil {
137 h.m = make(map[string][]*headerField)
138 }
139 140 f := newHeaderField(k, v, nil)
141 h.l = append(h.l, f)
142 h.m[k] = append(h.m[k], f)
143 }
144 145 // Get gets the first value associated with the given key. If there are no
146 // values associated with the key, Get returns "".
147 func (h *Header) Get(k string) string {
148 fields := h.m[textproto.CanonicalMIMEHeaderKey(k)]
149 if len(fields) == 0 {
150 return ""
151 }
152 return fields[len(fields)-1].v
153 }
154 155 // Raw gets the first raw header field associated with the given key.
156 //
157 // The returned bytes contain a complete field in the "Key: value" form,
158 // including trailing CRLF.
159 //
160 // The returned slice should not be modified and becomes invalid when the
161 // header is updated.
162 //
163 // An error is returned if the header field contains incorrect characters (see
164 // RFC 6532).
165 func (h *Header) Raw(k string) ([]byte, error) {
166 fields := h.m[textproto.CanonicalMIMEHeaderKey(k)]
167 if len(fields) == 0 {
168 return nil, nil
169 }
170 return fields[len(fields)-1].raw()
171 }
172 173 // Values returns all values associated with the given key.
174 //
175 // The returned slice should not be modified and becomes invalid when the
176 // header is updated.
177 func (h *Header) Values(k string) []string {
178 fields := h.m[textproto.CanonicalMIMEHeaderKey(k)]
179 if len(fields) == 0 {
180 return nil
181 }
182 l := make([]string, len(fields))
183 for i, field := range fields {
184 l[len(fields)-i-1] = field.v
185 }
186 return l
187 }
188 189 // Set sets the header fields associated with key to the single field value.
190 // It replaces any existing values associated with key.
191 func (h *Header) Set(k, v string) {
192 h.Del(k)
193 h.Add(k, v)
194 }
195 196 // Del deletes the values associated with key.
197 func (h *Header) Del(k string) {
198 k = textproto.CanonicalMIMEHeaderKey(k)
199 200 delete(h.m, k)
201 202 // Delete existing keys
203 for i := len(h.l) - 1; i >= 0; i-- {
204 if h.l[i].k == k {
205 h.l = append(h.l[:i], h.l[i+1:]...)
206 }
207 }
208 }
209 210 // Has checks whether the header has a field with the specified key.
211 func (h *Header) Has(k string) bool {
212 _, ok := h.m[textproto.CanonicalMIMEHeaderKey(k)]
213 return ok
214 }
215 216 // Copy creates an independent copy of the header.
217 func (h *Header) Copy() Header {
218 l := make([]*headerField, len(h.l))
219 copy(l, h.l)
220 m := makeHeaderMap(l)
221 return Header{l: l, m: m}
222 }
223 224 // Len returns the number of fields in the header.
225 func (h *Header) Len() int {
226 return len(h.l)
227 }
228 229 // Map returns all header fields as a map.
230 //
231 // This function is provided for interoperability with the standard library.
232 // If possible, Fields should be used instead to avoid loosing information.
233 // The map representation looses the ordering of the fields, the capitalization
234 // of the header keys, and the whitespace of the original header.
235 func (h *Header) Map() map[string][]string {
236 m := make(map[string][]string, h.Len())
237 fields := h.Fields()
238 for fields.Next() {
239 m[fields.Key()] = append(m[fields.Key()], fields.Value())
240 }
241 return m
242 }
243 244 // HeaderFields iterates over header fields. Its cursor starts before the first
245 // field of the header. Use Next to advance from field to field.
246 type HeaderFields interface {
247 // Next advances to the next header field. It returns true on success, or
248 // false if there is no next field.
249 Next() (more bool)
250 // Key returns the key of the current field.
251 Key() string
252 // Value returns the value of the current field.
253 Value() string
254 // Raw returns the raw current header field. See Header.Raw.
255 Raw() ([]byte, error)
256 // Del deletes the current field.
257 Del()
258 // Len returns the amount of header fields in the subset of header iterated
259 // by this HeaderFields instance.
260 //
261 // For Fields(), it will return the amount of fields in the whole header section.
262 // For FieldsByKey(), it will return the amount of fields with certain key.
263 Len() int
264 }
265 266 type headerFields struct {
267 h *Header
268 cur int
269 }
270 271 func (fs *headerFields) Next() bool {
272 fs.cur++
273 return fs.cur < len(fs.h.l)
274 }
275 276 func (fs *headerFields) index() int {
277 if fs.cur < 0 {
278 panic("message: HeaderFields method called before Next")
279 }
280 if fs.cur >= len(fs.h.l) {
281 panic("message: HeaderFields method called after Next returned false")
282 }
283 return len(fs.h.l) - fs.cur - 1
284 }
285 286 func (fs *headerFields) field() *headerField {
287 return fs.h.l[fs.index()]
288 }
289 290 func (fs *headerFields) Key() string {
291 return fs.field().k
292 }
293 294 func (fs *headerFields) Value() string {
295 return fs.field().v
296 }
297 298 func (fs *headerFields) Raw() ([]byte, error) {
299 return fs.field().raw()
300 }
301 302 func (fs *headerFields) Del() {
303 f := fs.field()
304 305 ok := false
306 for i, ff := range fs.h.m[f.k] {
307 if ff == f {
308 ok = true
309 fs.h.m[f.k] = append(fs.h.m[f.k][:i], fs.h.m[f.k][i+1:]...)
310 if len(fs.h.m[f.k]) == 0 {
311 delete(fs.h.m, f.k)
312 }
313 break
314 }
315 }
316 if !ok {
317 panic("message: field not found in Header.m")
318 }
319 320 fs.h.l = append(fs.h.l[:fs.index()], fs.h.l[fs.index()+1:]...)
321 fs.cur--
322 }
323 324 func (fs *headerFields) Len() int {
325 return len(fs.h.l)
326 }
327 328 // Fields iterates over all the header fields.
329 //
330 // The header may not be mutated while iterating, except using HeaderFields.Del.
331 func (h *Header) Fields() HeaderFields {
332 return &headerFields{h, -1}
333 }
334 335 type headerFieldsByKey struct {
336 h *Header
337 k string
338 cur int
339 }
340 341 func (fs *headerFieldsByKey) Next() bool {
342 fs.cur++
343 return fs.cur < len(fs.h.m[fs.k])
344 }
345 346 func (fs *headerFieldsByKey) index() int {
347 if fs.cur < 0 {
348 panic("message: headerfields.key or value called before next")
349 }
350 if fs.cur >= len(fs.h.m[fs.k]) {
351 panic("message: headerfields.key or value called after next returned false")
352 }
353 return len(fs.h.m[fs.k]) - fs.cur - 1
354 }
355 356 func (fs *headerFieldsByKey) field() *headerField {
357 return fs.h.m[fs.k][fs.index()]
358 }
359 360 func (fs *headerFieldsByKey) Key() string {
361 return fs.field().k
362 }
363 364 func (fs *headerFieldsByKey) Value() string {
365 return fs.field().v
366 }
367 368 func (fs *headerFieldsByKey) Raw() ([]byte, error) {
369 return fs.field().raw()
370 }
371 372 func (fs *headerFieldsByKey) Del() {
373 f := fs.field()
374 375 ok := false
376 for i := range fs.h.l {
377 if f == fs.h.l[i] {
378 ok = true
379 fs.h.l = append(fs.h.l[:i], fs.h.l[i+1:]...)
380 break
381 }
382 }
383 if !ok {
384 panic("message: field not found in Header.l")
385 }
386 387 fs.h.m[fs.k] = append(fs.h.m[fs.k][:fs.index()], fs.h.m[fs.k][fs.index()+1:]...)
388 if len(fs.h.m[fs.k]) == 0 {
389 delete(fs.h.m, fs.k)
390 }
391 fs.cur--
392 }
393 394 func (fs *headerFieldsByKey) Len() int {
395 return len(fs.h.m[fs.k])
396 }
397 398 // FieldsByKey iterates over all fields having the specified key.
399 //
400 // The header may not be mutated while iterating, except using HeaderFields.Del.
401 func (h *Header) FieldsByKey(k string) HeaderFields {
402 return &headerFieldsByKey{h, textproto.CanonicalMIMEHeaderKey(k), -1}
403 }
404 405 func readLineSlice(r *bufio.Reader, line []byte) ([]byte, error) {
406 for {
407 l, more, err := r.ReadLine()
408 line = append(line, l...)
409 if err != nil {
410 return line, err
411 }
412 413 if !more {
414 break
415 }
416 }
417 418 return line, nil
419 }
420 421 func isSpace(c byte) bool {
422 return c == ' ' || c == '\t'
423 }
424 425 func validHeaderKeyByte(b byte) bool {
426 c := int(b)
427 return c >= 33 && c <= 126 && c != ':'
428 }
429 430 // trim returns s with leading and trailing spaces and tabs removed.
431 // It does not assume Unicode or UTF-8.
432 func trim(s []byte) []byte {
433 i := 0
434 for i < len(s) && isSpace(s[i]) {
435 i++
436 }
437 n := len(s)
438 for n > i && isSpace(s[n-1]) {
439 n--
440 }
441 return s[i:n]
442 }
443 444 func hasContinuationLine(r *bufio.Reader) bool {
445 c, err := r.ReadByte()
446 if err != nil {
447 return false // bufio will keep err until next read.
448 }
449 r.UnreadByte()
450 return isSpace(c)
451 }
452 453 func readContinuedLineSlice(r *bufio.Reader) ([]byte, error) {
454 // Read the first line. We preallocate slice that it enough
455 // for most fields.
456 line, err := readLineSlice(r, make([]byte, 0, 256))
457 if err == io.EOF && len(line) == 0 {
458 // Header without a body
459 return nil, nil
460 } else if err != nil {
461 return nil, err
462 }
463 464 if len(line) == 0 { // blank line - no continuation
465 return line, nil
466 }
467 468 line = append(line, '\r', '\n')
469 470 // Read continuation lines.
471 for hasContinuationLine(r) {
472 line, err = readLineSlice(r, line)
473 if err != nil {
474 break // bufio will keep err until next read.
475 }
476 477 line = append(line, '\r', '\n')
478 }
479 480 return line, nil
481 }
482 483 func writeContinued(b *strings.Builder, l []byte) {
484 // Strip trailing \r, if any
485 if len(l) > 0 && l[len(l)-1] == '\r' {
486 l = l[:len(l)-1]
487 }
488 l = trim(l)
489 if len(l) == 0 {
490 return
491 }
492 if b.Len() > 0 {
493 b.WriteByte(' ')
494 }
495 b.Write(l)
496 }
497 498 // Strip newlines and spaces around newlines.
499 func trimAroundNewlines(v []byte) string {
500 var b strings.Builder
501 b.Grow(len(v))
502 for {
503 i := bytes.IndexByte(v, '\n')
504 if i < 0 {
505 writeContinued(&b, v)
506 break
507 }
508 writeContinued(&b, v[:i])
509 v = v[i+1:]
510 }
511 512 return b.String()
513 }
514 515 // ReadHeader reads a MIME header from r. The header is a sequence of possibly
516 // continued "Key: Value" lines ending in a blank line.
517 //
518 // To avoid denial of service attacks, the provided bufio.Reader should be
519 // reading from an io.LimitedReader or a similar Reader to bound the size of
520 // headers.
521 func ReadHeader(r *bufio.Reader) (Header, error) {
522 fs := make([]*headerField, 0, 32)
523 524 // The first line cannot start with a leading space.
525 if buf, err := r.Peek(1); err == nil && isSpace(buf[0]) {
526 line, err := readLineSlice(r, nil)
527 if err != nil {
528 return newHeader(fs), err
529 }
530 531 return newHeader(fs), fmt.Errorf("message: malformed MIME header initial line: %v", string(line))
532 }
533 534 for {
535 kv, err := readContinuedLineSlice(r)
536 if len(kv) == 0 {
537 return newHeader(fs), err
538 }
539 540 // Key ends at first colon; should not have trailing spaces but they
541 // appear in the wild, violating specs, so we remove them if present.
542 i := bytes.IndexByte(kv, ':')
543 if i < 0 {
544 return newHeader(fs), fmt.Errorf("message: malformed MIME header line: %v", string(kv))
545 }
546 547 keyBytes := trim(kv[:i])
548 549 // Verify that there are no invalid characters in the header key.
550 // See RFC 5322 Section 2.2
551 for _, c := range keyBytes {
552 if !validHeaderKeyByte(c) {
553 return newHeader(fs), fmt.Errorf("message: malformed MIME header key: %v", string(keyBytes))
554 }
555 }
556 557 key := textproto.CanonicalMIMEHeaderKey(string(keyBytes))
558 559 // As per RFC 7230 field-name is a token, tokens consist of one or more
560 // chars. We could return a an error here, but better to be liberal in
561 // what we accept, so if we get an empty key, skip it.
562 if key == "" {
563 continue
564 }
565 566 i++ // skip colon
567 v := kv[i:]
568 569 value := trimAroundNewlines(v)
570 fs = append(fs, newHeaderField(key, value, kv))
571 572 if err != nil {
573 return newHeader(fs), err
574 }
575 }
576 }
577 578 func foldLine(v string, maxlen int) (line, next string, ok bool) {
579 ok = true
580 581 // We'll need to fold before maxlen
582 foldBefore := maxlen + 1
583 foldAt := len(v)
584 585 var folding string
586 if foldBefore > len(v) {
587 // We reached the end of the string
588 if v[len(v)-1] != '\n' {
589 // If there isn't already a trailing CRLF, insert one
590 folding = "\r\n"
591 }
592 } else {
593 // Find the closest whitespace before maxlen
594 foldAt = strings.LastIndexAny(v[:foldBefore], " \t\n")
595 596 if foldAt == 0 {
597 // The whitespace we found was the previous folding WSP
598 foldAt = foldBefore - 1
599 } else if foldAt < 0 {
600 // We didn't find any whitespace, we have to insert one
601 foldAt = foldBefore - 2
602 }
603 604 switch v[foldAt] {
605 case ' ', '\t':
606 if v[foldAt-1] != '\n' {
607 folding = "\r\n" // The next char will be a WSP, don't need to insert one
608 }
609 case '\n':
610 folding = "" // There is already a CRLF, nothing to do
611 default:
612 // Another char, we need to insert CRLF + WSP. This will insert an
613 // extra space in the string, so this should be avoided if
614 // possible.
615 folding = "\r\n "
616 ok = false
617 }
618 }
619 620 return v[:foldAt] + folding, v[foldAt:], ok
621 }
622 623 const (
624 preferredHeaderLen = 76
625 maxHeaderLen = 998
626 )
627 628 // formatHeaderField formats a header field, ensuring each line is no longer
629 // than 76 characters. It tries to fold lines at whitespace characters if
630 // possible. If the header contains a word longer than this limit, it will be
631 // split.
632 func formatHeaderField(k, v string) string {
633 s := k + ": "
634 635 if v == "" {
636 return s + "\r\n"
637 }
638 639 first := true
640 for len(v) > 0 {
641 // If this is the first line, substract the length of the key
642 keylen := 0
643 if first {
644 keylen = len(s)
645 }
646 647 // First try with a soft limit
648 l, next, ok := foldLine(v, preferredHeaderLen-keylen)
649 if !ok {
650 // Folding failed to preserve the original header field value. Try
651 // with a larger, hard limit.
652 l, next, _ = foldLine(v, maxHeaderLen-keylen)
653 }
654 v = next
655 s += l
656 first = false
657 }
658 659 return s
660 }
661 662 // WriteHeader writes a MIME header to w.
663 func WriteHeader(w io.Writer, h Header) error {
664 for i := len(h.l) - 1; i >= 0; i-- {
665 f := h.l[i]
666 if rawField, err := f.raw(); err == nil {
667 if _, err := w.Write(rawField); err != nil {
668 return err
669 }
670 } else {
671 return fmt.Errorf("failed to write header field #%v (%q): %w", len(h.l)-i, f.k, err)
672 }
673 }
674 675 _, err := w.Write([]byte{'\r', '\n'})
676 return err
677 }
678