From e40317a3b43594c0629c8a0d1d569b4c8d55e2ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roman=20Smr=C5=BE?= <roman.smrz@seznam.cz>
Date: Thu, 23 May 2024 20:50:39 +0200
Subject: Plaintext data response to plaintext request

---
 src/Erebos/Network.hs | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

(limited to 'src/Erebos/Network.hs')

diff --git a/src/Erebos/Network.hs b/src/Erebos/Network.hs
index 744c476..cc111e2 100644
--- a/src/Erebos/Network.hs
+++ b/src/Erebos/Network.hs
@@ -108,7 +108,7 @@ defaultServerOptions = ServerOptions
 data Peer = Peer
     { peerAddress :: PeerAddress
     , peerServer_ :: Server
-    , peerConnection :: TVar (Either [(Bool, TransportPacket Ref, [TransportHeaderItem])] (Connection PeerAddress))
+    , peerConnection :: TVar (Either [(SecurityRequirement, TransportPacket Ref, [TransportHeaderItem])] (Connection PeerAddress))
     , peerIdentityVar :: TVar PeerIdentity
     , peerStorage_ :: Storage
     , peerInStorage :: PartialStorage
@@ -391,13 +391,17 @@ instance MonadFail PacketHandler where
 runPacketHandler :: Bool -> Peer -> PacketHandler () -> STM ()
 runPacketHandler secure peer@Peer {..} act = do
     let logd = writeTQueue $ serverErrorLog peerServer_
-    runExceptT (flip execStateT (PacketHandlerState peer [] [] []) $ unPacketHandler act) >>= \case
+    runExceptT (flip execStateT (PacketHandlerState peer [] [] [] False) $ unPacketHandler act) >>= \case
         Left err -> do
             logd $ "Error in handling packet from " ++ show peerAddress ++ ": " ++ err
         Right ph -> do
             when (not $ null $ phHead ph) $ do
                 let packet = TransportPacket (TransportHeader $ phHead ph) (phBody ph)
-                sendToPeerS' secure peer (phAckedBy ph) packet
+                    secreq = case (secure, phPlaintextReply ph) of
+                                  (True, _) -> EncryptedOnly
+                                  (False, False) -> PlaintextAllowed
+                                  (False, True) -> PlaintextOnly
+                sendToPeerS' secreq peer (phAckedBy ph) packet
 
 liftSTM :: STM a -> PacketHandler a
 liftSTM = PacketHandler . lift . lift
@@ -416,6 +420,7 @@ data PacketHandlerState = PacketHandlerState
     , phHead :: [TransportHeaderItem]
     , phAckedBy :: [TransportHeaderItem]
     , phBody :: [Ref]
+    , phPlaintextReply :: Bool
     }
 
 addHeader :: TransportHeaderItem -> PacketHandler ()
@@ -427,6 +432,9 @@ addAckedBy hs = modify $ \ph -> ph { phAckedBy = foldr appendDistinct (phAckedBy
 addBody :: Ref -> PacketHandler ()
 addBody r = modify $ \ph -> ph { phBody = r `appendDistinct` phBody ph }
 
+keepPlaintextReply :: PacketHandler ()
+keepPlaintextReply = modify $ \ph -> ph { phPlaintextReply = True }
+
 openStream :: PacketHandler RawStreamWriter
 openStream = do
     Peer {..} <- gets phPeer
@@ -489,6 +497,11 @@ handlePacket identity secure peer chanSvc svcs (TransportHeader headers) prefs =
                         partialRefFromDigest (peerInStorage peer) dgst
                     addHeader $ DataResponse dgst
                     addAckedBy [ Acknowledged dgst, Rejected dgst ]
+
+                    -- Plaintext request may indicate the peer has restarted/changed or
+                    -- otherwise lost the channel, so keep the reply plaintext as well.
+                    when (not secure) keepPlaintextReply
+
                     let bytes = lazyLoadBytes mref
                     -- TODO: MTU
                     if (secure && BL.length bytes > 500)
@@ -762,17 +775,17 @@ sendToPeerList peer parts = do
         ackedBy = concat [[ Acknowledged r, Rejected r, DataRequest r ] | r <- dgsts ]
     liftIO $ atomically $ sendToPeerS peer ackedBy packet
 
-sendToPeerS' :: Bool -> Peer -> [TransportHeaderItem] -> TransportPacket Ref -> STM ()
+sendToPeerS' :: SecurityRequirement -> Peer -> [TransportHeaderItem] -> TransportPacket Ref -> STM ()
 sendToPeerS' secure Peer {..} ackedBy packet = do
     readTVar peerConnection >>= \case
         Left xs -> writeTVar peerConnection $ Left $ (secure, packet, ackedBy) : xs
         Right conn -> writeFlow (connData conn) (secure, packet, ackedBy)
 
 sendToPeerS :: Peer -> [TransportHeaderItem] -> TransportPacket Ref -> STM ()
-sendToPeerS = sendToPeerS' True
+sendToPeerS = sendToPeerS' EncryptedOnly
 
 sendToPeerPlain :: Peer -> [TransportHeaderItem] -> TransportPacket Ref -> STM ()
-sendToPeerPlain = sendToPeerS' False
+sendToPeerPlain = sendToPeerS' PlaintextAllowed
 
 sendToPeerWith :: forall s m. (Service s, MonadIO m, MonadError String m) => Peer -> (ServiceState s -> ExceptT String IO (Maybe s, ServiceState s)) -> m ()
 sendToPeerWith peer fobj = do
-- 
cgit v1.2.3