kittemail/mailbox.go

package main

import (
	"io"
	"sync"
	"time"

	"encoding/json"

	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/backend/backendutil"
	"github.com/google/uuid"
)

var Delimiter = "/"

type Mailbox struct {
	Subscribed bool
	id         uuid.UUID
	name       string
	user       *User
}

func (mbox *Mailbox) Name() string {
	return mbox.name
}

func (mbox *Mailbox) Info() (*imap.MailboxInfo, error) {
	info := &imap.MailboxInfo{
		Delimiter: Delimiter,
		Name:      mbox.name,
	}
	return info, nil
}

func (mbox *Mailbox) uidNext() (uint32, error) {
	messages, err := Database.DB.Query("SELECT uid FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return 0, err
	}

	defer func() {
		err = messages.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	var uid uint32
	//goland:noinspection GoDfaErrorMayBeNotNil
	for messages.Next() {
		var messageID uint32
		err = messages.Scan(&messageID)
		if err != nil {
			return 0, err
		}
		if messageID > uid {
			uid = messageID
		}
	}

	uid++
	return uid, nil
}

func (mbox *Mailbox) flags() ([]string, error) {
	flagsMap := make(map[string]struct{})
	flagSlices, err := Database.DB.Query("SELECT flags FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return nil, err
	}

	defer func() {
		err = flagSlices.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	for flagSlices.Next() {
		var flagsRaw string
		err := flagSlices.Scan(&flagsRaw)
		if err != nil {
			return nil, err
		}

		var flags []string
		err = json.Unmarshal([]byte(flagsRaw), &flags)
		if err != nil {
			return nil, err
		}

		for _, f := range flags {
			// It's still unique, even if it is reassigned
			// Reassignment doesn't really matter since it's a 0 byte struct, and it would be less efficient to check if it's already in the map
			flagsMap[f] = struct{}{}
		}
	}

	var flags []string
	for f := range flagsMap {
		flags = append(flags, f)
	}
	return flags, nil
}

func (mbox *Mailbox) unseenSeqNum() (seqNum uint32, err error) {
	messages, err := Database.DB.Query("SELECT flags FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return 0, err
	}

	defer func() {
		err = messages.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	for messages.Next() {
		seqNum++

		var flagsRaw string
		err = messages.Scan(&flagsRaw)
		if err != nil {
			return 0, err
		}

		var flags []string
		err = json.Unmarshal([]byte(flagsRaw), &flags)
		if err != nil {
			return 0, err
		}

		seen := false
		for _, flag := range flags {
			if flag == imap.SeenFlag {
				seen = true
				break
			}
		}

		if !seen {
			return seqNum, nil
		}
	}

	return 0, nil
}

func (mbox *Mailbox) Status(items []imap.StatusItem) (status *imap.MailboxStatus, err error) {
	status = imap.NewMailboxStatus(mbox.name, items)

	status.Flags, err = mbox.flags()
	if err != nil {
		return
	}
	status.PermanentFlags = []string{"\\*"}
	status.UnseenSeqNum, err = mbox.unseenSeqNum()
	if err != nil {
		return
	}

	var messageCount int
	var hasCounted sync.Once

	var hasUidNext sync.Once
	var uidNext uint32

	var hasUnseen sync.Once
	var unseen uint32

	for _, name := range items {
		switch name {
		case imap.StatusMessages:
			hasCounted.Do(func() {
				err = Database.DB.QueryRow("SELECT COUNT(*) FROM messages WHERE mailbox = $1", mbox.id[:]).Scan(&messageCount)
			})
			if err != nil {
				return
			}
			status.Messages = uint32(messageCount)
		case imap.StatusUidNext:
			hasUidNext.Do(func() {
				uidNext, err = mbox.uidNext()
			})
			if err != nil {
				return
			}
			status.UidNext = uidNext
		case imap.StatusUidValidity:
			status.UidValidity = 1
		case imap.StatusRecent:
			status.Recent = 0 // TODO: implement most recent message
		case imap.StatusUnseen:
			hasUnseen.Do(func() {
				unseen, err = mbox.unseenSeqNum()
			})
			if err != nil {
				return
			}
			status.Unseen = unseen
		}
	}

	return
}

func (mbox *Mailbox) SetSubscribed(subscribed bool) error {
	_, err := Database.DB.Exec("UPDATE mailboxes SET subscribed = $1 WHERE id = $2", subscribed, mbox.id[:])
	if err != nil {
		return err
	}

	mbox.Subscribed = subscribed
	return nil
}

func (mbox *Mailbox) Check() error {
	return nil
}

func (mbox *Mailbox) ListMessages(useUid bool, seqSet *imap.SeqSet, items []imap.FetchItem, ch chan<- *imap.Message) error {
	defer close(ch)

	messages, err := Database.DB.Query("SELECT id, uid, created, bodySize, flags, owner FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return err
	}

	defer func() {
		err := messages.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	var seqNum uint32
	for messages.Next() {
		seqNum++

		var uid, size uint32
		var date time.Time
		var flagsRaw string
		var ownerRaw []byte
		var idRaw []byte
		err := messages.Scan(&idRaw, &uid, &date, &size, &flagsRaw, &ownerRaw)
		if err != nil {
			// Skip any emails that can't be read
			continue
		}

		msg, err := LoadRawMessage(idRaw, uid, date, size, flagsRaw, ownerRaw, mbox)
		if err != nil {
			// Skip any emails that fail to load from disk
			continue
		}

		var id uint32
		if useUid {
			id = msg.Uid
		} else {
			id = seqNum
		}
		if !seqSet.Contains(id) {
			// Skip any emails that aren't in the sequence set
			continue
		}

		m, err := msg.Fetch(seqNum, items)
		if err != nil {
			// Skip any emails that can't be represented as an IMAP message
			continue
		}

		ch <- m
	}

	return nil
}

func (mbox *Mailbox) SearchMessages(useUid bool, criteria *imap.SearchCriteria) ([]uint32, error) {
	messages, err := Database.DB.Query("SELECT id, uid, created, bodySize, flags, owner FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return nil, err
	}

	var ids []uint32

	defer func() {
		err := messages.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	// TODO: Make this thing multithreaded

	var seqNum uint32
	for messages.Next() {
		seqNum++

		var uid, size uint32
		var date time.Time
		var flagsRaw string
		var ownerRaw []byte
		var idRaw []byte
		err := messages.Scan(&idRaw, &uid, &date, &size, &flagsRaw, &ownerRaw)
		if err != nil {
			// Skip any emails that can't be read
			continue
		}

		msg, err := LoadRawMessage(idRaw, uid, date, size, flagsRaw, ownerRaw, mbox)
		if err != nil {
			// Skip any emails that fail to load from disk
			continue
		}

		ok, err := msg.Match(seqNum, criteria)
		if err != nil || !ok {
			closeErr := msg.Close()
			if closeErr != nil {
				log("Failed to close message body: "+closeErr.Error()+", resource leaks may occur", 1)
			}

			continue
		}

		var id uint32
		if useUid {
			id = msg.Uid
		} else {
			id = seqNum
		}

		closeErr := msg.Close()
		if closeErr != nil {
			log("Failed to close message body: "+closeErr.Error()+", resource leaks may occur", 1)
		}

		ids = append(ids, id)
	}

	return ids, nil
}

func (mbox *Mailbox) CreateMessage(flags []string, date time.Time, body imap.Literal) error {
	if date.IsZero() {
		date = time.Now()
	}

	b, err := io.ReadAll(body)
	if err != nil {
		return err
	}

	messageID := uuid.New()

	flagsRaw, err := json.Marshal(flags)
	if err != nil {
		return err
	}

	uid, err := mbox.uidNext()
	if err != nil {
		return err
	}

	_, err = Database.DB.Exec("INSERT INTO messages (mailbox, id, uid, created, bodySize, flags, owner) VALUES ($1, $2, $3, $4, $5, $6, $7)", mbox.id[:], messageID[:], uid, date, len(b), string(flagsRaw), mbox.user.sub[:])
	if err != nil {
		return err
	}

	return nil
}

func (mbox *Mailbox) UpdateMessagesFlags(useUid bool, seqset *imap.SeqSet, op imap.FlagsOp, flags []string) error {
	messages, err := Database.DB.Query("SELECT uid, flags FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return err
	}

	defer func() {
		err := messages.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	var seqNum uint32
	for messages.Next() {
		seqNum++

		var uid uint32
		var flagsRaw string
		err := messages.Scan(&uid, &flagsRaw)
		if err != nil {
			// Skip any emails that can't be read
			continue
		}

		var msgFlags []string
		err = json.Unmarshal([]byte(flagsRaw), &msgFlags)
		if err != nil {
			// Skip any emails that fail to unmarshal
			continue
		}

		var id uint32
		if useUid {
			id = uid
		} else {
			id = seqNum
		}
		if !seqset.Contains(id) {
			continue
		}

		flags = backendutil.UpdateFlags(flags, op, flags)

		flagsBytes, err := json.Marshal(flags)
		if err != nil {
			// Skip any emails that fail to marshal
			continue
		}

		_, err = Database.DB.Exec("UPDATE messages SET flags = $1 WHERE mailbox = $2 AND uid = $3", string(flagsBytes), mbox.id[:], uid)
		if err != nil {
			// Skip any emails that fail to update
			continue
		}
	}

	return nil
}

func (mbox *Mailbox) CopyMessages(useUid bool, seqset *imap.SeqSet, destName string) error {
	dest, err := mbox.user.GetRawMailbox(destName)
	if err != nil {
		return err
	}

	messages, err := Database.DB.Query("SELECT id, uid, created, bodySize, flags, owner FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return err
	}

	defer func() {
		err := messages.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	var seqNum uint32
	for messages.Next() {
		seqNum++

		var uid, size uint32
		var date time.Time
		var flagsRaw string
		var idRaw, ownerRaw []byte
		err := messages.Scan(&idRaw, &uid, &date, &size, &flagsRaw, &ownerRaw)
		if err != nil {
			// Skip any emails that can't be read
			continue
		}

		msg, err := LoadRawMessage(idRaw, uid, date, size, flagsRaw, ownerRaw, mbox)
		if err != nil {
			// Skip any emails that fail to load from disk
			continue
		}

		var id uint32
		if useUid {
			id = msg.Uid
		} else {
			id = seqNum
		}
		if !seqset.Contains(id) {
			continue
		}

		msgCopy := *msg
		msgCopy.Uid, err = dest.uidNext()
		if err != nil {
			closeErr := msg.Body.Close()
			if closeErr != nil {
				log("Failed to close message body: "+closeErr.Error()+", resource leaks may occur", 1)
			}
			return err
		}

		err = dest.CreateMessage(msgCopy.Flags, msgCopy.Date, msgCopy.getBody())
		if err != nil {
			closeErr := msg.Body.Close()
			if closeErr != nil {
				log("Failed to close message body: "+closeErr.Error()+", resource leaks may occur", 1)
			}
			return err
		}
	}

	return nil
}

func (mbox *Mailbox) Expunge() error {
	messages, err := Database.DB.Query("SELECT id FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return err
	}

	defer func() {
		err := messages.Close()
		if err != nil {
			log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
		}
	}()

	for messages.Next() {
		var idRaw []byte
		err := messages.Scan(&idRaw)
		if err != nil {
			return err
		}

		id, err := uuid.FromBytes(idRaw)
		if err != nil {
			return err
		}

		err = DeleteFile(id.String(), mbox.user.sub)
		if err != nil {
			return err
		}
	}

	_, err = Database.DB.Exec("DELETE FROM messages WHERE mailbox = $1", mbox.id[:])
	if err != nil {
		return err
	}

	return nil
}