Correctly isolate variable-width onion payloads.

2 files changed, 58 insertions, 26 deletions

diff --git a/OnionTransport.hs b/OnionTransport.hs
index 485f9d64..2a1003dc 100644
--- a/OnionTransport.hs
+++ b/OnionTransport.hs
@@ -43,23 +43,25 @@ import ToxAddress
 import qualified ToxCrypto
 import DHTTransport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo,DHTPublicKey)
 
+import Debug.Trace
 import Control.Arrow
 import Control.Concurrent.STM
-import qualified Data.ByteString as B
-         ;import Data.ByteString (ByteString)
+import qualified Data.ByteString  as B
+         ;import Data.ByteString  (ByteString)
 import Data.Coerce
+import Data.Function
 import Data.Functor.Contravariant
 import Data.Functor.Identity
 import Data.IP
 import Data.Maybe
 import Data.Monoid
-import Data.Serialize            as S
+import Data.Serialize             as S
+import Data.Type.Equality
 import Data.Typeable
 import Data.Word
+import GHC.Generics               ()
 import GHC.TypeLits
 import Network.Socket
-import GHC.Generics ()
-import Data.Type.Equality
 
 type HandleLo a = Maybe (Either String (ByteString, SockAddr)) -> IO a
 
@@ -163,25 +165,27 @@ encodeOnionAddr (msg,addr) = ( runPut (putOnionMsg msg >> putpath), saddr )
                     | OnionToMe a       <- addr = (a, return ())
 
 forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport -- HandleHi a -> IO a -> HandleLo a
-forwardOnions crypto udp = udp { awaitMessage = await' }
- where
-    -- forMe :: HandleHi
-    -- forThem :: handleLo
-    await' :: HandleLo a -> IO a
-    await' forThem = awaitMessage udp $ \case
+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
+      awaitMessage udp $ \case
         m@(Just (Right (bs,saddr))) -> case B.head bs of
-            0x80 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy  N0) crypto saddr udp (await' forThem)
-            0x81 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy  N1) crypto saddr udp (await' forThem)
-            0x82 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy  N2) crypto saddr udp (await' forThem)
-            0x8c -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy N3) crypto saddr udp (await' forThem)
-            0x8d -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy N2) crypto saddr udp (await' forThem)
-            0x8e -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy N1) crypto saddr udp (await' forThem)
-            _    -> forThem m
-        m -> forThem m
-
-forward :: forall c b b1. Serialize b =>
+            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
+            _    -> kont m
+        m -> kont m
+
+forward :: forall c b b1. (Serialize b, Show b) =>
               (Maybe (Either String b1) -> c) -> ByteString -> (b -> c) -> c
-forward forMe bs f = either (forMe . Just . Left) f $ decode $ B.tail bs
+forward kont bs f = either (kont . Just . Left) f $ decode $ B.tail bs
 
 class SumToThree a b
 
@@ -211,10 +215,21 @@ data OnionRequest n = OnionRequest
     , pathFromOwner :: ReturnPath n
     }
 
+deriving instance ( Show (Forwarding (ThreeMinus n) (OnionMessage Encrypted))
+                  , KnownNat (PeanoNat n)
+                  ) => Show (OnionRequest n)
+
 instance ( Serialize (Forwarding (ThreeMinus n) (OnionMessage Encrypted))
+         , Sized (ReturnPath n)
          , Serialize (ReturnPath n)
          ) => Serialize (OnionRequest n) where
-    get = OnionRequest <$> get <*> get <*> get
+    get = do
+        n24 <- get
+        cnt <- remaining
+        let fwdsize = case size :: Size (ReturnPath n) of ConstSize n -> cnt - n
+        fwd <- isolate fwdsize get
+        rpath <- get
+        return $ OnionRequest n24 fwd rpath
     put (OnionRequest n f p) = put n >> put f >> put p
 
 -- getRequest :: _
@@ -228,6 +243,8 @@ data OnionResponse n = OnionResponse
     , msgToOwner  :: OnionMessage Encrypted
     }
 
+deriving instance KnownNat (PeanoNat n) => Show (OnionResponse n)
+
 instance ( Serialize (ReturnPath n) ) => Serialize (OnionResponse n) where
     get = OnionResponse <$> get <*> (get >>= getOnionReply)
     put (OnionResponse p m) = put p >> putOnionMsg m
@@ -324,6 +341,17 @@ data Forwarding n msg where
     NotForwarded :: msg -> Forwarding N0 msg
     Forwarding :: Assym (Encrypted (Addressed (Forwarding n msg))) -> Forwarding (S n) msg
 
+instance Show msg => Show (Forwarding N0 msg) where
+    show (NotForwarded x) = "NotForwarded "++show x
+
+instance ( KnownNat (PeanoNat (S n))
+         , Show (Encrypted (Addressed (Forwarding n msg)))
+         ) => Show (Forwarding (S n) msg) where
+    show (Forwarding a) = unwords [ "Forwarding"
+                                  , "("++show (natVal (Proxy :: Proxy (PeanoNat (S n))))++")"
+                                  , show a
+                                  ]
+
 instance Sized msg => Sized (Forwarding N0 msg)
    where size = case size :: Size msg of
                  ConstSize n -> ConstSize n
@@ -342,13 +370,17 @@ instance (Serialize (Encrypted (Addressed (Forwarding n msg)))) => Serialize (Fo
     get = Forwarding <$> getAliasedAssym
     put (Forwarding x) = putAliasedAssym x
 
-handleOnionRequest :: LessThanThree n => proxy n -> TransportCrypto -> SockAddr -> UDPTransport -> IO a -> OnionRequest n -> IO a
+handleOnionRequest :: ( LessThanThree n
+                      , Sized (ReturnPath n)
+                      ) => proxy n -> TransportCrypto -> SockAddr -> UDPTransport -> IO a -> OnionRequest n -> IO a
 handleOnionRequest proxy crypto saddr udp kont (OnionRequest nonce msg rpath) = do
+    putStrLn "handleOnionRequest"
     (sym, snonce) <- atomically ( (,) <$> transportSymmetric crypto
                                       <*> transportNewNonce crypto )
     case peelOnion crypto msg of
         Left e -> do
             -- todo report encryption error
+            putStrLn $ "peelOnion: " ++ e
             kont
         Right (Addressed dst msg') -> do
             sendMessage udp dst (S.encode $ OnionRequest nonce msg' $ wrapSymmetric sym snonce saddr rpath)
diff --git a/ToxCrypto.hs b/ToxCrypto.hs
index 943cbea8..883aba28 100644
--- a/ToxCrypto.hs
+++ b/ToxCrypto.hs
@@ -137,9 +137,9 @@ instance Sized a => Serialize (Encrypted a) where
     put = putEncrypted
 
 instance Sized a => Sized (Encrypted a) where
-    size = case size of
+    size = case size :: Size a of
         ConstSize n -> ConstSize $ n + 16
-        VarSize f   -> VarSize $ \x -> f x + 16
+        VarSize _   -> VarSize $ \(Encrypted bs) -> B.length bs
 
 instance (Sized a, Sized b) => Sized (a,b) where
     size = case (size :: Size a, size :: Size b) of