proton-bridge/internal/transfer/provider_imap_source.go

// Copyright (c) 2020 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 transfer

import (
	"fmt"
	"io/ioutil"

	"github.com/emersion/go-imap"
)

type imapMessageInfo struct {
	id   string
	uid  uint32
	size uint32
}

const (
	imapPageSize     = uint32(2000)             // Optimized on Gmail.
	imapMaxFetchSize = uint32(50 * 1000 * 1000) // Size in octets. If 0, it will use one fetch per message.
)

// TransferTo exports messages based on rules to channel.
func (p *IMAPProvider) TransferTo(rules transferRules, progress *Progress, ch chan<- Message) {
	log.Info("Started transfer from IMAP to channel")
	defer log.Info("Finished transfer from IMAP to channel")

	imapMessageInfoMap := p.loadMessageInfoMap(rules, progress)

	for rule := range rules.iterateActiveRules() {
		log.WithField("rule", rule).Debug("Processing rule")
		messagesInfo := imapMessageInfoMap[rule.SourceMailbox.Name]
		p.transferTo(rule, messagesInfo, progress, ch)
	}
}

func (p *IMAPProvider) loadMessageInfoMap(rules transferRules, progress *Progress) map[string]map[string]imapMessageInfo {
	res := map[string]map[string]imapMessageInfo{}

	for rule := range rules.iterateActiveRules() {
		if progress.shouldStop() {
			break
		}

		mailboxName := rule.SourceMailbox.Name
		var mailbox *imap.MailboxStatus
		progress.callWrap(func() error {
			var err error
			mailbox, err = p.selectIn(mailboxName)
			return err
		})
		if mailbox.Messages == 0 {
			continue
		}

		messagesInfo := p.loadMessagesInfo(rule, progress, mailbox.UidValidity)
		res[rule.SourceMailbox.Name] = messagesInfo
		progress.updateCount(rule.SourceMailbox.Name, uint(len(messagesInfo)))
	}

	return res
}

func (p *IMAPProvider) loadMessagesInfo(rule *Rule, progress *Progress, uidValidity uint32) map[string]imapMessageInfo {
	messagesInfo := map[string]imapMessageInfo{}

	pageStart := uint32(1)
	pageEnd := imapPageSize
	for {
		if progress.shouldStop() {
			break
		}

		seqSet := &imap.SeqSet{}
		seqSet.AddRange(pageStart, pageEnd)

		items := []imap.FetchItem{imap.FetchUid, imap.FetchRFC822Size}
		if rule.HasTimeLimit() {
			items = append(items, imap.FetchEnvelope)
		}

		pageMsgCount := uint32(0)
		processMessageCallback := func(imapMessage *imap.Message) {
			pageMsgCount++
			if rule.HasTimeLimit() {
				t := imapMessage.Envelope.Date.Unix()
				if t != 0 && !rule.isTimeInRange(t) {
					log.WithField("uid", imapMessage.Uid).Debug("Message skipped due to time")
					return
				}
			}
			id := fmt.Sprintf("%s_%d:%d", rule.SourceMailbox.Name, uidValidity, imapMessage.Uid)
			messagesInfo[id] = imapMessageInfo{
				id:   id,
				uid:  imapMessage.Uid,
				size: imapMessage.Size,
			}
			progress.addMessage(id, rule)
		}

		progress.callWrap(func() error {
			return p.fetch(seqSet, items, processMessageCallback)
		})

		if pageMsgCount < imapPageSize {
			break
		}

		pageStart = pageEnd
		pageEnd += imapPageSize
	}

	return messagesInfo
}

func (p *IMAPProvider) transferTo(rule *Rule, messagesInfo map[string]imapMessageInfo, progress *Progress, ch chan<- Message) {
	progress.callWrap(func() error {
		_, err := p.selectIn(rule.SourceMailbox.Name)
		return err
	})

	seqSet := &imap.SeqSet{}
	seqSetSize := uint32(0)
	uidToID := map[uint32]string{}

	for _, messageInfo := range messagesInfo {
		if progress.shouldStop() {
			break
		}

		if seqSetSize != 0 && (seqSetSize+messageInfo.size) > imapMaxFetchSize {
			log.WithField("mailbox", rule.SourceMailbox.Name).WithField("seq", seqSet).WithField("size", seqSetSize).Debug("Fetching messages")

			p.exportMessages(rule, progress, ch, seqSet, uidToID)

			seqSet = &imap.SeqSet{}
			seqSetSize = 0
			uidToID = map[uint32]string{}
		}

		seqSet.AddNum(messageInfo.uid)
		seqSetSize += messageInfo.size
		uidToID[messageInfo.uid] = messageInfo.id
	}
}

func (p *IMAPProvider) exportMessages(rule *Rule, progress *Progress, ch chan<- Message, seqSet *imap.SeqSet, uidToID map[uint32]string) {
	section := &imap.BodySectionName{}
	items := []imap.FetchItem{imap.FetchUid, imap.FetchFlags, section.FetchItem()}

	processMessageCallback := func(imapMessage *imap.Message) {
		id, ok := uidToID[imapMessage.Uid]

		// Sometimes, server sends not requested messages.
		if !ok {
			log.WithField("uid", imapMessage.Uid).Warning("Message skipped: not requested")
			return
		}

		// Sometimes, server sends message twice, once with body and once without it.
		bodyReader := imapMessage.GetBody(section)
		if bodyReader == nil {
			log.WithField("uid", imapMessage.Uid).Warning("Message skipped: no body")
			return
		}

		body, err := ioutil.ReadAll(bodyReader)
		progress.messageExported(id, body, err)
		if err == nil {
			msg := p.exportMessage(rule, id, imapMessage, body)
			ch <- msg
		}
	}

	progress.callWrap(func() error {
		return p.uidFetch(seqSet, items, processMessageCallback)
	})
}

func (p *IMAPProvider) exportMessage(rule *Rule, id string, imapMessage *imap.Message, body []byte) Message {
	unread := true
	for _, flag := range imapMessage.Flags {
		if flag == imap.SeenFlag {
			unread = false
		}
	}

	return Message{
		ID:      id,
		Unread:  unread,
		Body:    body,
		Source:  rule.SourceMailbox,
		Targets: rule.TargetMailboxes,
	}
}