kittemail/message.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/google/uuid"
	"io"
	"strings"
	"sync"
	"time"

	"github.com/emersion/go-imap/v2"
	"github.com/emersion/go-imap/v2/imapserver"
	gomessage "github.com/emersion/go-message"
	"github.com/emersion/go-message/mail"
	"github.com/emersion/go-message/textproto"
)

type MessageBuffer struct {
	underlyingBuf io.ReadSeekCloser
	closed        bool
	parent        *message
}

func (b *MessageBuffer) Read(p []byte) (n int, err error) {
	if b.closed {
		return 0, io.ErrClosedPipe
	}
	return b.underlyingBuf.Read(p)
}

func (b *MessageBuffer) Seek(offset int64, whence int) (int64, error) {
	if b.closed {
		return 0, io.ErrClosedPipe
	}
	return b.underlyingBuf.Seek(offset, whence)
}

// Close releases the lock on the parent message. It does not close the underlying buffer.
// It is safe to call Close multiple times.
func (b *MessageBuffer) Close() {
	if b.closed {
		return
	}
	b.closed = true
	b.parent.bufMutex.Unlock()
}

// Reset resets the buffer to the initial state, once again reading from the beginning of the underlying buffer.
func (b *MessageBuffer) Reset() error {
	if b.closed {
		return io.ErrClosedPipe
	}
	_, err := b.underlyingBuf.Seek(0, io.SeekStart)
	return err
}

type message struct {
	// immutable
	uid           imap.UID
	underlyingBuf io.ReadSeekCloser
	bufMutex      sync.Mutex
	len           int64
	t             time.Time
	id            uuid.UUID

	// mutable, protected by Mailbox.mutex
	flags map[imap.Flag]struct{}
}

func (msg *message) buf() (*MessageBuffer, error) {
	msg.bufMutex.Lock()
	buf := &MessageBuffer{msg.underlyingBuf, false, msg}
	err := buf.Reset()
	if err != nil {
		msg.bufMutex.Unlock()
		return nil, err
	}
	return buf, nil
}

func (msg *message) fetch(w *imapserver.FetchResponseWriter, options *imap.FetchOptions) error {
	w.WriteUID(msg.uid)

	if options.Flags {
		w.WriteFlags(msg.flagList())
	}
	if options.InternalDate {
		w.WriteInternalDate(msg.t)
	}
	if options.RFC822Size {
		w.WriteRFC822Size(msg.len)
	}
	if options.Envelope {
		w.WriteEnvelope(msg.envelope())
	}
	if options.BodyStructure != nil {
		buf, err := msg.buf()
		if err != nil {
			return err
		}
		w.WriteBodyStructure(imapserver.ExtractBodyStructure(buf))
		buf.Close()
	}

	for _, bs := range options.BodySection {
		msgBuf, err := msg.buf()
		if err != nil {
			return err
		}
		buf := imapserver.ExtractBodySection(msgBuf, bs)
		wc := w.WriteBodySection(bs, int64(len(buf)))
		_, writeErr := wc.Write(buf)
		closeErr := wc.Close()
		msgBuf.Close()
		if writeErr != nil {
			return writeErr
		}
		if closeErr != nil {
			return closeErr
		}
	}

	for _, bs := range options.BinarySection {
		msgBuf, err := msg.buf()
		if err != nil {
			return err
		}
		buf := imapserver.ExtractBinarySection(msgBuf, bs)
		wc := w.WriteBinarySection(bs, int64(len(buf)))
		_, writeErr := wc.Write(buf)
		closeErr := wc.Close()
		msgBuf.Close()
		if writeErr != nil {
			return writeErr
		}
		if closeErr != nil {
			return closeErr
		}
	}

	for _, bss := range options.BinarySectionSize {
		buf, err := msg.buf()
		if err != nil {
			return err
		}
		n := imapserver.ExtractBinarySectionSize(buf, bss)
		w.WriteBinarySectionSize(bss, n)
		buf.Close()
	}

	return w.Close()
}

func (msg *message) envelope() *imap.Envelope {
	buf, err := msg.buf()
	if err != nil {
		return nil
	}
	br := bufio.NewReader(buf)
	header, err := textproto.ReadHeader(br)
	if err != nil {
		return nil
	}
	buf.Close()
	return imapserver.ExtractEnvelope(header)
}

func (msg *message) flagList() []imap.Flag {
	var flags []imap.Flag
	for flag := range msg.flags {
		flags = append(flags, flag)
	}
	return flags
}

func (msg *message) store(store *imap.StoreFlags) {
	switch store.Op {
	case imap.StoreFlagsSet:
		msg.flags = make(map[imap.Flag]struct{})
		fallthrough
	case imap.StoreFlagsAdd:
		for _, flag := range store.Flags {
			msg.flags[canonicalFlag(flag)] = struct{}{}
		}
	case imap.StoreFlagsDel:
		for _, flag := range store.Flags {
			delete(msg.flags, canonicalFlag(flag))
		}
	default:
		panic(fmt.Errorf("unknown STORE flag operation: %v", store.Op))
	}
}

func (msg *message) reader() (*gomessage.Entity, error) {
	buf, err := msg.buf()
	if err != nil {
		return nil, err
	}
	r, _ := gomessage.Read(buf)
	buf.Close()
	if r == nil {
		r, _ = gomessage.New(gomessage.Header{}, bytes.NewReader(nil))
	}
	return r, nil
}

func (msg *message) search(seqNum uint32, criteria *imap.SearchCriteria) (bool, error) {
	for _, seqSet := range criteria.SeqNum {
		if seqNum == 0 || !seqSet.Contains(seqNum) {
			return false, nil
		}
	}
	for _, uidSet := range criteria.UID {
		if !uidSet.Contains(msg.uid) {
			return false, nil
		}
	}
	if !matchDate(msg.t, criteria.Since, criteria.Before) {
		return false, nil
	}

	for _, flag := range criteria.Flag {
		if _, ok := msg.flags[canonicalFlag(flag)]; !ok {
			return false, nil
		}
	}
	for _, flag := range criteria.NotFlag {
		if _, ok := msg.flags[canonicalFlag(flag)]; ok {
			return false, nil
		}
	}

	if criteria.Larger != 0 && msg.len <= criteria.Larger {
		return false, nil
	}
	if criteria.Smaller != 0 && msg.len >= criteria.Smaller {
		return false, nil
	}

	msgReader, err := msg.reader()
	if err != nil {
		return false, err
	}

	header := mail.Header{Header: msgReader.Header}

	for _, fieldCriteria := range criteria.Header {
		if !matchHeaderFields(header.FieldsByKey(fieldCriteria.Key), fieldCriteria.Value) {
			return false, nil
		}
	}

	if !criteria.SentSince.IsZero() || !criteria.SentBefore.IsZero() {
		t, err := header.Date()
		if err != nil {
			return false, err
		} else if !matchDate(t, criteria.SentSince, criteria.SentBefore) {
			return false, nil
		}
	}

	for _, text := range criteria.Text {
		msgReader, err := msg.reader()
		if err != nil {
			return false, err
		}
		if !matchEntity(msgReader, text, true) {
			return false, nil
		}
	}
	for _, body := range criteria.Body {
		msgReader, err := msg.reader()
		if err != nil {
			return false, err
		}
		if !matchEntity(msgReader, body, false) {
			return false, nil
		}
	}

	for _, not := range criteria.Not {
		search, err := msg.search(seqNum, &not)
		if err != nil {
			return false, err
		}
		if search {
			return false, nil
		}
	}
	for _, or := range criteria.Or {
		search, err := msg.search(seqNum, &or[0])
		if err != nil {
			return false, err
		}
		search2, err := msg.search(seqNum, &or[1])
		if !search && !search2 {
			return false, nil
		}
	}

	return true, nil
}

func matchDate(t, since, before time.Time) bool {
	// We discard time zone information by setting it to UTC.
	// RFC 3501 explicitly requires zone unaware date comparison.
	t = time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)

	if !since.IsZero() && t.Before(since) {
		return false
	}
	if !before.IsZero() && !t.Before(before) {
		return false
	}
	return true
}

func matchHeaderFields(fields gomessage.HeaderFields, pattern string) bool {
	if pattern == "" {
		return fields.Len() > 0
	}

	pattern = strings.ToLower(pattern)
	for fields.Next() {
		v, _ := fields.Text()
		if strings.Contains(strings.ToLower(v), pattern) {
			return true
		}
	}
	return false
}

func matchEntity(e *gomessage.Entity, pattern string, includeHeader bool) bool {
	if pattern == "" {
		return true
	}

	if includeHeader && matchHeaderFields(e.Header.Fields(), pattern) {
		return true
	}

	if mr := e.MultipartReader(); mr != nil {
		for {
			part, err := mr.NextPart()
			if err == io.EOF {
				break
			} else if err != nil {
				return false
			}

			if matchEntity(part, pattern, includeHeader) {
				return true
			}
		}

		return false
	} else {
		t, _, err := e.Header.ContentType()
		if err != nil {
			return false
		}

		if !strings.HasPrefix(t, "text/") && !strings.HasPrefix(t, "message/") {
			return false
		}

		buf, err := io.ReadAll(e.Body)
		if err != nil {
			return false
		}

		return bytes.Contains(bytes.ToLower(buf), bytes.ToLower([]byte(pattern)))
	}
}

func canonicalFlag(flag imap.Flag) imap.Flag {
	return imap.Flag(strings.ToLower(string(flag)))
}