1 files changed, 25 insertions(+), 10 deletions(-)

M 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
 				)