2 files changed, 61 insertions, 23 deletions

diff --git a/OnionTransport.hs b/OnionTransport.hs
index 2a1003dc..96f9443f 100644
--- a/OnionTransport.hs
+++ b/OnionTransport.hs
@@ -46,6 +46,7 @@ import DHTTransport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo,DHTPublicKey)
 import Debug.Trace
 import Control.Arrow
 import Control.Concurrent.STM
+import Control.Monad
 import qualified Data.ByteString  as B
          ;import Data.ByteString  (ByteString)
 import Data.Coerce
@@ -62,6 +63,7 @@ import Data.Word
 import GHC.Generics               ()
 import GHC.TypeLits
 import Network.Socket
+import System.IO
 
 type HandleLo a = Maybe (Either String (ByteString, SockAddr)) -> IO a
 
@@ -159,10 +161,9 @@ putOnionMsg (OnionAnnounceResponse n8 n24 x) = put (0x84 :: Word8) >> put n8 >>
 putOnionMsg (OnionToRouteResponse         a) = putOnionAssym 0x86 (return ()) a
 
 encodeOnionAddr :: (OnionMessage Encrypted,OnionToOwner) -> (ByteString, SockAddr)
-encodeOnionAddr (msg,addr) = ( runPut (putOnionMsg msg >> putpath), saddr )
- where
-    (saddr,putpath) | OnionToOwner ni p <- addr = (nodeAddr ni, put p)
-                    | OnionToMe a       <- addr = (a, return ())
+encodeOnionAddr (msg,OnionToOwner ni p) = ( runPut $ putResponse (OnionResponse p msg)
+                                          , nodeAddr ni )
+encodeOnionAddr (msg,OnionToMe a) = ( runPut (putOnionMsg msg), a)
 
 forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport -- HandleHi a -> IO a -> HandleLo a
 forwardOnions crypto udp = udp { awaitMessage = forwardAwait crypto udp }
@@ -224,6 +225,7 @@ instance ( Serialize (Forwarding (ThreeMinus n) (OnionMessage Encrypted))
          , Serialize (ReturnPath n)
          ) => Serialize (OnionRequest n) where
     get = do
+        -- TODO share code with 'getOnionRequest'
         n24 <- get
         cnt <- remaining
         let fwdsize = case size :: Size (ReturnPath n) of ConstSize n -> cnt - n
@@ -288,6 +290,14 @@ type N1 = S N0
 type N2 = S N1
 type N3 = S N2
 
+class KnownPeanoNat n where
+    peanoVal :: p n -> Int
+
+instance KnownPeanoNat N0 where
+    peanoVal _ = 0
+instance KnownPeanoNat n => KnownPeanoNat (S n) where
+    peanoVal _ = 1 + peanoVal (Proxy :: Proxy n)
+
 type family PeanoNat p where
     PeanoNat N0    = 0
     PeanoNat (S n) = 1 + PeanoNat n
@@ -339,7 +349,7 @@ instance KnownNat (PeanoNat n) => Show (ReturnPath n) where
 
 data Forwarding n msg where
     NotForwarded :: msg -> Forwarding N0 msg
-    Forwarding :: Assym (Encrypted (Addressed (Forwarding n msg))) -> Forwarding (S n) msg
+    Forwarding :: PublicKey -> Encrypted (Addressed (Forwarding n msg)) -> Forwarding (S n) msg
 
 instance Show msg => Show (Forwarding N0 msg) where
     show (NotForwarded x) = "NotForwarded "++show x
@@ -347,10 +357,11 @@ instance Show msg => Show (Forwarding N0 msg) where
 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
-                                  ]
+    show (Forwarding k a) = unwords [ "Forwarding"
+                                    , "("++show (natVal (Proxy :: Proxy (PeanoNat (S n))))++")"
+                                    , show (key2id k)
+                                    , show a
+                                    ]
 
 instance Sized msg => Sized (Forwarding N0 msg)
    where size = case size :: Size msg of
@@ -358,32 +369,36 @@ instance Sized msg => Sized (Forwarding N0 msg)
                  VarSize f -> VarSize $ \(NotForwarded x) -> f x
 
 instance Sized (Forwarding n msg) => Sized (Forwarding (S n) msg)
-   where size = case size :: Size (Assym (Encrypted (Addressed (Forwarding n msg)))) of
-                 ConstSize n -> ConstSize n
-                 VarSize f -> VarSize $ \(Forwarding a) -> f a
+   where size = ConstSize 32
+                <> contramap (\(Forwarding _ e) -> e)
+                             (size :: Size (Encrypted (Addressed (Forwarding n msg))))
 
 instance Serialize msg => Serialize (Forwarding N0 msg) where
     get = NotForwarded <$> get
     put (NotForwarded msg) = put msg
 
 instance (Serialize (Encrypted (Addressed (Forwarding n msg)))) => Serialize (Forwarding (S n) msg) where
-    get = Forwarding <$> getAliasedAssym
-    put (Forwarding x) = putAliasedAssym x
+    get = Forwarding <$> getPublicKey <*> get
+    put (Forwarding k x) = putPublicKey k >> put x
 
-handleOnionRequest :: ( LessThanThree n
+handleOnionRequest :: forall a proxy n.
+                      ( LessThanThree n
+                      , KnownPeanoNat 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"
+    let n = peanoVal rpath
+    hPutStrLn stderr $ "handleOnionRequest " ++ show n
     (sym, snonce) <- atomically ( (,) <$> transportSymmetric crypto
                                       <*> transportNewNonce crypto )
-    case peelOnion crypto msg of
+    case peelOnion crypto nonce msg of
         Left e -> do
             -- todo report encryption error
-            putStrLn $ "peelOnion: " ++ e
+            hPutStrLn stderr $ unwords [ "peelOnion:", show n, either show show (either4or6 saddr), e]
             kont
         Right (Addressed dst msg') -> do
-            sendMessage udp dst (S.encode $ OnionRequest nonce msg' $ wrapSymmetric sym snonce saddr rpath)
+            hPutStrLn stderr $ unwords [ "peelOnion:", show n, either show show (either4or6 saddr), "SUCCESS"]
+            sendMessage udp dst (runPut $ putRequest $ OnionRequest nonce msg' $ wrapSymmetric sym snonce saddr rpath)
             kont
 
 wrapSymmetric :: Serialize (ReturnPath n) =>
@@ -397,20 +412,23 @@ peelSymmetric sym (ReturnPath nonce e) = decryptSymmetric sym nonce e >>= decode
 
 peelOnion :: Serialize (Addressed (Forwarding n t))
               => TransportCrypto
+              -> Nonce24
               -> Forwarding (S n) t
               -> Either String (Addressed (Forwarding n t))
-peelOnion crypto (Forwarding fwd) =
-    fmap runIdentity $ uncomposed $ decryptMessage crypto (assymNonce fwd) (Right fwd)
+peelOnion crypto nonce (Forwarding k fwd) =
+    fmap runIdentity $ uncomposed $ decryptMessage crypto nonce (Right $ Assym k nonce fwd)
 
-handleOnionResponse :: Serialize (ReturnPath n) => proxy (S n) -> TransportCrypto -> SockAddr -> UDPTransport -> IO a -> OnionResponse (S n) -> IO a
+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
     sym <- atomically $ transportSymmetric crypto
     case peelSymmetric sym path of
         Left e -> do
             -- todo report encryption error
+            let n = peanoVal path
+            hPutStrLn stderr $ unwords [ "peelSymmetric:", show n, either show show (either4or6 saddr), e]
             kont
         Right (Addressed dst path') -> do
-            sendMessage udp dst (S.encode $ OnionResponse path' msg)
+            sendMessage udp dst (runPut $ putResponse $ OnionResponse path' msg)
             kont
 
 
@@ -419,6 +437,7 @@ data AnnounceRequest = AnnounceRequest
     , announceSeeking :: NodeId  -- Public key we are searching for
     , announceKey     :: NodeId  -- Public key that we want those sending back data packets to use
     }
+ deriving Show
 
 instance Sized AnnounceRequest where size = ConstSize (32*3)
 
@@ -436,9 +455,24 @@ getOnionRequest = do
     path <- get
     return (a,path)
 
+putRequest :: (KnownPeanoNat n, Serialize (OnionRequest n)) => OnionRequest n -> Put
+putRequest req = do
+    putWord8 $ 0x80 + fromIntegral (peanoVal req)
+    put req
+
+putResponse :: (KnownPeanoNat n, Serialize (OnionResponse n)) => OnionResponse n -> Put
+putResponse resp = do
+    let tag = 0x8f - fromIntegral (peanoVal resp)
+    -- OnionResponse N0 is an alias for the OnionMessage Encrypted type which includes a tag
+    -- in it's Serialize instance.
+    when (tag /= 0x8f) (putWord8 tag)
+    put resp
+
+
 data KeyRecord = NotStored    Nonce32
                | SendBackKey  PublicKey
                | Acknowledged Nonce32
+ deriving Show
 
 instance Sized KeyRecord where size = ConstSize 33
 
@@ -457,6 +491,7 @@ data AnnounceResponse = AnnounceResponse
     { is_stored     :: KeyRecord
     , announceNodes :: SendNodes
     }
+ deriving Show
 
 instance Sized AnnounceResponse where
     size = contramap is_stored size <> contramap announceNodes size
diff --git a/ToxCrypto.hs b/ToxCrypto.hs
index 883aba28..7797da70 100644
--- a/ToxCrypto.hs
+++ b/ToxCrypto.hs
@@ -281,6 +281,9 @@ instance Show Nonce8 where
 newtype Nonce32 = Nonce32 ByteString
  deriving (Eq, Ord, ByteArrayAccess, Data)
 
+instance Show Nonce32 where
+    showsPrec d nonce = quoted (mappend $ bin2hex nonce)
+
 instance Serialize Nonce32 where
     get = Nonce32 <$> getBytes 32
     put (Nonce32 bs) = putByteString bs