Plaintext data response to plaintext request

2 files changed, 36 insertions, 16 deletions

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
diff --git a/src/Erebos/Network/Protocol.hs b/src/Erebos/Network/Protocol.hs
index 6c974bd..27e05ba 100644
--- a/src/Erebos/Network/Protocol.hs
+++ b/src/Erebos/Network/Protocol.hs
@@ -3,6 +3,7 @@ module Erebos.Network.Protocol (
     transportToObject,
     TransportHeader(..),
     TransportHeaderItem(..),
+    SecurityRequirement(..),
 
     WaitingRef(..),
     WaitingRefCallback,
@@ -93,6 +94,11 @@ data TransportHeaderItem
 newtype Cookie = Cookie ByteString
     deriving (Eq, Show)
 
+data SecurityRequirement = PlaintextOnly
+                         | PlaintextAllowed
+                         | EncryptedOnly
+    deriving (Eq, Ord)
+
 isHeaderItemAcknowledged :: TransportHeaderItem -> Bool
 isHeaderItemAcknowledged = \case
     Acknowledged {} -> False
@@ -165,10 +171,10 @@ data GlobalState addr = (Eq addr, Show addr) => GlobalState
 data Connection addr = Connection
     { cGlobalState :: GlobalState addr
     , cAddress :: addr
-    , cDataUp :: Flow (Bool, TransportPacket PartialObject) (Bool, TransportPacket Ref, [TransportHeaderItem])
-    , cDataInternal :: Flow (Bool, TransportPacket Ref, [TransportHeaderItem]) (Bool, TransportPacket PartialObject)
+    , cDataUp :: Flow (Bool, TransportPacket PartialObject) (SecurityRequirement, TransportPacket Ref, [TransportHeaderItem])
+    , cDataInternal :: Flow (SecurityRequirement, TransportPacket Ref, [TransportHeaderItem]) (Bool, TransportPacket PartialObject)
     , cChannel :: TVar ChannelState
-    , cSecureOutQueue :: TQueue (Bool, TransportPacket Ref, [TransportHeaderItem])
+    , cSecureOutQueue :: TQueue (SecurityRequirement, TransportPacket Ref, [TransportHeaderItem])
     , cMaxInFlightPackets :: TVar Int
     , cReservedPackets :: TVar Int
     , cSentPackets :: TVar [SentPacket]
@@ -180,7 +186,7 @@ data Connection addr = Connection
 connAddress :: Connection addr -> addr
 connAddress = cAddress
 
-connData :: Connection addr -> Flow (Bool, TransportPacket PartialObject) (Bool, TransportPacket Ref, [TransportHeaderItem])
+connData :: Connection addr -> Flow (Bool, TransportPacket PartialObject) (SecurityRequirement, TransportPacket Ref, [TransportHeaderItem])
 connData = cDataUp
 
 connGetChannel :: Connection addr -> STM ChannelState
@@ -734,13 +740,13 @@ processOutgoing gs@GlobalState {..} = do
                     | isJust mbch = do
                         acks <- readTVar cToAcknowledge
                         if null acks then retry
-                                     else return ((True, TransportPacket (TransportHeader []) [], []), Nothing)
+                                     else return ((EncryptedOnly, TransportPacket (TransportHeader []) [], []), Nothing)
                     | otherwise = retry
 
             ((secure, packet@(TransportPacket (TransportHeader hitems) content), plainAckedBy), mbReserved) <-
                 checkOutstanding <|> checkDataInternal <|> checkAcknowledgements
 
-            when (isNothing mbch && secure) $ do
+            when (isNothing mbch && secure >= EncryptedOnly) $ do
                 writeTQueue cSecureOutQueue (secure, packet, plainAckedBy)
 
             acknowledge <- case mbch of
@@ -759,16 +765,17 @@ processOutgoing gs@GlobalState {..} = do
                                  [] -> return ()
                                  xs -> sequence_ $ map (streamAccepted conn) xs
 
-                mbs <- case mbch of
-                    Just ch -> do
+                mbs <- case (secure, mbch) of
+                    (PlaintextOnly, _)          -> return $ Just (BL.toStrict plain, plainAckedBy)
+                    (PlaintextAllowed, Nothing) -> return $ Just (BL.toStrict plain, plainAckedBy)
+                    (_, Just ch) -> do
                         runExceptT (channelEncrypt ch $ BL.toStrict $ 0x00 `BL.cons` plain) >>= \case
                             Right (ctext, counter) -> do
                                 let isAcked = any isHeaderItemAcknowledged hitems
                                 return $ Just (0x80 `B.cons` ctext, if isAcked then [ AcknowledgedSingle $ fromIntegral counter ] else [])
                             Left err -> do atomically $ gLog $ "Failed to encrypt data: " ++ err
                                            return Nothing
-                    Nothing | secure    -> return Nothing
-                            | otherwise -> return $ Just (BL.toStrict plain, plainAckedBy)
+                    (EncryptedOnly, Nothing)    -> return Nothing
 
                 case mbs of
                     Just (bs, ackedBy) -> do