1 files changed, 75 insertions, 18 deletions

diff --git a/src/Network/Tox/Onion/Transport.hs b/src/Network/Tox/Onion/Transport.hs
index 160b99f7..f6d9ca31 100644
--- a/src/Network/Tox/Onion/Transport.hs
+++ b/src/Network/Tox/Onion/Transport.hs
@@ -50,6 +50,7 @@ module Network.Tox.Onion.Transport
     , selectAlias
     , RouteId(..)
     , routeId
+    , rewrap
     ) where
 
 import Network.Address (fromSockAddr,toSockAddr,setPort,either4or6,sockAddrPort)
@@ -87,6 +88,8 @@ import qualified Text.ParserCombinators.ReadP as RP
 import Data.Hashable
 import DPut
 import DebugTag
+import Data.Word64Map (fitsInInt)
+import Data.Bits (shiftR,shiftL)
 
 type HandleLo a = Maybe (Either String (ByteString, SockAddr)) -> IO a
 
@@ -270,20 +273,20 @@ encodeOnionAddr crypto getRoute (msg,OnionDestination _ ni (Just rid)) = do
     return x
 
 
-forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport
-forwardOnions crypto udp = udp { awaitMessage = forwardAwait crypto udp }
+forwardOnions :: TransportCrypto -> UDPTransport -> (Int -> OnionResponse N1 -> IO ()) {- ^ TCP relay send -} -> UDPTransport
+forwardOnions crypto udp sendTCP = udp { awaitMessage = forwardAwait crypto udp sendTCP }
 
-forwardAwait :: TransportCrypto -> UDPTransport -> HandleLo a -> IO a
-forwardAwait crypto udp kont = do
+forwardAwait :: TransportCrypto -> UDPTransport -> (Int -> OnionResponse N1 -> IO ()) {- ^ TCP relay send -} -> HandleLo a -> IO a
+forwardAwait crypto udp sendTCP kont = do
     fix $ \another -> do
       awaitMessage udp $ \case
         m@(Just (Right (bs,saddr))) -> case B.head bs of
-            0x80 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N0) crypto saddr udp another
-            0x81 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N1) crypto saddr udp another
-            0x82 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N2) crypto saddr udp another
-            0x8c -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N3) crypto saddr udp another
-            0x8d -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N2) crypto saddr udp another
-            0x8e -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N1) crypto saddr udp another
+            0x80 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N0) crypto (Addressed saddr) udp another
+            0x81 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N1) crypto (Addressed saddr) udp another
+            0x82 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N2) crypto (Addressed saddr) udp another
+            0x8c -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N3) crypto saddr udp sendTCP another
+            0x8d -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N2) crypto saddr udp sendTCP another
+            0x8e -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N1) crypto saddr udp sendTCP another
             _    -> kont m
         m -> kont m
 
@@ -392,6 +395,7 @@ instance (Sized (ReturnPath n)) => Sized (OnionResponse (S n)) where
     size = contramap pathToOwner size <> contramap msgToOwner size
 
 data Addressed a = Addressed { sockAddr :: SockAddr, unaddressed :: a }
+                 | TCPIndex { tcpIndex :: Int, unaddressed :: a }
  deriving (Eq,Show)
 
 instance Sized a => Sized (Addressed a) where
@@ -419,9 +423,24 @@ putForwardAddr saddr = fromMaybe (return $ error "unsupported SockAddr family")
             IPv6 ip6 -> S.put (0x0a :: Word8) >> S.put ip6
         S.put port
 
+addrToIndex :: SockAddr -> Int
+addrToIndex (SockAddrInet6 _ _ (lo, hi, _, _) _) =
+    if fitsInInt (Proxy :: Proxy Word64)
+        then fromIntegral lo + (fromIntegral hi `shiftL` 32)
+        else fromIntegral lo
+addrToIndex _ = 0
+
+indexToAddr :: Int -> SockAddr
+indexToAddr x = SockAddrInet6 0 0 (fromIntegral x, fromIntegral (x `shiftR` 32),0,0) 0
+
 instance Serialize a => Serialize (Addressed a) where
-    get = Addressed <$> getForwardAddr <*> get
+    get = do saddr <- getForwardAddr
+             a <- get
+             case sockAddrPort saddr of
+                Just 0 -> return $ TCPIndex (addrToIndex saddr) a
+                _      -> return $ Addressed saddr a
     put (Addressed addr x) = putForwardAddr addr >> put x
+    put (TCPIndex  idx  x) = putForwardAddr (indexToAddr idx) >> put x
 
 data N0
 data S n
@@ -529,31 +548,55 @@ instance (Serialize (Encrypted (Addressed (Forwarding n msg)))) => Serialize (Fo
     get = Forwarding <$> getPublicKey <*> get
     put (Forwarding k x) = putPublicKey k >> put x
 
+rewrap :: (ThreeMinus n ~ S (ThreeMinus (S n)),
+                 Serialize (ReturnPath n),
+                 Serialize
+                   (Forwarding (ThreeMinus (S n)) (OnionMessage Encrypted))) =>
+                TransportCrypto
+                -> (forall x. x -> Addressed x)
+                -> OnionRequest n
+                -> IO (Either String (OnionRequest (S n), SockAddr))
+rewrap crypto saddr (OnionRequest nonce msg rpath) = do
+    (sym, snonce) <- atomically ( (,) <$> transportSymmetric crypto
+                                      <*> transportNewNonce crypto )
+    peeled <- peelOnion crypto nonce msg
+    return $ peeled >>= \case
+        Addressed dst msg'
+          -> Right (OnionRequest nonce msg' $ wrapSymmetric sym snonce saddr rpath, dst)
+        _ -> Left "Onion forward to TCP client?"
+
 handleOnionRequest :: forall a proxy n.
                       ( LessThanThree n
                       , KnownPeanoNat n
                       , Sized (ReturnPath n)
                       , Typeable n
-                      ) => proxy n -> TransportCrypto -> SockAddr -> UDPTransport -> IO a -> OnionRequest n -> IO a
+                      ) => proxy n -> TransportCrypto -> (forall x. x -> Addressed x) -> UDPTransport -> IO a -> OnionRequest n -> IO a
 handleOnionRequest proxy crypto saddr udp kont (OnionRequest nonce msg rpath) = do
     let n = peanoVal rpath
     dput XOnion $ "handleOnionRequest " ++ show n
     (sym, snonce) <- atomically ( (,) <$> transportSymmetric crypto
                                       <*> transportNewNonce crypto )
     peeled <- peelOnion crypto nonce msg
+    let showDestination = case saddr () of
+            Addressed a _ -> either show show $ either4or6 a
+            TCPIndex  i _ -> "TCP" ++ show [i]
+
     case peeled of
         Left e -> do
             -- todo report encryption error
-            dput XOnion $ unwords [ "peelOnion:", show n, either show show (either4or6 saddr), e]
+            dput XOnion $ unwords [ "peelOnion:", show n, showDestination, e]
             kont
         Right (Addressed dst msg') -> do
-            dput XOnion $ unwords [ "peelOnion:", show n, either show show (either4or6 saddr), "-->", either show show (either4or6 dst), "SUCCESS"]
+            dput XOnion $ unwords [ "peelOnion:", show n, showDestination, "-->", either show show (either4or6 dst), "SUCCESS"]
             sendMessage udp dst (runPut $ putRequest $ OnionRequest nonce msg' $ wrapSymmetric sym snonce saddr rpath)
             kont
+        Right (TCPIndex {}) -> do
+            dput XUnexpected "handleOnionRequest: Onion forward to TCP client?"
+            kont
 
 wrapSymmetric :: Serialize (ReturnPath n) =>
-                 SymmetricKey -> Nonce24 -> SockAddr -> ReturnPath n -> ReturnPath (S n)
-wrapSymmetric sym n saddr rpath = ReturnPath n $ encryptSymmetric sym n (encodePlain $ Addressed saddr rpath)
+                 SymmetricKey -> Nonce24 -> (forall x. x -> Addressed x) -> ReturnPath n -> ReturnPath (S n)
+wrapSymmetric sym n saddr rpath = ReturnPath n $ encryptSymmetric sym n (encodePlain $ saddr rpath)
 
 peelSymmetric :: Serialize (Addressed (ReturnPath n))
                  => SymmetricKey -> ReturnPath (S n) -> Either String (Addressed (ReturnPath n))
@@ -568,8 +611,16 @@ peelOnion :: Serialize (Addressed (Forwarding n t))
 peelOnion crypto nonce (Forwarding k fwd) = do
     fmap runIdentity . uncomposed <$> decryptMessage crypto (dhtKey crypto) nonce (Right $ Asymm k nonce fwd)
 
-handleOnionResponse :: (KnownPeanoNat n, Sized (ReturnPath n), Serialize (ReturnPath n)) => proxy (S n) -> TransportCrypto -> SockAddr -> UDPTransport -> IO a -> OnionResponse (S n) -> IO a
-handleOnionResponse proxy crypto saddr udp kont (OnionResponse path msg) = do
+handleOnionResponse :: (KnownPeanoNat n, Sized (ReturnPath n), Serialize (ReturnPath n), Typeable n) =>
+                        proxy (S n)
+                        -> TransportCrypto
+                        -> SockAddr
+                        -> UDPTransport
+                        -> (Int -> OnionResponse N1 -> IO ()) -- ^ TCP-relay onion send.
+                        -> IO a
+                        -> OnionResponse (S n)
+                        -> IO a
+handleOnionResponse proxy crypto saddr udp sendTCP kont (OnionResponse path msg) = do
     sym <- atomically $ transportSymmetric crypto
     case peelSymmetric sym path of
         Left e -> do
@@ -580,6 +631,12 @@ handleOnionResponse proxy crypto saddr udp kont (OnionResponse path msg) = do
         Right (Addressed dst path') -> do
             sendMessage udp dst (runPut $ putResponse $ OnionResponse path' msg)
             kont
+        Right (TCPIndex dst path') -> do
+            -- This should only occur for OnionResponse N1
+            case gcast (OnionResponse path' msg) of
+                Just supported -> sendTCP dst supported
+                Nothing        -> dput XUnexpected "handleOnionResponse: TCP-bound message not supported."
+            kont
 
 
 data AnnounceRequest = AnnounceRequest