retry: add option to store outgoing messages in db

Author: tulir

client.go

@@ -150,6 +150,10 @@ type Client struct {
 	// PreRetryCallback is called before a retry receipt is accepted.
 	// If it returns false, the accepting will be cancelled and the retry receipt will be ignored.
 	PreRetryCallback func(receipt *events.Receipt, id types.MessageID, retryCount int, msg *waE2E.Message) bool
+	// Should whatsmeow store recently sent messages in the database so that retry receipts can be accepted
+	// even if the process is restarted? If false, only the in-memory cache and GetMessageForRetry will be used.
+	UseRetryMessageStore bool
+	lastRetryStoreClear  time.Time
 
 	// PrePairCallback is called before pairing is completed. If it returns false, the pairing will be cancelled and
 	// the client will disconnect.

internals.go

@@ -47,7 +47,34 @@ func (rm RecentMessage) IsEmpty() bool {
 	return rm.wa == nil && rm.fb == nil
 }
 
-func (cli *Client) addRecentMessage(to types.JID, id types.MessageID, wa *waE2E.Message, fb *waMsgApplication.MessageApplication) {
+func (cli *Client) addRecentMessage(ctx context.Context, to types.JID, id types.MessageID, wa *waE2E.Message, fb *waMsgApplication.MessageApplication) error {
+	if cli.UseRetryMessageStore {
+		var buf []byte
+		var format string
+		var err error
+		if wa != nil {
+			buf, err = proto.Marshal(wa)
+			format = "wa"
+		} else if fb != nil {
+			buf, err = proto.Marshal(fb)
+			format = "fb"
+		}
+		if err != nil {
+			return fmt.Errorf("failed to marshal message for retry store: %w", err)
+		}
+		if buf != nil {
+			err = cli.Store.EventBuffer.AddOutgoingEvent(ctx, to, id, format, buf)
+			if err != nil {
+				return fmt.Errorf("failed to add message to retry store: %w", err)
+			}
+			if time.Since(cli.lastRetryStoreClear) > 12*time.Hour {
+				err = cli.Store.EventBuffer.DeleteOldOutgoingEvents(ctx)
+				if err != nil {
+					return fmt.Errorf("failed to clear old messages from retry store: %w", err)
+				}
+			}
+		}
+	}
 	cli.recentMessagesLock.Lock()
 	key := recentMessageKey{to, id}
 	if cli.recentMessagesList[cli.recentMessagesPtr].ID != "" {
@@ -60,6 +87,7 @@ func (cli *Client) addRecentMessage(to types.JID, id types.MessageID, wa *waE2E.
 		cli.recentMessagesPtr = 0
 	}
 	cli.recentMessagesLock.Unlock()
+	return nil
 }
 
 func (cli *Client) getRecentMessage(to types.JID, id types.MessageID) RecentMessage {
@@ -91,6 +119,13 @@ func (cli *Client) getMessageForRetry(ctx context.Context, receipt *events.Recei
 			return &msg, nil
 		}
 	}
+	if cli.UseRetryMessageStore {
+		format, buf, err := cli.Store.EventBuffer.GetOutgoingEvent(ctx, receipt.Chat, altChat, messageID)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get message from retry store: %w", err)
+		}
+		return parseRecentMessage(format, buf)
+	}
 	waMsg := cli.GetMessageForRetry(receipt.Sender, receipt.Chat, messageID)
 	if waMsg != nil {
 		cli.Log.Debugf("Found message in GetMessageForRetry to accept retry receipt for %s/%s from %s", receipt.Chat, messageID, receipt.Sender)
@@ -99,6 +134,25 @@ func (cli *Client) getMessageForRetry(ctx context.Context, receipt *events.Recei
 	return nil, nil
 }
 
+func parseRecentMessage(format string, buf []byte) (*RecentMessage, error) {
+	var rm RecentMessage
+	var err error
+	switch format {
+	case "wa":
+		rm.wa = &waE2E.Message{}
+		err = proto.Unmarshal(buf, rm.wa)
+	case "fb":
+		rm.fb = &waMsgApplication.MessageApplication{}
+		err = proto.Unmarshal(buf, rm.fb)
+	default:
+		err = fmt.Errorf("unknown format in retry store: %s", format)
+	}
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal payload in retry store: %w", err)
+	}
+	return &rm, nil
+}
+
 const recreateSessionTimeout = 1 * time.Hour
 
 func (cli *Client) shouldRecreateSession(ctx context.Context, retryCount int, jid types.JID) (reason string, recreate bool) {

retry.go

@@ -373,10 +373,12 @@ func (cli *Client) SendMessage(ctx context.Context, to types.JID, message *waE2E
 	resp.DebugTimings.Queue = time.Since(start)
 	defer cli.messageSendLock.Unlock()
 
-	respChan := cli.waitResponse(req.ID)
 	// Peer message retries aren't implemented yet
 	if !req.Peer {
-		cli.addRecentMessage(to, req.ID, message, nil)
+		err = cli.addRecentMessage(ctx, to, req.ID, message, nil)
+		if err != nil {
+			return
+		}
 	}
 
 	if message.GetMessageContextInfo().GetMessageSecret() != nil {
@@ -387,6 +389,8 @@ func (cli *Client) SendMessage(ctx context.Context, to types.JID, message *waE2E
 			cli.Log.Debugf("Stored message secret key for outgoing message %s", req.ID)
 		}
 	}
+
+	respChan := cli.waitResponse(req.ID)
 	var phash string
 	var data []byte
 	switch to.Server {

send.go

@@ -137,10 +137,13 @@ func (cli *Client) SendFBMessage(
 	resp.DebugTimings.Queue = time.Since(start)
 	defer cli.messageSendLock.Unlock()
 
-	respChan := cli.waitResponse(req.ID)
 	if !req.Peer {
-		cli.addRecentMessage(to, req.ID, nil, messageAppProto)
+		err = cli.addRecentMessage(ctx, to, req.ID, nil, messageAppProto)
+		if err != nil {
+			return
+		}
 	}
+	respChan := cli.waitResponse(req.ID)
 	var phash string
 	var data []byte
 	switch to.Server {

sendfb.go

@@ -275,3 +275,15 @@ func (n *NoopStore) PutManyLIDMappings(ctx context.Context, mappings []LIDMappin
 func (n *NoopStore) PutLIDMapping(ctx context.Context, lid types.JID, jid types.JID) error {
 	return n.Error
 }
+
+func (n *NoopStore) DeleteOldOutgoingEvents(ctx context.Context) error {
+	return nil
+}
+
+func (n *NoopStore) GetOutgoingEvent(ctx context.Context, chatJID, altChatJID types.JID, id types.MessageID) (string, []byte, error) {
+	return "", nil, nil
+}
+
+func (n *NoopStore) AddOutgoingEvent(ctx context.Context, chatJID types.JID, id types.MessageID, format string, plaintext []byte) error {
+	return nil
+}

store/noop.go

@@ -1028,3 +1028,33 @@ func (s *SQLStore) DeleteOldBufferedHashes(ctx context.Context) error {
 	_, err := s.db.Exec(ctx, deleteOldBufferedHashesQuery, time.Now().Add(-14*24*time.Hour).UnixMilli())
 	return err
 }
+
+const (
+	getOutgoingEventQuery = `
+		SELECT format, plaintext FROM whatsmeow_retry_buffer WHERE our_jid=$1 AND (chat_jid=$2 OR chat_jid=$3) AND message_id=$4
+	`
+	addOutgoingEventQuery = `
+		INSERT INTO whatsmeow_retry_buffer (our_jid, chat_jid, message_id, format, plaintext, timestamp)
+		VALUES ($1, $2, $3, $4, $5, $6)
+		ON CONFLICT (our_jid, chat_jid, message_id) DO UPDATE
+			SET format=excluded.format, plaintext=excluded.plaintext, timestamp=excluded.timestamp
+	`
+	deleteOldOutgoingEventsQuery = `
+		DELETE FROM whatsmeow_retry_buffer WHERE our_jid=$1 AND timestamp < $2
+	`
+)
+
+func (s *SQLStore) GetOutgoingEvent(ctx context.Context, chatJID, altChatJID types.JID, id types.MessageID) (format string, result []byte, err error) {
+	err = s.db.QueryRow(ctx, getOutgoingEventQuery, s.JID, chatJID, altChatJID, id).Scan(&format, &result)
+	return
+}
+
+func (s *SQLStore) AddOutgoingEvent(ctx context.Context, chatJID types.JID, id types.MessageID, format string, plaintext []byte) error {
+	_, err := s.db.Exec(ctx, addOutgoingEventQuery, s.JID, chatJID, id, format, plaintext, time.Now().UnixMilli())
+	return err
+}
+
+func (s *SQLStore) DeleteOldOutgoingEvents(ctx context.Context) error {
+	_, err := s.db.Exec(ctx, deleteOldOutgoingEventsQuery, s.JID, time.Now().Add(-7*24*time.Hour).UnixMilli())
+	return err
+}

store/sqlstore/store.go

@@ -1,4 +1,4 @@
--- v0 -> v12 (compatible with v8+): Latest schema
+-- v0 -> v13 (compatible with v8+): Latest schema
 CREATE TABLE whatsmeow_device (
 	jid TEXT PRIMARY KEY,
 	lid TEXT,
@@ -158,3 +158,17 @@ CREATE TABLE whatsmeow_event_buffer (
 	PRIMARY KEY (our_jid, ciphertext_hash),
 	FOREIGN KEY (our_jid) REFERENCES whatsmeow_device(jid) ON DELETE CASCADE ON UPDATE CASCADE
 );
+
+CREATE TABLE whatsmeow_retry_buffer (
+	our_jid    TEXT   NOT NULL,
+	chat_jid   TEXT   NOT NULL,
+	message_id TEXT   NOT NULL,
+	format     TEXT   NOT NULL,
+	plaintext  bytea  NOT NULL,
+	timestamp  BIGINT NOT NULL,
+
+	PRIMARY KEY (our_jid, chat_jid, message_id),
+	FOREIGN KEY (our_jid) REFERENCES whatsmeow_device(jid) ON DELETE CASCADE ON UPDATE CASCADE
+);
+
+CREATE INDEX whatsmeow_retry_buffer_timestamp_idx ON whatsmeow_retry_buffer (our_jid, timestamp);

store/sqlstore/upgrades/00-latest-schema.sql

@@ -0,0 +1,14 @@
+-- v13 (compatible with v8+): Add buffer for outgoing events to accept retry receipts
+CREATE TABLE whatsmeow_retry_buffer (
+	our_jid    TEXT   NOT NULL,
+	chat_jid   TEXT   NOT NULL,
+	message_id TEXT   NOT NULL,
+	format     TEXT   NOT NULL,
+	plaintext  bytea  NOT NULL,
+	timestamp  BIGINT NOT NULL,
+
+	PRIMARY KEY (our_jid, chat_jid, message_id),
+	FOREIGN KEY (our_jid) REFERENCES whatsmeow_device(jid) ON DELETE CASCADE ON UPDATE CASCADE
+);
+
+CREATE INDEX whatsmeow_retry_buffer_timestamp_idx ON whatsmeow_retry_buffer (our_jid, timestamp);

store/sqlstore/upgrades/13-retry-buffer.sql

@@ -158,6 +158,10 @@ type EventBuffer interface {
 	DoDecryptionTxn(ctx context.Context, fn func(context.Context) error) error
 	ClearBufferedEventPlaintext(ctx context.Context, ciphertextHash [32]byte) error
 	DeleteOldBufferedHashes(ctx context.Context) error
+
+	GetOutgoingEvent(ctx context.Context, chatJID, altChatJID types.JID, id types.MessageID) (string, []byte, error)
+	AddOutgoingEvent(ctx context.Context, chatJID types.JID, id types.MessageID, format string, plaintext []byte) error
+	DeleteOldOutgoingEvents(ctx context.Context) error
 }
 
 type LIDMapping struct {