GODT-1146: Refactor header filtering

internal/imap/mailbox_header.go

@@ -0,0 +1,104 @@
+// Copyright (c) 2021 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.
+//
+// ProtonMail Bridge is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// ProtonMail Bridge is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
+
+package imap
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"strings"
+
+	"github.com/emersion/go-imap"
+	"github.com/pkg/errors"
+)
+
+func filterHeader(header []byte, section *imap.BodySectionName) []byte {
+	// Empty section.Fields means BODY[HEADER] was requested so we should return the full header.
+	if len(section.Fields) == 0 {
+		return header
+	}
+
+	fieldMap := make(map[string]struct{})
+
+	for _, field := range section.Fields {
+		fieldMap[strings.ToLower(field)] = struct{}{}
+	}
+
+	return filterHeaderLines(header, func(field string) bool {
+		_, ok := fieldMap[strings.ToLower(field)]
+
+		if section.NotFields {
+			ok = !ok
+		}
+
+		return ok
+	})
+}
+
+func filterHeaderLines(header []byte, wantField func(string) bool) []byte {
+	var res []byte
+
+	for _, line := range headerLines(header) {
+		if len(bytes.TrimSpace(line)) == 0 {
+			res = append(res, line...)
+		} else {
+			split := bytes.SplitN(line, []byte(": "), 2)
+
+			if len(split) != 2 {
+				continue
+			}
+
+			if wantField(string(bytes.ToLower(split[0]))) {
+				res = append(res, line...)
+			}
+		}
+	}
+
+	return res
+}
+
+// NOTE: This sucks because we trim and split stuff here already, only to do it again when we use this function!
+func headerLines(header []byte) [][]byte {
+	var lines [][]byte
+
+	r := bufio.NewReader(bytes.NewReader(header))
+
+	for {
+		b, err := r.ReadBytes('\n')
+		if err != nil {
+			if err != io.EOF {
+				panic(errors.Wrap(err, "failed to read header line"))
+			}
+
+			break
+		}
+
+		switch {
+		case len(bytes.TrimSpace(b)) == 0:
+			lines = append(lines, b)
+
+		case len(bytes.SplitN(b, []byte(": "), 2)) != 2:
+			lines[len(lines)-1] = append(lines[len(lines)-1], b...)
+
+		default:
+			lines = append(lines, b)
+		}
+	}
+
+	return lines
+}