more robustness fixes

1 files changed, 15 insertions, 10 deletions

diff --git a/SybilLimit.hs b/SybilLimit.hs
index 9befb31..f2b983c 100644
--- a/SybilLimit.hs
+++ b/SybilLimit.hs
@@ -139,6 +139,7 @@ updateRoute i (k,j) fs = jfromk $ ktoj fs
     setFrom i x f = f { routesFrom = IntMap.insert i x (routesFrom f)}
 
 
+-- | Note: an exception will be thrown if srcId is not a valid friend node.
 forwardMessage :: ThisNode -> (FriendNode -> IntMap Int) -> (NodeId,PeerMessage) -> (NodeId, PeerMessage)
 forwardMessage me next (srcId,msg) = msg'
       where msg' = ( dest, msg { rrCount = rrCount msg + 1 } )
@@ -186,29 +187,33 @@ reactToMessage w (srcId, msg@(RegistrationQuery {})) me =
     resp = RegistrationResponse { rrSuspect = rrSuspect msg
                                 , rrRegisteredFrom = rrRegisteredFrom msg
                                 , rrValidRegistration = check }
-    key = friendById me Map.! rrRegisteredFrom msg
-    check = registeredTerminal (friends me ! key) == Just (rrSuspect msg)
+    check = fromMaybe False $ do
+        key <- Map.lookup (rrRegisteredFrom msg) $ friendById me
+        return $ registeredTerminal (friends me ! key) == Just (rrSuspect msg)
 
 reactToMessage w (srcId, msg) me =
-    if rrCount msg == w then terminate
-                        else MessageReaction me [forwardMessage me next (srcId,msg)] []
+    case Map.lookup srcId $ friendById me of
+        Nothing -> MessageReaction me [] []
+        Just srcNo -> 
+            if rrCount msg == w
+                then terminate srcNo
+                else MessageReaction me [forwardMessage me next (srcId,msg)] []
  where
     (terminate,next) = case msg of
                         ForwardRRMessage {} -> ( terminateForward, routesTo )
                         BackwardRRMessage {} -> ( terminateBackward, routesFrom )
 
-    terminateForward = MessageReaction me { friends = friends' }
-                                       [pong]
-                                       []
+    terminateForward srcNo = MessageReaction me { friends = friends' }
+                                             [pong]
+                                             []
       where pong =( srcId, BackwardRRMessage { rrCount = 1
                                              , rrRoute = rrRoute msg
                                              , rrTail = (srcId, selfId me)
                                              } )
-            friends' = IntMap.adjust adj suspectNo $ friends me
-            suspectNo = friendById me Map.! rrSuspect msg
+            friends' = IntMap.adjust adj srcNo $ friends me
             adj f = f { registeredTerminal = Just $ rrSuspect msg }
 
-    terminateBackward = MessageReaction me' [] []
+    terminateBackward _ = MessageReaction me' [] []
       where me' = me { routeTails =
                         IntMap.insert (rrRoute msg) (rrTail msg)
                                      $ routeTails me