messageinfo: report message sizes

Report sizes of the message across all backends.

Signed-off-by: Koni Marti <koni.marti@gmail.com>
Acked-by: Robin Jarry <robin@jarry.cc>

5 files changed, 66 insertions, 4 deletions

diff --git a/worker/imap/fetch.go b/worker/imap/fetch.go
index a8245784..f209d7f0 100644
--- a/worker/imap/fetch.go
+++ b/worker/imap/fetch.go
@@ -42,6 +42,7 @@ func (imapw *IMAPWorker) handleFetchMessageHeaders(
 		imap.FetchInternalDate,
 		imap.FetchFlags,
 		imap.FetchUid,
+		imap.FetchRFC822Size,
 		section.FetchItem(),
 	}
 	imapw.handleFetchMessages(msg, toFetch, items,
@@ -66,6 +67,7 @@ func (imapw *IMAPWorker) handleFetchMessageHeaders(
 				InternalDate:  _msg.InternalDate,
 				RFC822Headers: header,
 				Refs:          parse.MsgIDList(header, "references"),
+				Size:          _msg.Size,
 				Uid:           _msg.Uid,
 			}
 			imapw.worker.PostMessage(&types.MessageInfo{
diff --git a/worker/lib/parse.go b/worker/lib/parse.go
index 3baaa459..ac86292f 100644
--- a/worker/lib/parse.go
+++ b/worker/lib/parse.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"regexp"
 	"strings"
 	"time"
@@ -381,3 +382,12 @@ func ReadMessage(r io.Reader) (*message.Entity, error) {
 	}
 	return entity, nil
 }
+
+// FileSize returns the size of the file specified by name
+func FileSize(name string) (uint32, error) {
+	fileInfo, err := os.Stat(name)
+	if err != nil {
+		return 0, fmt.Errorf("failed to obtain fileinfo: %w", err)
+	}
+	return uint32(fileInfo.Size()), nil
+}
diff --git a/worker/maildir/message.go b/worker/maildir/message.go
index a941624a..0b653fa6 100644
--- a/worker/maildir/message.go
+++ b/worker/maildir/message.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/emersion/go-maildir"
 
+	"git.sr.ht/~rjarry/aerc/log"
 	"git.sr.ht/~rjarry/aerc/models"
 	"git.sr.ht/~rjarry/aerc/worker/lib"
 )
@@ -73,7 +74,19 @@ func (m Message) Remove() error {
 
 // MessageInfo populates a models.MessageInfo struct for the message.
 func (m Message) MessageInfo() (*models.MessageInfo, error) {
-	return lib.MessageInfo(m)
+	info, err := lib.MessageInfo(m)
+	if err != nil {
+		return nil, err
+	}
+	// if size retrieval fails, just return info and log error
+	if name, err := m.dir.Filename(m.key); err != nil {
+		log.Errorf("failed to obtain filename: %v", err)
+	} else {
+		if info.Size, err = lib.FileSize(name); err != nil {
+			log.Errorf("failed to obtain file size: %v", err)
+		}
+	}
+	return info, nil
 }
 
 // MessageHeaders populates a models.MessageInfo struct for the message with
diff --git a/worker/mbox/worker.go b/worker/mbox/worker.go
index 01a1425a..fbdbec36 100644
--- a/worker/mbox/worker.go
+++ b/worker/mbox/worker.go
@@ -15,6 +15,7 @@ import (
 	"git.sr.ht/~rjarry/aerc/worker/handlers"
 	"git.sr.ht/~rjarry/aerc/worker/lib"
 	"git.sr.ht/~rjarry/aerc/worker/types"
+	"github.com/miolini/datacounter"
 )
 
 func init() {
@@ -155,7 +156,7 @@ func (w *mboxWorker) handleMessage(msg types.WorkerMessage) error {
 				reterr = err
 				break
 			}
-			msgInfo, err := lib.MessageInfo(m)
+			msgInfo, err := messageInfo(m, true)
 			if err != nil {
 				w.worker.PostMessageInfoError(msg, uid, err)
 				break
@@ -404,13 +405,19 @@ func sortUids(folder *container, uids []uint32,
 	criteria []*types.SortCriterion,
 ) ([]uint32, error) {
 	var infos []*models.MessageInfo
+	needSize := false
+	for _, item := range criteria {
+		if item.Field == types.SortSize {
+			needSize = true
+		}
+	}
 	for _, uid := range uids {
 		m, err := folder.Message(uid)
 		if err != nil {
 			log.Errorf("could not get message %v", err)
 			continue
 		}
-		info, err := lib.MessageInfo(m)
+		info, err := messageInfo(m, needSize)
 		if err != nil {
 			log.Errorf("could not get message info %v", err)
 			continue
@@ -419,3 +426,20 @@ func sortUids(folder *container, uids []uint32,
 	}
 	return lib.Sort(infos, criteria)
 }
+
+func messageInfo(m lib.RawMessage, needSize bool) (*models.MessageInfo, error) {
+	info, err := lib.MessageInfo(m)
+	if err != nil {
+		return nil, err
+	}
+	if needSize {
+		if r, err := m.NewReader(); err == nil {
+			var buf bytes.Buffer
+			ctr := datacounter.NewWriterCounter(&buf)
+			if _, err := io.Copy(ctr, r); err == nil {
+				info.Size = uint32(ctr.Count())
+			}
+		}
+	}
+	return info, nil
+}
diff --git a/worker/notmuch/message.go b/worker/notmuch/message.go
index 28731d3d..b6d7f54c 100644
--- a/worker/notmuch/message.go
+++ b/worker/notmuch/message.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/emersion/go-maildir"
 
+	"git.sr.ht/~rjarry/aerc/log"
 	"git.sr.ht/~rjarry/aerc/models"
 	"git.sr.ht/~rjarry/aerc/worker/lib"
 	notmuch "git.sr.ht/~rjarry/aerc/worker/notmuch/lib"
@@ -34,7 +35,19 @@ func (m *Message) NewReader() (io.ReadCloser, error) {
 
 // MessageInfo populates a models.MessageInfo struct for the message.
 func (m *Message) MessageInfo() (*models.MessageInfo, error) {
-	return lib.MessageInfo(m)
+	info, err := lib.MessageInfo(m)
+	if err != nil {
+		return nil, err
+	}
+	// if size retrieval fails, just return info and log error
+	if name, err := m.Filename(); err != nil {
+		log.Errorf("failed to obtain filename: %v", err)
+	} else {
+		if info.Size, err = lib.FileSize(name); err != nil {
+			log.Errorf("failed to obtain file size: %v", err)
+		}
+	}
+	return info, nil
 }
 
 // NewBodyPartReader creates a new io.Reader for the requested body part(s) of