IdentityUpdate updates self and requires ACK

2 files changed, 28 insertions, 18 deletions

diff --git a/src/Identity.hs b/src/Identity.hs
index 0e3e318..b81228f 100644
--- a/src/Identity.hs
+++ b/src/Identity.hs
@@ -23,6 +23,7 @@ import Control.Monad
 import Control.Monad.Except
 import qualified Control.Monad.Identity as I
 
+import Data.Either
 import Data.Foldable
 import Data.Function
 import Data.List
@@ -214,15 +215,23 @@ toComposedIdentity idt = idt { idData_ = toList $ idDataF idt
                              }
 
 
-updateIdentitySets :: Foldable m => [(Stored (Signed IdentityData), Set (Stored (Signed IdentityData)))] -> Identity m -> ComposedIdentity
-updateIdentitySets updates orig@Identity { idData_ = idata } =
-    case validateIdentityF $ map update $ toList idata of
-         Just updated -> updated { idOwner_ = updateIdentitySets updates <$> idOwner_ updated }
-         Nothing -> toComposedIdentity orig
-    where update x = foldl (\y (y', set) -> if y `S.member` set then y' else y) x updates
-
 updateIdentity :: Foldable m => [Stored (Signed IdentityData)] -> Identity m -> ComposedIdentity
-updateIdentity = updateIdentitySets . map (\u -> (u, ancestors [u]))
+updateIdentity [] orig = toComposedIdentity orig
+updateIdentity updates orig =
+    case validateIdentityF $ filterAncestors (ourUpdates ++ idata) of
+         -- need to filter ancestors here as validateIdentityF currently stores the whole list in idData_
+         Just updated -> updated
+             { idOwner_ = updateIdentity ownerUpdates <$> idOwner_ updated
+             , idUpdates_ = ownerUpdates
+             }
+         Nothing -> toComposedIdentity orig
+    where idata = toList $ idData_ orig
+          ilen = length idata
+          (ourUpdates, ownerUpdates) = partitionEithers $ flip map (filterAncestors $ updates ++ idUpdates_ orig) $
+              -- if an update is related to anything in idData_, use it here, otherwise push to owners
+              \u -> if length (filterAncestors (u : idata)) < ilen + 1
+                       then Left u
+                       else Right u
 
 updateOwners :: [Stored (Signed IdentityData)] -> Identity m -> Identity m
 updateOwners updates orig@Identity { idOwner_ = Just owner, idUpdates_ = cupdates } =
diff --git a/src/Network.hs b/src/Network.hs
index 01caef7..c9a2d8b 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -294,15 +294,16 @@ startServer opt origHead logd' services = do
             let announceUpdate identity = do
                     st <- derivePartialStorage storage
                     let selfRef = partialRef st $ storedRef $ idData identity
-                        updateRefs = map (partialRef st . storedRef) $ idUpdates identity
-                        hitems = AnnounceSelf selfRef : map AnnounceUpdate updateRefs
+                        updateRefs = selfRef : map (partialRef st . storedRef) (idUpdates identity)
+                        ackedBy = concat [[ Acknowledged r, Rejected r, DataRequest r ] | r <- updateRefs ]
+                        hitems = map AnnounceUpdate updateRefs
                         packet = TransportPacket (TransportHeader $  hitems) []
 
                     ps <- readMVar peers
                     forM_ ps $ \peer -> atomically $ do
                         ((,) <$> readTVar (peerIdentityVar peer) <*> readTVar (peerChannel peer)) >>= \case
                             (PeerIdentityFull _, ChannelEstablished _) ->
-                                writeTQueue (peerOutQueue peer) (True, [], packet)
+                                writeTQueue (peerOutQueue peer) (True, ackedBy, packet)
                             _  -> return ()
 
             let shareState self shared peer = do
@@ -591,7 +592,7 @@ handlePacket origHead identity secure peer chanSvc svcs (TransportHeader headers
                             addHeader $ AnnounceSelf ref
                             writeTVarP (peerIdentityVar peer) $ PeerIdentityRef wref idwait
                             liftSTM $ writeTChan (serverChanPeer $ peerServer peer) peer
-                        _ -> addHeader $ Acknowledged pref
+                        _ -> return ()
 
             AnnounceUpdate ref -> do
                 readTVarP (peerIdentityVar peer) >>= \case
@@ -716,9 +717,9 @@ finalizedChannel peer oh self = do
     -- Identity update
     do
         let selfRef = partialRef ist $ storedRef $ idData $ self
-            updateRefs = map (partialRef ist . storedRef) $ idUpdates $ self
-        sendToPeerS peer [] $ flip TransportPacket [] $ TransportHeader $
-            AnnounceSelf selfRef : map AnnounceUpdate updateRefs
+            updateRefs = selfRef : map (partialRef ist . storedRef) (idUpdates self)
+            ackedBy = concat [[ Acknowledged r, Rejected r, DataRequest r ] | r <- updateRefs ]
+        sendToPeerS peer ackedBy $ flip TransportPacket [] $ TransportHeader $ map AnnounceUpdate updateRefs
 
     -- Shared state
     readTVar (peerIdentityVar peer) >>= \case
@@ -741,7 +742,7 @@ handleIdentityAnnounce :: UnifiedIdentity -> Peer -> Ref -> WaitingRefCallback
 handleIdentityAnnounce self peer ref = liftIO $ atomically $ do
     let validateAndUpdate upds act = case validateIdentity $ wrappedLoad ref of
             Just pid' -> do
-                let pid = updateOwners upds pid'
+                let pid = fromMaybe pid' $ toUnifiedIdentity (updateIdentity upds pid')
                 writeTVar (peerIdentityVar peer) $ PeerIdentityFull pid
                 writeTChan (serverChanPeer $ peerServer peer) peer
                 act pid
@@ -766,9 +767,9 @@ handleIdentityUpdate :: Peer -> Ref -> WaitingRefCallback
 handleIdentityUpdate peer ref = liftIO $ atomically $ do
     pidentity <- readTVar (peerIdentityVar peer)
     if  | PeerIdentityFull pid <- pidentity
+        , Just pid' <- toUnifiedIdentity $ updateIdentity [wrappedLoad ref] pid
         -> do
-            writeTVar (peerIdentityVar peer) $ PeerIdentityFull $
-                updateOwners [wrappedLoad ref] pid
+            writeTVar (peerIdentityVar peer) $ PeerIdentityFull pid'
             writeTChan (serverChanPeer $ peerServer peer) peer
 
         | otherwise -> return ()