1 files changed, 44 insertions, 41 deletions

diff --git a/src/Network/Tox/Onion/Transport.hs b/src/Network/Tox/Onion/Transport.hs
index 34ba23f6..e52c8faa 100644
--- a/src/Network/Tox/Onion/Transport.hs
+++ b/src/Network/Tox/Onion/Transport.hs
@@ -20,6 +20,7 @@ module Network.Tox.Onion.Transport
     , parseDataToRoute
     , encodeDataToRoute
     , forwardOnions
+    , AliasSelector(..)
     , OnionDestination(..)
     , OnionMessage(..)
     , Rendezvous(..)
@@ -39,7 +40,7 @@ module Network.Tox.Onion.Transport
     , OnionRoute(..)
     , N3
     , onionKey
-    , onionNodeInfo
+    , onionAliasSelector
     ) where
 
 import Network.Address (fromSockAddr,toSockAddr,setPort,either4or6,sockAddrPort)
@@ -100,17 +101,27 @@ msgNonce (OnionAnnounceResponse _ n24 _) = n24
 msgNonce (OnionToRoute _ a)              = assymNonce a
 msgNonce (OnionToRouteResponse a)        = assymNonce a
 
+data AliasSelector = SearchingAlias | AnnouncingAlias NodeId
+    deriving (Eq,Ord,Show)
+
 data OnionDestination r
-    = OnionToOwner     NodeInfo (ReturnPath N3) -- ^ Somebody else's path to us.
-    | OnionDestination NodeInfo (Maybe r)       -- ^ Our own onion-path.
+    = OnionToOwner
+        { onionNodeInfo   :: NodeInfo
+        , onionReturnPath :: ReturnPath N3 -- ^ Somebody else's path to us.
+        }
+    | OnionDestination
+        { onionAliasSelector' :: AliasSelector
+        , onionNodeInfo       :: NodeInfo
+        , onionRouteSpec      :: Maybe r    -- ^ Our own onion-path.
+        }
     deriving Show
 
-onionNodeInfo :: OnionDestination r -> NodeInfo
-onionNodeInfo (OnionToOwner     ni _) = ni
-onionNodeInfo (OnionDestination ni _) = ni
+onionAliasSelector :: OnionDestination r -> AliasSelector
+onionAliasSelector (OnionToOwner    {}                         ) = SearchingAlias
+onionAliasSelector (OnionDestination{onionAliasSelector' = sel}) = sel
 
-onionKey :: OnionDestination r -> Maybe PublicKey
-onionKey od = Just $ id2key . nodeId $ onionNodeInfo od
+onionKey :: OnionDestination r -> PublicKey
+onionKey od = id2key . nodeId $ onionNodeInfo od
 
 instance Sized (OnionMessage Encrypted) where
     size = VarSize $ \case
@@ -152,7 +163,7 @@ onion bs saddr getf = do (f,(assym,ret3)) <- runGet ((,) <$> getf <*> getOnionRe
                          oaddr <- onionToOwner assym ret3 saddr
                          return (f assym, oaddr)
 
-parseOnionAddr :: (SockAddr -> Nonce8 -> IO (Maybe (NodeInfo,r)))
+parseOnionAddr :: (SockAddr -> Nonce8 -> IO (Maybe (OnionDestination r)))
                   -> (ByteString, SockAddr)
                   -> IO (Either (OnionMessage Encrypted,OnionDestination r)
                                 (ByteString,SockAddr))
@@ -167,11 +178,11 @@ parseOnionAddr lookupSender (msg,saddr)
             Just (Right msg@(OnionAnnounceResponse n8 _ _)) -> do
                 maddr <- lookupSender saddr n8
                 maybe (return right) -- Response unsolicited or too late.
-                      (return . Left . \(ni,r) -> (msg,OnionDestination ni (Just r)))
+                      (return . Left . \od -> (msg,od))
                       maddr
             Just (Right msg@(OnionToRouteResponse asym)) -> do
                 let ni = asymNodeInfo saddr asym
-                return $ Left (msg, OnionDestination ni Nothing)
+                return $ Left (msg, OnionDestination SearchingAlias ni Nothing)
             _ -> return right
 
 getOnionReply :: Word8 -> Maybe (Get (OnionMessage Encrypted))
@@ -191,10 +202,10 @@ encodeOnionAddr :: (NodeInfo -> r -> IO (Maybe OnionRoute))
 encodeOnionAddr _ (msg,OnionToOwner ni p) =
     return $ Just ( runPut $ putResponse (OnionResponse p msg)
                   , nodeAddr ni )
-encodeOnionAddr _        (msg,OnionDestination _  Nothing   ) = do
+encodeOnionAddr _ (msg,OnionDestination _ _ Nothing) = do
     hPutStrLn stderr $ "ONION encode missing routeid"
     return Nothing
-encodeOnionAddr getRoute (msg,OnionDestination ni (Just rid)) = do
+encodeOnionAddr getRoute (msg,OnionDestination _ ni (Just rid)) = do
     let go route = do
             return ( runPut $ putRequest $ wrapForRoute msg ni route
                    , nodeAddr $ routeNodeA route)
@@ -599,26 +610,19 @@ instance Sized OnionData where
             VarSize f   -> f req
 
 
-selectKey :: TransportCrypto -> OnionMessage Identity -> OnionDestination r -> IO (SecretKey, PublicKey)
-selectKey crypto
-          (OnionAnnounce a@Assym { assymData = Identity (AnnounceRequest _ pkey akey, _) })
-          rpath
-    | (akey /= zeroID) = atomically $ do
-                            ks <- filter (\(sk,pk) -> pk == id2key pkey)
-                                    <$> readTVar (userKeys crypto)
-                            maybe (return $ aliasKey crypto rpath)
-                                  return
-                                  (listToMaybe ks)
+selectKey :: TransportCrypto -> OnionMessage f -> OnionDestination r -> IO (SecretKey, PublicKey)
+selectKey crypto _ rpath@(OnionDestination (AnnouncingAlias pkey) _ _)
+    = atomically $ do ks <- filter (\(sk,pk) -> pk == id2key pkey)
+                                <$> readTVar (userKeys crypto)
+                      maybe (return $ aliasKey crypto rpath)
+                            return
+                            (listToMaybe ks)
 selectKey crypto msg rpath = return $ aliasKey crypto rpath
 
 encrypt :: TransportCrypto -> OnionMessage Identity -> OnionDestination r -> IO (OnionMessage Encrypted, OnionDestination r)
 encrypt crypto msg rpath = do
-    (skey,pkey) <- selectKey crypto msg rpath
-    let skey = fst $ aliasKey crypto rpath
-
-        -- The OnionToMe case shouldn't happen, but we'll use our own public
-        -- key in this situation.
-        okey = fromMaybe (transportPublic crypto) $ onionKey rpath
+    (skey,pkey) <- selectKey crypto msg rpath -- source key
+    let okey = onionKey rpath                 -- destination key
     return ( transcode ( (. (runIdentity . either id assymData))
                                          . encryptMessage skey okey)
                                        msg
@@ -632,11 +636,13 @@ encryptMessage skey destKey n a = ToxCrypto.encrypt secret plain
     plain  = encodePlain a
 
 decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionDestination r -> IO (Either String (OnionMessage Identity, OnionDestination r))
-decrypt crypto msg addr = return $ do
-    msg <- sequenceMessage $ transcode (\n -> decryptMessage (aliasKey crypto addr) n . left (senderkey addr)) msg
+decrypt crypto msg addr = do
+    (skey,pkey) <- selectKey crypto msg addr
+    return $ do
+    msg <- sequenceMessage $ transcode (\n -> decryptMessage (skey,pkey) n . left (senderkey addr)) msg
     Right (msg, addr)
 
-senderkey :: OnionDestination r -> t -> (Maybe PublicKey, t)
+senderkey :: OnionDestination r -> t -> (PublicKey, t)
 senderkey addr e = (onionKey addr, e)
 
 aliasKey :: TransportCrypto -> OnionDestination r -> (SecretKey,PublicKey)
@@ -649,17 +655,14 @@ dhtKey crypto = (transportSecret &&& transportPublic) crypto
 decryptMessage :: Serialize x =>
                   (SecretKey,PublicKey)
                   -> Nonce24
-                  -> Either (Maybe PublicKey, Encrypted x)
+                  -> Either (PublicKey, Encrypted x)
                             (Assym (Encrypted x))
                   -> (Either String ∘ Identity) x
-decryptMessage crypto n arg
-    | Just secret <- msecret = plain $ ToxCrypto.decrypt secret e
-    | otherwise              = Composed $ Left "decryptMessage: Unknown sender"
+decryptMessage crypto n arg = plain $ ToxCrypto.decrypt secret e
  where
-    msecret  = do sender <- mkey
-                  Just $ computeSharedSecret (fst crypto) sender n
-    (mkey,e) = either id (Just . senderKey &&& assymData) arg
-    plain    = Composed . fmap Identity . (>>= decodePlain)
+    secret     = computeSharedSecret (fst crypto) sender n
+    (sender,e) = either id (senderKey &&& assymData) arg
+    plain      = Composed . fmap Identity . (>>= decodePlain)
 
 
 sequenceMessage :: Applicative m => OnionMessage (m ∘ f) -> m (OnionMessage f)
@@ -751,4 +754,4 @@ encodeDataToRoute :: TransportCrypto
 encodeDataToRoute crypto (dta, Rendezvous pub ni)
     = Just ( OnionToRoute pub -- Public key of destination node
                           dta
-           , OnionDestination ni Nothing )
+           , OnionDestination SearchingAlias ni Nothing )