Avoid leaking unencrypted messages on PEP timeout

e2ee/fetch.go

@@ -25,7 +25,14 @@ const FetchTimeout = 15 * time.Second
 // entirely by whatever PEP node names backend itself reports via
 // DeviceListNode/BundleNode, it has no OMEMO-specific knowledge of its own,
 // so it works unchanged for any future Backend implementation.
-func FetchAndIngestPeer(component *xmpp.Component, backend Backend, owner PeerID, peer PeerID) error {
+//
+// ctx should be a long-lived, process-scoped context (e.g.
+// gateway.ShutdownCtx), not context.Background() - each PEP request derives
+// its own FetchTimeout-bounded child from it (see fetchOneItem), and
+// deriving from a shutdown-aware parent means process teardown cancels any
+// in-flight request immediately instead of leaving it to run out its own
+// timeout regardless of teardown.
+func FetchAndIngestPeer(ctx context.Context, component *xmpp.Component, backend Backend, owner PeerID, peer PeerID) error {
 	_, ownerBareJID, ok := SplitOwnedPeer(owner)
 	if !ok {
 		return fmt.Errorf("e2ee: FetchAndIngestPeer: %q is not an OwnedPeer", owner)
@@ -39,7 +46,7 @@ func FetchAndIngestPeer(component *xmpp.Component, backend Backend, owner PeerID
 		return fmt.Errorf("e2ee: NegotiatedVariant: %w", err)
 	}
 
-	deviceListItem, err := fetchOneItem(component, ownerBareJID, string(peer), backend.DeviceListNode(variant))
+	deviceListItem, err := fetchOneItem(ctx, component, ownerBareJID, string(peer), backend.DeviceListNode(variant))
 	if err != nil {
 		return fmt.Errorf("e2ee: fetch device list for %s: %w", peer, err)
 	}
@@ -68,7 +75,7 @@ func FetchAndIngestPeer(component *xmpp.Component, backend Backend, owner PeerID
 		if known[id] {
 			continue
 		}
-		bundleItem, err := fetchOneItem(component, ownerBareJID, string(peer), backend.BundleNode(variant, id))
+		bundleItem, err := fetchOneItem(ctx, component, ownerBareJID, string(peer), backend.BundleNode(variant, id))
 		if err != nil {
 			lastErr = fmt.Errorf("e2ee: fetch bundle for %s:%s: %w", peer, id, err)
 			continue
@@ -94,14 +101,14 @@ func FetchAndIngestPeer(component *xmpp.Component, backend Backend, owner PeerID
 // fetchOneItem requests the single most recent item of node from jid's PEP
 // service, sent as if from ownerBareJID (the bridged chat's pseudo-JID doing
 // the fetching), and returns its payload node.
-func fetchOneItem(component *xmpp.Component, ownerBareJID, jid, node string) (*stanza.Node, error) {
+func fetchOneItem(parent context.Context, component *xmpp.Component, ownerBareJID, jid, node string) (*stanza.Node, error) {
 	iq, err := stanza.NewItemsRequest(jid, node, 1)
 	if err != nil {
 		return nil, err
 	}
 	iq.Attrs.From = ownerBareJID
 
-	ctx, cancel := context.WithTimeout(context.Background(), FetchTimeout)
+	ctx, cancel := context.WithTimeout(parent, FetchTimeout)
 	defer cancel()
 
 	ch, err := component.SendIQ(ctx, iq)

telegram/handlers.go

@@ -409,8 +409,12 @@ func (c *Client) updateMessageContent(update *client.UpdateMessageContent) {
 			// OMEMO encrypt hook - same rules as SendMessageToGateway's (see
 			// its doc comment): personal-chat pseudo-JIDs only, gated on the
 			// per-chat active flag, once per logical edit rather than per
-			// jids resource copy.
+			// jids resource copy. If OMEMO is active but a session can't be
+			// established or encryption fails, the real edited content must
+			// NEVER go out in the clear - omemoFailed sends
+			// gateway.OMEMOSendFailedBody instead.
 			var envelope *e2ee.Envelope
+			omemoFailed := false
 			if !isMUC {
 				if backend, ok := gateway.E2EE.Backend(); ok {
 					owner := e2ee.OwnedPeer(c.Session.Login, gateway.CHATJID(update.ChatId, false))
@@ -418,8 +422,10 @@ func (c *Client) updateMessageContent(update *client.UpdateMessageContent) {
 						peer := e2ee.PeerID(c.jid)
 						if err := c.ensureOMEMOSession(backend, owner, peer); err != nil {
 							log.Error(errors.Wrap(err, "Failed to establish OMEMO session"))
+							omemoFailed = true
 						} else if env, err := backend.Encrypt(owner, []e2ee.PeerID{peer}, []byte(text.String())); err != nil {
 							log.Error(errors.Wrap(err, "Failed to encrypt OMEMO message"))
+							omemoFailed = true
 						} else {
 							envelope = &env
 						}
@@ -427,15 +433,20 @@ func (c *Client) updateMessageContent(update *client.UpdateMessageContent) {
 				}
 			}
 
+			body := text.String()
+			if omemoFailed {
+				body = gateway.OMEMOSendFailedBody
+			}
+
 			for _, jid := range jids {
 				if safeToSend {
 					gateway.SendMessage(jid, from, c.xmpp,
-						gateway.SMBody(text.String()), gateway.SMId(id), gateway.SMReplaceId(replaceId),
+						gateway.SMBody(body), gateway.SMId(id), gateway.SMReplaceId(replaceId),
 						gateway.SMIsCarbon(isCarbon), gateway.SMIsGroupchat(isMUC), gateway.SMOriginalFrom(originalFrom),
 						gateway.SMOMEMOEnvelope(envelope),
 					)
 				} else {
-					gateway.SendMUCAnnouncement(jid, from, text.String(), nickname, id, c.xmpp)
+					gateway.SendMUCAnnouncement(jid, from, body, nickname, id, c.xmpp)
 				}
 			}
 		}()

telegram/utils.go

@@ -2071,26 +2071,36 @@ func (c *Client) SendMessageToGateway(chatId int64, message *client.Message, id
 	// resources all share the same OMEMO device fan-out. OOB file links are
 	// intentionally left in the clear (encrypting a URL string doesn't
 	// protect the file itself); only auxText, and text when there's no OOB
-	// swap, are ever encrypted.
+	// swap, are ever encrypted. If OMEMO is active but a session can't be
+	// established or encryption fails (e.g. the peer never answers the PEP
+	// bundle fetch - see ensureOMEMOSession/gateway.ShutdownCtx), the real
+	// content must NEVER go out in the clear as a fallback - omemoFailed
+	// short-circuits the send loop below to gateway.OMEMOSendFailedBody
+	// instead.
 	var textEnvelope, auxEnvelope *e2ee.Envelope
-	if !isGroupchat {
+	omemoFailed := false
+	needsOMEMOEncryption := (oob == "" && text != "") || auxText != ""
+	if !isGroupchat && needsOMEMOEncryption {
 		if backend, ok := gateway.E2EE.Backend(); ok {
 			owner := e2ee.OwnedPeer(c.Session.Login, gateway.CHATJID(chatId, false))
 			if active, _ := backend.Enabled(owner); active {
 				peer := e2ee.PeerID(c.jid)
 				if err := c.ensureOMEMOSession(backend, owner, peer); err != nil {
 					log.Error(errors.Wrap(err, "Failed to establish OMEMO session"))
+					omemoFailed = true
 				} else {
 					if oob == "" && text != "" {
 						if env, err := backend.Encrypt(owner, []e2ee.PeerID{peer}, []byte(text)); err != nil {
 							log.Error(errors.Wrap(err, "Failed to encrypt OMEMO message"))
+							omemoFailed = true
 						} else {
 							textEnvelope = &env
 						}
 					}
-					if auxText != "" {
+					if !omemoFailed && auxText != "" {
 						if env, err := backend.Encrypt(owner, []e2ee.PeerID{peer}, []byte(auxText)); err != nil {
 							log.Error(errors.Wrap(err, "Failed to encrypt OMEMO auxiliary message"))
+							omemoFailed = true
 						} else {
 							auxEnvelope = &env
 						}
@@ -2107,6 +2117,11 @@ func (c *Client) SendMessageToGateway(chatId int64, message *client.Message, id
 			gateway.SMOriginalFrom(originalFrom), gateway.SMMamQueryId(mamQueryId),
 			gateway.SMMucJID(mucJID), gateway.SMMucUserItem(mucUserItem),
 		}
+		if omemoFailed {
+			gateway.SendMessage(jid, from, c.xmpp, append(commonArgs,
+				gateway.SMBody(gateway.OMEMOSendFailedBody), gateway.SMId(sId), gateway.SMStanzaId(stanzaId))...)
+			continue
+		}
 		gateway.SendMessage(jid, from, c.xmpp, append(commonArgs,
 			gateway.SMBody(text), gateway.SMId(sId), gateway.SMStanzaId(stanzaId), gateway.SMOOB(oob),
 			gateway.SMOMEMOEnvelope(textEnvelope))...)
@@ -2132,7 +2147,7 @@ func (c *Client) ensureOMEMOSession(backend e2ee.Backend, owner e2ee.PeerID, pee
 	if len(devices) > 0 {
 		return nil
 	}
-	return e2ee.FetchAndIngestPeer(c.xmpp, backend, owner, peer)
+	return e2ee.FetchAndIngestPeer(gateway.ShutdownCtx, c.xmpp, backend, owner, peer)
 }
 
 // SendDelayedMUCMessage is used to send MUC history via the legacy method or MAM

xmpp/component.go

@@ -287,6 +287,11 @@ func SaveSessions() {
 func Close(component *xmpp.Component) {
 	log.Error("Disconnecting...")
 
+	// cancel any in-flight background operations (e.g. e2ee bundle
+	// fetches) derived from gateway.ShutdownCtx, rather than letting them
+	// run out their own timeout regardless of teardown
+	gateway.CancelShutdown()
+
 	sessionLock.Lock()
 	// close all sessions
 	for _, session := range sessions {

xmpp/gateway/gateway.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"encoding/xml"
 	"github.com/pkg/errors"
 	"strconv"
@@ -78,6 +79,15 @@ var IdsDB badger.IdsDB
 // feature didn't exist" rather than nil-checking E2EE itself.
 var E2EE *e2ee.Manager
 
+// ShutdownCtx is cancelled once (see xmpp.Close) when the component begins
+// tearing down. Long-running background operations that would otherwise
+// wait out their own timeout regardless of process lifecycle - e.g. e2ee's
+// PEP bundle fetches - should derive their per-call timeout from this
+// (context.WithTimeout(gateway.ShutdownCtx, ...)) rather than
+// context.Background(), so teardown cancels them immediately instead of
+// leaving them to run for up to their full timeout past Close().
+var ShutdownCtx, CancelShutdown = context.WithCancel(context.Background())
+
 // DirtySessions denotes that some Telegram session configurations
 // were changed and need to be re-flushed to the YamlDB
 var DirtySessions = false

xmpp/gateway/omemo.go

@@ -17,6 +17,13 @@ import (
 // actual machine-readable hint).
 const omemoFallbackBody = "[This message is OMEMO end-to-end encrypted]"
 
+// OMEMOSendFailedBody is sent in place of the real message body when OMEMO
+// is active for a chat but encryption couldn't go through (no session
+// could be established, e.g. the peer never answered a PEP bundle
+// request) - callers (telegram/utils.go, telegram/handlers.go) must never
+// fall back to sending the real plaintext in this case, only this notice.
+const OMEMOSendFailedBody = "[Failed to send: could not establish an OMEMO session]"
+
 // omemoStanzaExtension converts a just-produced OMEMO e2ee.Envelope's raw
 // bytes (an omemo.WireEnvelope, gob-encoded - see e2ee/omemo/envelope.go)
 // into the wire stanza.MsgExtension for its version, plus the XEP-0380 EME