XEP-0308 message editing

5 files changed, 184 insertions, 18 deletions

diff --git a/badger/ids.go b/badger/ids.go
index 80fb9ad..079887a 100644
--- a/badger/ids.go
+++ b/badger/ids.go
@@ -121,6 +121,104 @@ func splitTgByteString(val []byte) (int64, int64, error) {
 	return tgChatId, tgMsgId, err
 }
 
+// ReplaceIdPair replaces an old entry by XMPP ID with both new XMPP and Tg ID
+func (db *IdsDB) ReplaceIdPair(tgAccount, xmppAccount, oldXmppId, newXmppId string, newMsgId int64) error {
+	// read old pair
+	chatId, oldMsgId, err := db.GetByXmppId(tgAccount, xmppAccount, oldXmppId)
+	if err != nil {
+		return err
+	}
+
+	bPrefix := toKeyPrefix(tgAccount, xmppAccount)
+
+	bOldTgId := toTgByteString(chatId, oldMsgId)
+	bOldXmppId := toXmppByteString(oldXmppId)
+	bOldTgKey := toByteKey(bPrefix, bOldTgId, "tg")
+	bOldXmppKey := toByteKey(bPrefix, bOldXmppId, "xmpp")
+
+	bTgId := toTgByteString(chatId, newMsgId)
+	bXmppId := toXmppByteString(newXmppId)
+	bTgKey := toByteKey(bPrefix, bTgId, "tg")
+	bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")
+
+	return db.db.Update(func(txn *badger.Txn) error {
+		// save new pair
+		if err := txn.Set(bTgKey, bXmppId); err != nil {
+			return err
+		}
+		if err := txn.Set(bXmppKey, bTgId); err != nil {
+			return err
+		}
+		// delete old pair
+		if err := txn.Delete(bOldTgKey); err != nil {
+			return err
+		}
+		return txn.Delete(bOldXmppKey)
+	})
+}
+
+// ReplaceXmppId replaces an old XMPP ID with new XMPP ID and keeps Tg ID intact
+func (db *IdsDB) ReplaceXmppId(tgAccount, xmppAccount, oldXmppId, newXmppId string) error {
+	// read old Tg IDs
+	chatId, msgId, err := db.GetByXmppId(tgAccount, xmppAccount, oldXmppId)
+	if err != nil {
+		return err
+	}
+
+	bPrefix := toKeyPrefix(tgAccount, xmppAccount)
+
+	bOldXmppId := toXmppByteString(oldXmppId)
+	bOldXmppKey := toByteKey(bPrefix, bOldXmppId, "xmpp")
+
+	bTgId := toTgByteString(chatId, msgId)
+	bXmppId := toXmppByteString(newXmppId)
+	bTgKey := toByteKey(bPrefix, bTgId, "tg")
+	bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")
+
+	return db.db.Update(func(txn *badger.Txn) error {
+		// save new pair
+		if err := txn.Set(bTgKey, bXmppId); err != nil {
+			return err
+		}
+		if err := txn.Set(bXmppKey, bTgId); err != nil {
+			return err
+		}
+		// delete old xmpp->tg entry
+		return txn.Delete(bOldXmppKey)
+	})
+}
+
+// ReplaceTgId replaces an old Tg ID with new Tg ID and keeps Tg chat ID and XMPP ID intact
+func (db *IdsDB) ReplaceTgId(tgAccount, xmppAccount string, chatId, oldMsgId, newMsgId int64) error {
+	// read old XMPP ID
+	xmppId, err := db.GetByTgIds(tgAccount, xmppAccount, chatId, oldMsgId)
+	if err != nil {
+		return err
+	}
+
+	bPrefix := toKeyPrefix(tgAccount, xmppAccount)
+
+	bOldTgId := toTgByteString(chatId, oldMsgId)
+	bOldTgKey := toByteKey(bPrefix, bOldTgId, "tg")
+
+	bTgId := toTgByteString(chatId, newMsgId)
+	bXmppId := toXmppByteString(xmppId)
+	bTgKey := toByteKey(bPrefix, bTgId, "tg")
+	bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")
+
+	return db.db.Update(func(txn *badger.Txn) error {
+		// save new pair
+		if err := txn.Set(bTgKey, bXmppId); err != nil {
+			return err
+		}
+		if err := txn.Set(bXmppKey, bTgId); err != nil {
+			return err
+		}
+		// delete old tg->xmpp entry
+		return txn.Delete(bOldTgKey)
+	})
+}
+
 // Gc compacts the value log
 func (db *IdsDB) Gc() {
 	db.db.RunValueLogGC(0.7)
diff --git a/telegram/handlers.go b/telegram/handlers.go
index bd768ae..6de59e7 100644
--- a/telegram/handlers.go
+++ b/telegram/handlers.go
@@ -272,6 +272,11 @@ func (c *Client) updateAuthorizationState(update *client.UpdateAuthorizationStat
 
 // clean uploaded files
 func (c *Client) updateMessageSendSucceeded(update *client.UpdateMessageSendSucceeded) {
+	log.Debugf("replace message %v with %v", update.OldMessageId, update.Message.Id)
+	if err := gateway.IdsDB.ReplaceTgId(c.Session.Login, c.jid, update.Message.ChatId, update.OldMessageId, update.Message.Id); err != nil {
+		log.Error("failed to replace %v with %v: %v", update.OldMessageId, update.Message.Id, err.Error())
+	}
+
 	file, _ := c.contentToFile(update.Message.Content)
 	if file != nil && file.Local != nil {
 		c.cleanTempFile(file.Local.Path)
diff --git a/telegram/utils.go b/telegram/utils.go
index dd63248..ecce6da 100644
--- a/telegram/utils.go
+++ b/telegram/utils.go
@@ -991,14 +991,14 @@ func (c *Client) PrepareOutgoingMessageContent(text string) client.InputMessageC
 	return c.prepareOutgoingMessageContent(text, nil)
 }
 
-// ProcessOutgoingMessage executes commands or sends messages to mapped chats
-func (c *Client) ProcessOutgoingMessage(chatID int64, text string, returnJid string, id string, replyId int64) {
+// ProcessOutgoingMessage executes commands or sends messages to mapped chats, returns message id
+func (c *Client) ProcessOutgoingMessage(chatID int64, text string, returnJid string, replyId int64, replaceId int64) int64 {
 	if !c.Online() {
 		// we're offline
-		return
+		return 0
 	}
 
-	if strings.HasPrefix(text, "/") || strings.HasPrefix(text, "!") {
+	if replaceId == 0 && (strings.HasPrefix(text, "/") || strings.HasPrefix(text, "!")) {
 		// try to execute commands
 		response, isCommand := c.ProcessChatCommand(chatID, text)
 		if response != "" {
@@ -1006,7 +1006,7 @@ func (c *Client) ProcessOutgoingMessage(chatID int64, text string, returnJid str
 		}
 		// do not send on success
 		if isCommand {
-			return
+			return 0
 		}
 	}
 
@@ -1014,7 +1014,7 @@ func (c *Client) ProcessOutgoingMessage(chatID int64, text string, returnJid str
 
 	// quotations
 	var reply int64
-	if replyId == 0 {
+	if replaceId == 0 && replyId == 0 {
 		replySlice := replyRegex.FindStringSubmatch(text)
 		if len(replySlice) > 1 {
 			reply, _ = strconv.ParseInt(replySlice[1], 10, 64)
@@ -1069,24 +1069,29 @@ func (c *Client) ProcessOutgoingMessage(chatID int64, text string, returnJid str
 
 	content := c.prepareOutgoingMessageContent(text, file)
 
+	if replaceId != 0 {
+		tgMessage, err := c.client.EditMessageText(&client.EditMessageTextRequest{
+			ChatId:              chatID,
+			MessageId:           replaceId,
+			InputMessageContent: content,
+		})
+		if err != nil {
+			c.returnError(returnJid, chatID, "Not edited", err)
+			return 0
+		}
+		return tgMessage.Id
+	}
+
 	tgMessage, err := c.client.SendMessage(&client.SendMessageRequest{
 		ChatId:              chatID,
 		ReplyToMessageId:    reply,
 		InputMessageContent: content,
 	})
 	if err != nil {
-		gateway.SendTextMessage(
-			returnJid,
-			strconv.FormatInt(chatID, 10),
-			fmt.Sprintf("Not sent: %s", err.Error()),
-			c.xmpp,
-		)
-	} else {
-		err = gateway.IdsDB.Set(c.Session.Login, c.jid, tgMessage.ChatId, tgMessage.Id, id)
-		if err != nil {
-			log.Errorf("Failed to save ids %v/%v %v", tgMessage.ChatId, tgMessage.Id, id)
-		}
+		c.returnError(returnJid, chatID, "Not sent", err)
+		return 0
 	}
+	return tgMessage.Id
 }
 
 func (c *Client) returnMessage(returnJid string, chatID int64, text string) {
diff --git a/xmpp/extensions/extensions.go b/xmpp/extensions/extensions.go
index 78de47d..2d547af 100644
--- a/xmpp/extensions/extensions.go
+++ b/xmpp/extensions/extensions.go
@@ -180,6 +180,12 @@ type ClientMessage struct {
 	Extensions []stanza.MsgExtension `xml:",omitempty"`
 }
 
+// Replace is from XEP-0308
+type Replace struct {
+	XMLName xml.Name `xml:"urn:xmpp:message-correct:0 replace"`
+	Id      string   `xml:"id,attr"`
+}
+
 // Namespace is a namespace!
 func (c PresenceNickExtension) Namespace() string {
 	return c.XMLName.Space
@@ -225,6 +231,11 @@ func (c ComponentPrivilege) Namespace() string {
 	return c.XMLName.Space
 }
 
+// Namespace is a namespace!
+func (c Replace) Namespace() string {
+	return c.XMLName.Space
+}
+
 // Name is a packet name
 func (ClientMessage) Name() string {
 	return "message"
@@ -291,4 +302,10 @@ func init() {
 		"urn:xmpp:privilege:1",
 		"privilege",
 	}, ComponentPrivilege{})
+
+	// message edit
+	stanza.TypeRegistry.MapExtension(stanza.PKTMessage, xml.Name{
+		"urn:xmpp:message-correct:0",
+		"replace",
+	}, Replace{})
 }
diff --git a/xmpp/handlers.go b/xmpp/handlers.go
index 5178acd..51fd831 100644
--- a/xmpp/handlers.go
+++ b/xmpp/handlers.go
@@ -102,10 +102,13 @@ func HandleMessage(s xmpp.Sender, p stanza.Packet) {
 		if ok {
 			var reply extensions.Reply
 			var fallback extensions.Fallback
+			var replace extensions.Replace
 			msg.Get(&reply)
 			msg.Get(&fallback)
+			msg.Get(&replace)
 			log.Debugf("reply: %#v", reply)
 			log.Debugf("fallback: %#v", fallback)
+			log.Debugf("replace: %#v", replace)
 
 			var replyId int64
 			var err error
@@ -138,8 +141,46 @@ func HandleMessage(s xmpp.Sender, p stanza.Packet) {
 					text = text[:start] + text[end:]
 				}
 			}
+			var replaceId int64
+			if replace.Id != "" {
+				chatId, msgId, err := gateway.IdsDB.GetByXmppId(session.Session.Login, bare, replace.Id)
+				if err == nil {
+					if chatId != toID {
+						gateway.SendTextMessage(msg.From, strconv.FormatInt(toID, 10), "<ERROR: Chat mismatch>", component)
+						return
+					}
+					replaceId = msgId
+					log.Debugf("replace tg: %#v %#v", chatId, msgId)
+				} else {
+					gateway.SendTextMessage(msg.From, strconv.FormatInt(toID, 10), "<ERROR: Could not find matching message to edit>", component)
+					return
+				}
+			}
 
-			session.ProcessOutgoingMessage(toID, text, msg.From, msg.Id, replyId)
+			tgMessageId := session.ProcessOutgoingMessage(toID, text, msg.From, replyId, replaceId)
+			if tgMessageId != 0 {
+				if replaceId != 0 {
+					// not needed (is it persistent among clients though?)
+					/* err = gateway.IdsDB.ReplaceIdPair(session.Session.Login, bare, replace.Id, msg.Id, tgMessageId)
+					if err != nil {
+						log.Errorf("Failed to replace id %v with %v %v", replace.Id, msg.Id, tgMessageId)
+					} */
+				} else {
+					err = gateway.IdsDB.Set(session.Session.Login, bare, toID, tgMessageId, msg.Id)
+					if err != nil {
+						log.Errorf("Failed to save ids %v/%v %v", toID, tgMessageId, msg.Id)
+					}
+				}
+			} else {
+				/*
+				// if a message failed to edit on Telegram side, match new XMPP ID with old Telegram ID anyway
+				if replaceId != 0 {
+					err = gateway.IdsDB.ReplaceXmppId(session.Session.Login, bare, replace.Id, msg.Id)
+					if err != nil {
+						log.Errorf("Failed to replace id %v with %v", replace.Id, msg.Id)
+					}
+				} */
+			}
 			return
 		} else {
 			toJid, err := stanza.NewJid(msg.To)