Protocol: use cookies during whole plaintext phase

3 files changed, 111 insertions, 65 deletions

diff --git a/src/Network.hs b/src/Network.hs
index 5455c07..8cb2ed5 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -420,9 +420,9 @@ handlePacket identity secure peer chanSvc svcs (TransportHeader headers) prefs =
             [ [ storedRef sidentity ]
             , map storedRef $ idUpdates identity
             , case ochannel of
-                   ChannelOurRequest req  -> [ storedRef req ]
-                   ChannelOurAccept acc _ -> [ storedRef acc ]
-                   _                      -> []
+                   ChannelOurRequest _ req  -> [ storedRef req ]
+                   ChannelOurAccept _ acc _ -> [ storedRef acc ]
+                   _                        -> []
             ]
 
     runPacketHandler secure peer $ do
@@ -430,7 +430,7 @@ handlePacket identity secure peer chanSvc svcs (TransportHeader headers) prefs =
         forM_ headers $ \case
             Acknowledged dgst -> do
                 liftSTM (getPeerChannel peer) >>= \case
-                    ChannelOurAccept acc ch | refDigest (storedRef acc) == dgst -> do
+                    ChannelOurAccept _ acc ch | refDigest (storedRef acc) == dgst -> do
                         liftSTM $ finalizedChannel peer ch identity
                     _ -> return ()
 
@@ -474,21 +474,22 @@ handlePacket identity secure peer chanSvc svcs (TransportHeader headers) prefs =
                     _ -> return ()
 
             TrChannelRequest dgst -> do
-                let process = do
+                let process cookie = do
                         addHeader $ Acknowledged dgst
                         wref <- newWaitingRef dgst $ handleChannelRequest peer identity
-                        liftSTM $ setPeerChannel peer $ ChannelPeerRequest wref
+                        liftSTM $ setPeerChannel peer $ ChannelPeerRequest cookie wref
                     reject = addHeader $ Rejected dgst
 
                 liftSTM (getPeerChannel peer) >>= \case
-                    ChannelNone {} -> process
-                    ChannelCookieWait {} -> process
-                    ChannelCookieReceived {} -> process
-                    ChannelOurRequest our | dgst < refDigest (storedRef our) -> process
-                                          | otherwise -> reject
-                    ChannelPeerRequest {} -> process
+                    ChannelNone {} -> return ()
+                    ChannelCookieWait {} -> return ()
+                    ChannelCookieReceived cookie -> process $ Just cookie
+                    ChannelCookieConfirmed cookie -> process $ Just cookie
+                    ChannelOurRequest mbcookie our | dgst < refDigest (storedRef our) -> process mbcookie
+                                                   | otherwise -> reject
+                    ChannelPeerRequest mbcookie _ -> process mbcookie
                     ChannelOurAccept {} -> reject
-                    ChannelEstablished {} -> process
+                    ChannelEstablished {} -> process Nothing
 
             TrChannelAccept dgst -> do
                 let process = do
@@ -498,10 +499,11 @@ handlePacket identity secure peer chanSvc svcs (TransportHeader headers) prefs =
                     ChannelNone {} -> reject
                     ChannelCookieWait {} -> reject
                     ChannelCookieReceived {} -> reject
+                    ChannelCookieConfirmed {} -> reject
                     ChannelOurRequest {} -> process
                     ChannelPeerRequest {} -> process
-                    ChannelOurAccept our _ | dgst < refDigest (storedRef our) -> process
-                                           | otherwise -> addHeader $ Rejected dgst
+                    ChannelOurAccept _ our _ | dgst < refDigest (storedRef our) -> process
+                                             | otherwise -> addHeader $ Rejected dgst
                     ChannelEstablished {} -> process
 
             ServiceType _ -> return ()
@@ -538,10 +540,10 @@ setupChannel identity peer upid = do
             ]
     liftIO $ atomically $ do
         getPeerChannel peer >>= \case
-            ChannelCookieReceived {} -> do
+            ChannelCookieConfirmed cookie -> do
                 sendToPeerPlain peer [ Acknowledged reqref, Rejected reqref ] $
                     TransportPacket (TransportHeader hitems) [storedRef req]
-                setPeerChannel peer $ ChannelOurRequest req
+                setPeerChannel peer $ ChannelOurRequest (Just cookie) req
             _ -> return ()
 
 handleChannelRequest :: Peer -> UnifiedIdentity -> Ref -> WaitingRefCallback
@@ -550,8 +552,8 @@ handleChannelRequest peer identity req = do
         (acc, ch) <- flip runReaderT (peerStorage peer) $ acceptChannelRequest identity upid (wrappedLoad req)
         liftIO $ atomically $ do
             getPeerChannel peer >>= \case
-                ChannelPeerRequest wr | wrDigest wr == refDigest req -> do
-                    setPeerChannel peer $ ChannelOurAccept acc ch
+                ChannelPeerRequest mbcookie wr | wrDigest wr == refDigest req -> do
+                    setPeerChannel peer $ ChannelOurAccept mbcookie acc ch
                     let accref = refDigest $ storedRef acc
                         header = TrChannelAccept accref
                         ackedBy = [ Acknowledged accref, Rejected accref ]
diff --git a/src/Network/Protocol.hs b/src/Network/Protocol.hs
index dc33296..e286a67 100644
--- a/src/Network/Protocol.hs
+++ b/src/Network/Protocol.hs
@@ -59,6 +59,7 @@ protocolVersions = [protocolVersion]
 data TransportPacket a = TransportPacket TransportHeader [a]
 
 data TransportHeader = TransportHeader [TransportHeaderItem]
+    deriving (Show)
 
 data TransportHeaderItem
     = Acknowledged RefDigest
@@ -75,10 +76,10 @@ data TransportHeaderItem
     | TrChannelAccept RefDigest
     | ServiceType ServiceID
     | ServiceRef RefDigest
-    deriving (Eq)
+    deriving (Eq, Show)
 
 newtype Cookie = Cookie ByteString
-    deriving (Eq)
+    deriving (Eq, Show)
 
 transportToObject :: PartialStorage -> TransportHeader -> PartialObject
 transportToObject st (TransportHeader items) = Rec $ map single items
@@ -171,11 +172,12 @@ wrDigest = refDigest . wrefPartial
 
 
 data ChannelState = ChannelNone
-                  | ChannelCookieWait
-                  | ChannelCookieReceived Cookie
-                  | ChannelOurRequest (Stored ChannelRequest)
-                  | ChannelPeerRequest WaitingRef
-                  | ChannelOurAccept (Stored ChannelAccept) Channel
+                  | ChannelCookieWait -- sent initiation, waiting for response
+                  | ChannelCookieReceived Cookie -- received cookie, but no cookie echo yet
+                  | ChannelCookieConfirmed Cookie -- received cookie echo, no need to send from our side
+                  | ChannelOurRequest (Maybe Cookie) (Stored ChannelRequest)
+                  | ChannelPeerRequest (Maybe Cookie) WaitingRef
+                  | ChannelOurAccept (Maybe Cookie) (Stored ChannelAccept) Channel
                   | ChannelEstablished Channel
 
 
@@ -272,9 +274,9 @@ processIncoming gs@GlobalState {..} = do
     mbch <- case mbconn of
         Nothing -> return Nothing
         Just conn -> readTVar (cChannel conn) >>= return . \case
-            ChannelEstablished ch -> Just ch
-            ChannelOurAccept _ ch -> Just ch
-            _                     -> Nothing
+            ChannelEstablished ch   -> Just ch
+            ChannelOurAccept _ _ ch -> Just ch
+            _                       -> Nothing
 
     return $ do
         let deserialize = liftEither . runExcept . deserializeObjects gStorage . BL.fromStrict
@@ -324,9 +326,43 @@ processIncoming gs@GlobalState {..} = do
 
 processPacket :: GlobalState addr -> Either addr (Connection addr) -> Bool -> TransportPacket a -> IO (Maybe (Connection addr, Maybe (TransportPacket a)))
 processPacket gs@GlobalState {..} econn secure packet@(TransportPacket (TransportHeader header) _) = if
+    -- Established secure communication
     | Right conn <- econn, secure
     -> return $ Just (conn, Just packet)
 
+    -- Plaintext communication with cookies to prove origin
+    | cookie:_ <- mapMaybe (\case CookieEcho x -> Just x; _ -> Nothing) header
+    -> verifyCookie gs addr cookie >>= \case
+        True -> do
+            atomically $ do
+                conn@Connection {..} <- getConnection gs addr
+                channel <- readTVar cChannel
+                let received = listToMaybe $ mapMaybe (\case CookieSet x -> Just x; _ -> Nothing) header
+                case received `mplus` channelCurrentCookie channel of
+                    Just current -> do
+                        cookieEchoReceived gs conn mbpid current
+                        return $ Just (conn, Just packet)
+                    Nothing -> do
+                        gLog $ show addr <> ": missing cookie set, dropping " <> show header
+                        return $ Nothing
+
+        False -> do
+            atomically $ gLog $ show addr <> ": cookie verification failed, dropping " <> show header
+            return Nothing
+
+    -- Response to initiation packet
+    | cookie:_ <- mapMaybe (\case CookieSet x -> Just x; _ -> Nothing) header
+    , Just _ <- version
+    , Right conn@Connection {..} <- econn
+    -> do
+        atomically $ readTVar cChannel >>= \case
+            ChannelCookieWait -> do
+                writeTVar cChannel $ ChannelCookieReceived cookie
+                writeFlow gControlFlow (NewConnection conn mbpid)
+                return $ Just (conn, Nothing)
+            _ -> return Nothing
+
+    -- Initiation packet
     | _:_ <- mapMaybe (\case Initiation x -> Just x; _ -> Nothing) header
     , Just ver <- version
     -> do
@@ -341,33 +377,7 @@ processPacket gs@GlobalState {..} econn secure packet@(TransportPacket (Transpor
             writeFlow gDataFlow (addr, reply)
         return Nothing
 
-    | cookie:_ <- mapMaybe (\case CookieSet x -> Just x; _ -> Nothing) header
-    , Just _ <- version
-    , Right conn@Connection {..} <- econn
-    -> do
-        atomically $ readTVar cChannel >>= \case
-            ChannelCookieWait -> do
-                writeTVar cChannel $ ChannelCookieReceived cookie
-                writeFlow gControlFlow (NewConnection conn mbpid)
-                return $ Just (conn, Nothing)
-            _ -> return Nothing
-
-    | Right conn <- econn
-    -> return $ Just (conn, Just packet)
-
-    | cookie:_ <- mapMaybe (\case CookieEcho x -> Just x; _ -> Nothing) header
-    , Just _ <- version
-    -> verifyCookie gs addr cookie >>= \case
-        True -> do
-            conn <- atomically $ findConnection gs addr >>= \case
-                Just conn -> return conn
-                Nothing -> do
-                    conn <- newConnection gs addr
-                    writeFlow gControlFlow (NewConnection conn mbpid)
-                    return conn
-            return $ Just (conn, Just packet)
-        False -> return Nothing
-
+    -- Announce packet outside any connection
     | dgst:_ <- mapMaybe (\case AnnounceSelf x -> Just x; _ -> Nothing) header
     , Just _ <- version
     -> do
@@ -377,13 +387,46 @@ processPacket gs@GlobalState {..} econn secure packet@(TransportPacket (Transpor
                 writeFlow gControlFlow $ ReceivedAnnounce addr dgst
         return Nothing
 
-    | otherwise -> return Nothing
+    | otherwise -> do
+        atomically $ gLog $ show addr <> ": dropping packet " <> show header
+        return Nothing
 
   where
     addr = either id cAddress econn
     mbpid = listToMaybe $ mapMaybe (\case AnnounceSelf dgst -> Just dgst; _ -> Nothing) header
     version = listToMaybe $ filter (\v -> ProtocolVersion v `elem` header) protocolVersions
 
+channelCurrentCookie :: ChannelState -> Maybe Cookie
+channelCurrentCookie = \case
+    ChannelCookieReceived cookie -> Just cookie
+    ChannelCookieConfirmed cookie -> Just cookie
+    ChannelOurRequest mbcookie _ -> mbcookie
+    ChannelPeerRequest mbcookie _ -> mbcookie
+    ChannelOurAccept mbcookie _ _ -> mbcookie
+    _ -> Nothing
+
+cookieEchoReceived :: GlobalState addr -> Connection addr -> Maybe RefDigest -> Cookie -> STM ()
+cookieEchoReceived GlobalState {..} conn@Connection {..} mbpid cookieSet = do
+    readTVar cChannel >>= \case
+        ChannelNone -> newConn
+        ChannelCookieWait -> newConn
+        ChannelCookieReceived {} -> update
+        _ -> return ()
+  where
+    update = do
+        writeTVar cChannel $ ChannelCookieConfirmed cookieSet
+    newConn = do
+        update
+        writeFlow gControlFlow (NewConnection conn mbpid)
+
+generateCookieHeaders :: GlobalState addr -> addr -> ChannelState -> IO [TransportHeaderItem]
+generateCookieHeaders gs addr ch = catMaybes <$> sequence [ echoHeader, setHeader ]
+  where
+    echoHeader = return $ CookieEcho <$> channelCurrentCookie ch
+    setHeader = case ch of
+        ChannelCookieWait {} -> Just . CookieSet <$> createCookie gs addr
+        ChannelCookieReceived {} -> Just . CookieSet <$> createCookie gs addr
+        _ -> return Nothing
 
 createCookie :: GlobalState addr -> addr -> IO Cookie
 createCookie GlobalState {} addr = return (Cookie $ BC.pack $ show addr)
@@ -430,15 +473,16 @@ processOutgoing gs@GlobalState {..} = do
             when (isNothing mbch && secure) $ do
                 writeTQueue cSecureOutQueue (secure, packet, ackedBy)
 
-            header <- readTVar cChannel >>= return . TransportHeader . \case
-                ChannelCookieReceived cookie -> CookieEcho cookie : ProtocolVersion protocolVersion : hitems
-                _                            -> hitems
-
-            let plain = BL.concat $
-                    (serializeObject $ transportToObject gStorage header)
-                    : map lazyLoadBytes content
+            channel <- readTVar cChannel
 
             return $ do
+                cookieHeaders <- generateCookieHeaders gs cAddress channel
+                let header = TransportHeader $ cookieHeaders ++ hitems
+
+                let plain = BL.concat $
+                        (serializeObject $ transportToObject gStorage header)
+                        : map lazyLoadBytes content
+
                 mbs <- case mbch of
                     Just ch -> do
                         runExceptT (channelEncrypt ch $ BL.toStrict $ 0x00 `BL.cons` plain) >>= \case
diff --git a/src/Service.hs b/src/Service.hs
index a3a19a4..580c17d 100644
--- a/src/Service.hs
+++ b/src/Service.hs
@@ -98,7 +98,7 @@ data SomeStorageWatcher s = forall a. Eq a => SomeStorageWatcher (Stored LocalSt
 
 
 newtype ServiceID = ServiceID UUID
-    deriving (Eq, Ord, StorableUUID)
+    deriving (Eq, Ord, Show, StorableUUID)
 
 mkServiceID :: String -> ServiceID
 mkServiceID = maybe (error "Invalid service ID") ServiceID . U.fromString