2 files changed, 17 insertions, 30 deletions

diff --git a/worker/maildir/message.go b/worker/maildir/message.go
index 5c6c9307..55e2613b 100644
--- a/worker/maildir/message.go
+++ b/worker/maildir/message.go
@@ -53,43 +53,29 @@ func (m Message) SetFlags(flags []maildir.Flag) error {
 	return m.dir.SetFlags(m.key, flags)
 }
 
-// MarkReplied either adds or removes the maildir.FlagReplied flag from the
-// message.
-func (m Message) MarkReplied(answered bool) error {
+// SetOneFlag enables or disables a single message flag on the message.
+func (m Message) SetOneFlag(flag maildir.Flag, enable bool) error {
 	flags, err := m.Flags()
 	if err != nil {
 		return fmt.Errorf("could not read previous flags: %v", err)
 	}
-	if answered {
-		flags = append(flags, maildir.FlagReplied)
+	if enable {
+		flags = append(flags, flag)
 		return m.SetFlags(flags)
 	}
 	var newFlags []maildir.Flag
-	for _, flag := range flags {
-		if flag != maildir.FlagReplied {
-			newFlags = append(newFlags, flag)
+	for _, oldFlag := range flags {
+		if oldFlag != flag {
+			newFlags = append(newFlags, oldFlag)
 		}
 	}
 	return m.SetFlags(newFlags)
 }
 
-// MarkRead either adds or removes the maildir.FlagSeen flag from the message.
-func (m Message) MarkRead(seen bool) error {
-	flags, err := m.Flags()
-	if err != nil {
-		return fmt.Errorf("could not read previous flags: %v", err)
-	}
-	if seen {
-		flags = append(flags, maildir.FlagSeen)
-		return m.SetFlags(flags)
-	}
-	var newFlags []maildir.Flag
-	for _, flag := range flags {
-		if flag != maildir.FlagSeen {
-			newFlags = append(newFlags, flag)
-		}
-	}
-	return m.SetFlags(newFlags)
+// MarkReplied either adds or removes the maildir.FlagReplied flag from the
+// message.
+func (m Message) MarkReplied(answered bool) error {
+	return m.SetOneFlag(maildir.FlagReplied, answered)
 }
 
 // Remove deletes the email immediately.
diff --git a/worker/maildir/worker.go b/worker/maildir/worker.go
index f14672e5..467f2aa0 100644
--- a/worker/maildir/worker.go
+++ b/worker/maildir/worker.go
@@ -193,8 +193,8 @@ func (w *Worker) handleMessage(msg types.WorkerMessage) error {
 		return w.handleFetchFullMessages(msg)
 	case *types.DeleteMessages:
 		return w.handleDeleteMessages(msg)
-	case *types.ReadMessages:
-		return w.handleReadMessages(msg)
+	case *types.FlagMessages:
+		return w.handleFlagMessages(msg)
 	case *types.AnsweredMessages:
 		return w.handleAnsweredMessages(msg)
 	case *types.CopyMessages:
@@ -473,7 +473,7 @@ func (w *Worker) handleAnsweredMessages(msg *types.AnsweredMessages) error {
 	return nil
 }
 
-func (w *Worker) handleReadMessages(msg *types.ReadMessages) error {
+func (w *Worker) handleFlagMessages(msg *types.FlagMessages) error {
 	for _, uid := range msg.Uids {
 		m, err := w.c.Message(*w.selected, uid)
 		if err != nil {
@@ -481,8 +481,9 @@ func (w *Worker) handleReadMessages(msg *types.ReadMessages) error {
 			w.err(msg, err)
 			continue
 		}
-		if err := m.MarkRead(msg.Read); err != nil {
-			w.worker.Logger.Printf("could not mark message as read: %v", err)
+		flag := flagToMaildir[msg.Flag]
+		if err := m.SetOneFlag(flag, msg.Enable); err != nil {
+			w.worker.Logger.Printf("could change flag %v to %v on message: %v", flag, msg.Enable, err)
 			w.err(msg, err)
 			continue
 		}