1 files changed, 32 insertions, 25 deletions

diff --git a/src/Network/Tox/Onion/Transport.hs b/src/Network/Tox/Onion/Transport.hs
index 8c4df694..a521c39e 100644
--- a/src/Network/Tox/Onion/Transport.hs
+++ b/src/Network/Tox/Onion/Transport.hs
@@ -19,7 +19,7 @@ module Network.Tox.Onion.Transport
     ( parseOnionAddr
     , encodeOnionAddr
     , forwardOnions
-    , OnionToOwner(..)
+    , OnionDestination(..)
     , OnionMessage(..)
     , DataToRoute(..)
     , AnnounceResponse(..)
@@ -41,7 +41,7 @@ import Network.QueryResponse
 import Crypto.Tox hiding (encrypt,decrypt)
 import Network.Tox.NodeId
 import qualified Crypto.Tox as ToxCrypto
-import Network.Tox.DHT.Transport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo,DHTPublicKey)
+import Network.Tox.DHT.Transport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo,DHTPublicKey,asymNodeInfo)
 
 import Debug.Trace
 import Control.Arrow
@@ -87,11 +87,11 @@ deriving instance ( Show (f (AnnounceRequest, Nonce8))
                   , Show (f DataToRoute)
                   ) => Show (OnionMessage f)
 
-data OnionToOwner = OnionToOwner NodeInfo (ReturnPath N3)
-                  | OnionToMe SockAddr -- SockAddr is immediate peer in route
+data OnionDestination = OnionToOwner NodeInfo (ReturnPath N3) -- ^ Somebody else's path to us.
+                      | OnionDestination NodeInfo             -- ^ Our own onion-path.
     deriving Show
 
-onionKey :: OnionToOwner -> Maybe PublicKey
+onionKey :: OnionDestination -> Maybe PublicKey
 onionKey (OnionToOwner ni _) = Just $ id2key (nodeId ni)
 onionKey _                   = Nothing
 
@@ -120,7 +120,7 @@ instance Serialize (OnionMessage Encrypted) where
     put (OnionAnnounceResponse n8 n24 x) = putWord8 0x84 >> put n8 >> put n24 >> put x
     put (OnionToRouteResponse a)         = putWord8 0x86 >> putAliasedAssym a
 
-onionToOwner :: Assym a -> ReturnPath N3 -> SockAddr -> Either String OnionToOwner
+onionToOwner :: Assym a -> ReturnPath N3 -> SockAddr -> Either String OnionDestination
 onionToOwner assym ret3 saddr = do
     ni <- nodeInfo (key2id $ senderKey assym) saddr
     return $ OnionToOwner ni ret3
@@ -131,18 +131,18 @@ onion :: Sized msg =>
                ByteString
                -> SockAddr
                -> Get (Assym (Encrypted msg) -> t)
-               -> Either String (t, OnionToOwner)
+               -> Either String (t, OnionDestination)
 onion bs saddr getf = do (f,(assym,ret3)) <- runGet ((,) <$> getf <*> getOnionRequest) bs
                          oaddr <- onionToOwner assym ret3 saddr
                          return (f assym, oaddr)
 
 
-parseOnionAddr :: (ByteString, SockAddr) -> Either (OnionMessage Encrypted,OnionToOwner) (ByteString,SockAddr)
+parseOnionAddr :: (ByteString, SockAddr) -> Either (OnionMessage Encrypted,OnionDestination) (ByteString,SockAddr)
 parseOnionAddr (msg,saddr)
     | Just (typ,bs) <- B.uncons msg
     , let right = Right (msg,saddr)
           query = either (const right) Left
-          response = either (const right) (Left . (, OnionToMe saddr))
+          response = either (const right) (Left . \msg -> ( msg , replyAlias saddr msg ))
     = case typ of
     0x83 -> query $ onion bs saddr (pure OnionAnnounce)            -- Announce Request
     0x85 -> query $ onion bs saddr (OnionToRoute <$> getPublicKey) -- Onion Data Request
@@ -154,22 +154,28 @@ getOnionReply :: Word8 -> Get (OnionMessage Encrypted)
 getOnionReply 0x84 = OnionAnnounceResponse <$> get <*> get <*> get
 getOnionReply 0x86 = OnionToRouteResponse <$> getOnionAssym
 
+replyAlias :: SockAddr -> OnionMessage Encrypted -> OnionDestination
+replyAlias saddr (OnionAnnounceResponse  _ _ _)
+    = OnionDestination
+    $ either (error "replyAlias: bad protocol") id
+    $ nodeInfo zeroID saddr -- TODO OnionAnnounceResponse has no sender key
+replyAlias saddr (OnionToRouteResponse asym)
+    = OnionDestination $ asymNodeInfo saddr asym
+
 putOnionMsg :: OnionMessage Encrypted -> Put
 putOnionMsg (OnionAnnounce                a) = putOnionAssym 0x83 (return ()) a
 putOnionMsg (OnionToRoute pubkey          a) = putOnionAssym 0x85 (putPublicKey pubkey) a
 putOnionMsg (OnionAnnounceResponse n8 n24 x) = put (0x84 :: Word8) >> put n8 >> put n24 >> put x
 putOnionMsg (OnionToRouteResponse         a) = putOnionAssym 0x86 (return ()) a
 
-encodeOnionAddr :: (OnionMessage Encrypted,OnionToOwner) -> (ByteString, SockAddr)
+encodeOnionAddr :: (OnionMessage Encrypted,OnionDestination) -> (ByteString, SockAddr)
 encodeOnionAddr (msg,OnionToOwner ni p) = ( runPut $ putResponse (OnionResponse p msg)
                                           , nodeAddr ni )
-encodeOnionAddr (msg,OnionToMe a) = ( runPut (putOnionMsg msg), a)
+encodeOnionAddr (msg,OnionDestination a) = ( runPut (putOnionMsg msg), nodeAddr a) -- TODO: Construct (OnionRequest N0)?
 
-forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport -- HandleHi a -> IO a -> HandleLo a
+forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport
 forwardOnions crypto udp = udp { awaitMessage = forwardAwait crypto udp }
 
--- forMe :: HandleHi
--- forThem :: handleLo
 forwardAwait :: TransportCrypto -> UDPTransport -> HandleLo a -> IO a
 forwardAwait crypto udp kont = do
     fix $ \another -> do
@@ -520,26 +526,27 @@ instance Sized OnionData where
                              -- should be treated as variable sized.
             VarSize f   -> f dhtpk
 
-encrypt :: TransportCrypto -> OnionMessage Identity -> OnionToOwner -> (OnionMessage Encrypted, OnionToOwner)
-encrypt crypto msg rpath = ( transcode (encryptMessage crypto okey) msg
+encrypt :: TransportCrypto -> OnionMessage Identity -> OnionDestination -> (OnionMessage Encrypted, OnionDestination)
+encrypt crypto msg rpath = ( transcode ( (. (runIdentity . either id assymData))
+                                         . encryptMessage skey okey)
+                                       msg
                            , rpath)
     where
+        skey = transportSecret crypto
+
         -- The OnionToMe case shouldn't happen, but we'll use our own public
         -- key in this situation.
         okey = fromMaybe (transportPublic crypto) $ onionKey rpath
 
+
 encryptMessage :: Serialize a =>
-                  TransportCrypto -> PublicKey -> Nonce24 -> Either (Identity a) (Assym (Identity a)) -> Encrypted a
-encryptMessage crypto destKey n (Right a) = ToxCrypto.encrypt secret plain
- where
-    secret = computeSharedSecret (transportSecret crypto) destKey n
-    plain  = encodePlain $ runIdentity $ assymData a
-encryptMessage crypto destKey n (Left x) = ToxCrypto.encrypt secret plain
+                  SecretKey -> PublicKey -> Nonce24 -> a -> Encrypted a
+encryptMessage skey destKey n a = ToxCrypto.encrypt secret plain
  where
-    secret = computeSharedSecret (transportSecret crypto) destKey n
-    plain  = encodePlain $ runIdentity $ x
+    secret = computeSharedSecret skey destKey n
+    plain  = encodePlain a
 
-decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionToOwner -> Either String (OnionMessage Identity, OnionToOwner)
+decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionDestination -> Either String (OnionMessage Identity, OnionDestination)
 decrypt crypto msg addr = (, addr) <$> (sequenceMessage $ transcode (decryptMessage crypto) msg)
 
 decryptMessage :: Serialize x =>