Redirect to edit draft after saving message

2 files changed, 26 insertions, 13 deletions

diff --git a/plugins/base/imap.go b/plugins/base/imap.go
index c7718e9..f15105c 100644
--- a/plugins/base/imap.go
+++ b/plugins/base/imap.go
@@ -9,6 +9,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+	nettextproto "net/textproto"
 
 	"github.com/dustin/go-humanize"
 	"github.com/emersion/go-imap"
@@ -570,20 +571,20 @@ func markMessageAnswered(conn *imapclient.Client, mboxName string, uid uint32) e
 	return conn.UidStore(seqSet, item, flags, nil)
 }
 
-func appendMessage(c *imapclient.Client, msg *OutgoingMessage, mboxType mailboxType) (saved bool, err error) {
+func appendMessage(c *imapclient.Client, msg *OutgoingMessage, mboxType mailboxType) (*MailboxInfo, uint32, error) {
 	mbox, err := getMailboxByType(c, mboxType)
 	if err != nil {
-		return false, err
+		return nil, 0, err
 	}
 	if mbox == nil {
-		return false, nil
+		return nil, 0, fmt.Errorf("Unable to resolve mailbox")
 	}
 
 	// IMAP needs to know in advance the final size of the message, so
 	// there's no way around storing it in a buffer here.
 	var buf bytes.Buffer
 	if err := msg.WriteTo(&buf); err != nil {
-		return false, err
+		return nil, 0, err
 	}
 
 	flags := []string{imap.SeenFlag}
@@ -591,10 +592,20 @@ func appendMessage(c *imapclient.Client, msg *OutgoingMessage, mboxType mailboxT
 		flags = append(flags, imap.DraftFlag)
 	}
 	if err := c.Append(mbox.Name, flags, time.Now(), &buf); err != nil {
-		return false, err
+		return nil, 0, err
+	}
+	criteria := &imap.SearchCriteria{
+		Header: make(nettextproto.MIMEHeader),
+	}
+	criteria.Header.Add("Message-Id", msg.MessageID)
+	if uids, err := c.UidSearch(criteria); err != nil {
+		return nil, 0, err
+	} else {
+		if len(uids) != 1 {
+			panic(fmt.Errorf("Duplicate message ID"))
+		}
+		return mbox, uids[0], nil
 	}
-
-	return true, nil
 }
 
 func deleteMessage(c *imapclient.Client, mboxName string, uid uint32) error {
diff --git a/plugins/base/routes.go b/plugins/base/routes.go
index 6ffdcf6..1f3b8b4 100644
--- a/plugins/base/routes.go
+++ b/plugins/base/routes.go
@@ -507,7 +507,7 @@ func submitCompose(ctx *alps.Context, msg *OutgoingMessage, options *composeOpti
 	}
 
 	err = ctx.Session.DoIMAP(func(c *imapclient.Client) error {
-		if _, err := appendMessage(c, msg, mailboxSent); err != nil {
+		if _, _, err := appendMessage(c, msg, mailboxSent); err != nil {
 			return err
 		}
 		if draft := options.Draft; draft != nil {
@@ -620,14 +620,15 @@ func handleCompose(ctx *alps.Context, msg *OutgoingMessage, options *composeOpti
 		}
 
 		if saveAsDraft {
+			var (
+				drafts *MailboxInfo
+				uid    uint32
+			)
 			err = ctx.Session.DoIMAP(func(c *imapclient.Client) error {
-				copied, err := appendMessage(c, msg, mailboxDrafts)
+				drafts, uid, err = appendMessage(c, msg, mailboxDrafts)
 				if err != nil {
 					return err
 				}
-				if !copied {
-					return fmt.Errorf("no Draft mailbox found")
-				}
 				if draft := options.Draft; draft != nil {
 					if err := deleteMessage(c, draft.Mailbox, draft.Uid); err != nil {
 						return err
@@ -638,7 +639,8 @@ func handleCompose(ctx *alps.Context, msg *OutgoingMessage, options *composeOpti
 			if err != nil {
 				return fmt.Errorf("failed to save message to Draft mailbox: %v", err)
 			}
-			return ctx.Redirect(http.StatusFound, "/mailbox/INBOX")
+			return ctx.Redirect(http.StatusFound, fmt.Sprintf(
+				"/message/%s/%d/edit?part=1", drafts.Name, uid))
 		} else {
 			return submitCompose(ctx, msg, options)
 		}