From 3e18b12f377db5aee2a8975ac55c938d9ba83d0b Mon Sep 17 00:00:00 2001
From: Stephen Paul Weber <singpolyma@singpolyma.net>
Date: Mon, 16 May 2022 10:59:31 -0500
Subject: [PATCH] Allow direct message route to notify of removal

In case the backend is removing the user for policy reasons, etc, allows us to
clear out the message route setting when the backend is no longer listening.
---
 Main.hs | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/Main.hs b/Main.hs
index 8819e74..f5b21e4 100644
--- a/Main.hs
+++ b/Main.hs
@@ -356,6 +356,13 @@ mapBody f (m@Message { messagePayloads = payloads }) =
 		) payloads
 	}
 
+deleteDirectMessageRoute db userJid = do
+	DB.del db (DB.byJid userJid ["direct-message-route"])
+	mcheoJid <- fmap (parseJID =<<) $ DB.get db (DB.byJid userJid ["cheoJid"])
+	forM_ mcheoJid $ \cheoJid -> do
+		DB.del db (DB.byJid userJid ["cheoJid"])
+		DB.srem db (DB.byNode cheoJid ["owners"]) [bareTxt userJid]
+
 unregisterDirectMessageRoute db componentJid userJid route = do
 	maybeCheoJid <- (parseJID =<<) <$> DB.get db (DB.byJid userJid ["cheoJid"])
 	forM_ maybeCheoJid $ \cheoJid -> do
@@ -683,12 +690,24 @@ handleRegister db componentJid iq@(IQ { iqType = IQSet, iqFrom = Just from }) qu
 handleRegister db componentJid iq@(IQ { iqType = IQSet, iqPayload = Just payload, iqFrom = Just from }) query
 	| [passwordEl] <- isNamed (fromString "{jabber:iq:register}password") =<< elementChildren query = do
 		handleVerificationCode db componentJid (mconcat $ elementText passwordEl) iq from
-handleRegister db componentJid iq@(IQ { iqFrom = Just from, iqType = IQSet }) query
+handleRegister db componentJid iq@(IQ { iqTo = Just to, iqFrom = Just from, iqType = IQSet }) query
 	| [_] <- isNamed (fromString "{jabber:iq:register}remove") =<< elementChildren query = do
-		tel <- fromMaybe mempty <$> DB.get db (DB.byJid from ["registered"])
-		forM_ (telToJid tel (formatJID componentJid)) $ \cheoJid ->
-			DB.del db (DB.byNode cheoJid ["registered"])
-		DB.del db (DB.byJid from ["registered"])
+		fromBackend <- case XMPP.parseJID =<< (unescapeJid . XMPP.strNode <$> XMPP.jidNode to) of
+			Just unproxiedTo -> do
+				route <- (XMPP.parseJID <=< id) <$> DB.get db (DB.byJid unproxiedTo ["direct-message-route"])
+				if fmap bareTxt route == Just (bareTxt from) then do
+					deleteDirectMessageRoute db unproxiedTo
+					return True
+				else
+					return False
+			Nothing -> return False
+
+		when (not fromBackend) $ do
+			tel <- fromMaybe mempty <$> DB.get db (DB.byJid from ["registered"])
+			forM_ (telToJid tel (formatJID componentJid)) $ \cheoJid ->
+				DB.del db (DB.byNode cheoJid ["registered"])
+			DB.del db (DB.byJid from ["registered"])
+
 		return [mkStanzaRec $ iq {
 			iqTo = iqFrom iq,
 			iqFrom = iqTo iq,
@@ -2127,11 +2146,7 @@ main = do
 							return ()
 						Nothing -> do
 							maybeExistingRoute <- (parseJID =<<) <$> DB.get db (DB.byJid userJid' ["direct-message-route"])
-							DB.del db (DB.byJid userJid' ["direct-message-route"])
-							mcheoJid <- fmap (parseJID =<<) $ DB.get db (DB.byJid userJid' ["cheoJid"])
-							forM_ mcheoJid $ \cheoJid -> do
-								DB.del db (DB.byJid userJid' ["cheoJid"])
-								DB.srem db (DB.byNode cheoJid ["owners"]) [bareTxt userJid']
+							deleteDirectMessageRoute db userJid'
 							forM_ maybeExistingRoute $ \existingRoute ->
 								atomically . writeTChan sendToComponent . mkStanzaRec =<< unregisterDirectMessageRoute db componentJid userJid' existingRoute
 				)
-- 
2.38.5