1 files changed, 38 insertions, 18 deletions

diff --git a/src/Data/Tox/Relay.hs b/src/Data/Tox/Relay.hs
index bd0e5968..f801d1cd 100644
--- a/src/Data/Tox/Relay.hs
+++ b/src/Data/Tox/Relay.hs
@@ -1,8 +1,9 @@
-{-# LANGUAGE DeriveDataTypeable #-}
-{-# LANGUAGE FlexibleInstances  #-}
-{-# LANGUAGE KindSignatures     #-}
-{-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE ConstraintKinds            #-}
+{-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE KindSignatures             #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
 module Data.Tox.Relay where
 
 import Data.ByteString as B
@@ -16,11 +17,30 @@ import qualified Rank2
 import Crypto.Tox
 import Network.Tox.Onion.Transport
 
+newtype ConId = ConId Word8
+    deriving (Eq,Show,Ord,Data,Serialize)
+
+badcon :: ConId
+badcon = ConId 0
+
+-- Maps to a range -120 .. 119
+c2key :: ConId -> Maybe Int
+c2key (ConId x) | x < 16    = Nothing
+                | otherwise = Just $ case divMod (x - 15) 2 of
+                                (q,0) -> negate $ fromIntegral q
+                                (q,1) -> fromIntegral q
+
+-- Maps to range 16 .. 255
+-- negatives become odds
+key2c :: Int -> ConId
+key2c y = ConId $ if y < 0 then 15 + fromIntegral (negate y * 2)
+                           else 16 + fromIntegral (y * 2)
+
 data RelayPacket
     = RoutingRequest PublicKey
-    | RoutingResponse Word8 PublicKey
-    | ConnectNotification Word8
-    | DisconnectNotification Word8
+    | RoutingResponse ConId PublicKey -- 0 for refusal, 16-255 for success.
+    | ConnectNotification ConId
+    | DisconnectNotification ConId
     | RelayPing Word64
     | RelayPong Word64
     | OOBSend PublicKey ByteString
@@ -28,12 +48,12 @@ data RelayPacket
     | OnionPacket (OnionRequest N0)
     | OnionPacketResponse (OnionResponse N1)
     -- 0x0A through 0x0F reserved for future use.
-    | RelayData Word8 ByteString -- Word8 is a connection id.  Encoded as number 16 to 255.
+    | RelayData ConId ByteString -- Word8 is a connection id.  Encoded as number 16 to 255.
  deriving (Eq,Ord,Show,Data)
 
 packetNumber :: RelayPacket -> Word8
-packetNumber (RelayData conid _) = conid -- 0 to 15 not allowed.
-packetNumber rp                  = fromIntegral $ pred $ constrIndex $ toConstr rp
+packetNumber (RelayData (ConId conid) _) = conid -- 0 to 15 not allowed.
+packetNumber rp                          = fromIntegral $ pred $ constrIndex $ toConstr rp
 
 instance Sized RelayPacket where
     size = mappend (ConstSize 1) $ VarSize $ \x -> case x of
@@ -59,24 +79,24 @@ instance Serialize RelayPacket where
         pktid <- getWord8
         case pktid of
             0 -> RoutingRequest <$> getPublicKey
-            1 -> RoutingResponse <$> getWord8 <*> getPublicKey
-            2 -> ConnectNotification <$> getWord8
-            3 -> DisconnectNotification <$> getWord8
+            1 -> RoutingResponse <$> get <*> getPublicKey
+            2 -> ConnectNotification <$> get
+            3 -> DisconnectNotification <$> get
             4 -> RelayPing <$> getWord64be
             5 -> RelayPong <$> getWord64be
             6 -> OOBSend <$> getPublicKey <*> (remaining >>= getBytes)
             7 -> OOBRecv <$> getPublicKey <*> (remaining >>= getBytes)
             8 -> OnionPacket <$> get
             9 -> OnionPacketResponse <$> get
-            conid -> RelayData conid <$> (remaining >>= getBytes)
+            conid -> RelayData (ConId conid) <$> (remaining >>= getBytes)
 
     put rp = do
         putWord8 $ packetNumber rp
         case rp of
             RoutingRequest k             -> putPublicKey k
-            RoutingResponse rpid k       -> putWord8 rpid >> putPublicKey k
-            ConnectNotification conid    -> putWord8 conid
-            DisconnectNotification conid -> putWord8 conid
+            RoutingResponse rpid k       -> put rpid >> putPublicKey k
+            ConnectNotification conid    -> put conid
+            DisconnectNotification conid -> put conid
             RelayPing pingid             -> putWord64be pingid
             RelayPong pingid             -> putWord64be pingid
             OOBSend k bs                 -> putPublicKey k >> putByteString bs