header.go raw

   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