1 files changed, 34 insertions, 13 deletions

diff --git a/dht/src/Codec/AsciiKey256.hs b/dht/src/Codec/AsciiKey256.hs
index 6040f454..ee17b7c1 100644
--- a/dht/src/Codec/AsciiKey256.hs
+++ b/dht/src/Codec/AsciiKey256.hs
@@ -3,6 +3,7 @@ module Codec.AsciiKey256 where
 
 import Control.Applicative
 import Control.Monad
+import Data.Bits
 import qualified Data.ByteArray               as BA
          ;import Data.ByteArray               as BA (ByteArrayAccess)
 import qualified Data.ByteString              as B
@@ -17,6 +18,7 @@ import qualified Data.Text                    as T
          ;import Data.Text                    (Text)
 import Data.Word
 import Foreign.Ptr
+import System.IO.Unsafe
 import qualified Text.ParserCombinators.ReadP as RP
 
 stripSuffix :: Text -> Text -> Maybe Text
@@ -67,14 +69,33 @@ parseToken32 str = fmap (BA.drop 1) $ Base64.decode $ C8.pack $ 'A':map (nmtoken
 showToken32 :: ByteArrayAccess bin => bin -> String
 showToken32 bs = map (nmtoken64 True) $ C8.unpack $ BA.drop 1 $ Base64.encode $ BA.cons 0 $ BA.convert bs
 
+foreign import ccall "crc4itu" c_crc4itu :: Word8        -- ^ init crc
+                                            -> Ptr Word8 -- ^ data to checksum
+                                            -> Int32     -- ^ length of data
+                                            -> IO Word8  -- crc in low 4 bits
+
+-- | CRC4-ITU. Return crc in lowest 4 bits.
+crc4itu :: ByteArrayAccess ba => Word8 -- ^ Existing CRC value (usually 0x00) before process a new one.
+                                 -> ba -- ^ Data to checksum.
+                                 -> Word8
+crc4itu crc0 b = unsafePerformIO $ BA.withByteArray b $ \p ->
+    c_crc4itu crc0 p (fromIntegral $ BA.length b)
 
 -- | Parse 52-digit z-base32 token into 32-byte bytestring.
 parse32Token32 :: String -> Either String ByteString
-parse32Token32 str = fmap (BA.drop 1) $ Base32.decode $ C8.pack $ 'y':map (fixupDigit32 . toLower) (take 52 str)
-
--- | Encode 62-digit z-base32 token from 32-byte bytestring.
+parse32Token32 str = do
+    bs <- Base32.decode $ C8.pack $ map (fixupDigit32 . toLower) (take 52 str) ++ "y"
+    case BA.splitAt 32 bs of
+        (key,mac) | crc4itu 0 key == shiftR (BA.index mac 0) 4
+            -> Right key
+        _   -> Left "Failed cyclic redundancy check."
+
+-- | Encode 52-digit z-base32 token from 32-byte bytestring.
 show32Token32 :: ByteArrayAccess bin => bin -> String
-show32Token32 bs = map (nmtoken64 True) $ C8.unpack $ BA.drop 1 $ Base32.encode $ BA.cons 0 $ BA.convert bs
+show32Token32 bs = C8.unpack $ B.take 52 $ Base32.encode (b `B.snoc` shiftL crc 4)
+ where
+    b = BA.convert bs
+    crc = crc4itu 0 bs
 
 
 readsPrecKey256 :: (ByteString -> Maybe a) -> [Char] -> [(a, [Char])]
@@ -96,11 +117,11 @@ parseKey256 nidstr = do
         let nidbs  = C8.pack nidstr
             (bs,_) = Base16.decode nidbs
             enid   = case C8.length nidbs of
-                        52 -> Base32.decode (C8.pack $ 'y' : map (fixupDigit32 . toLower) nidstr)
-                        43 -> Base64.decode (C8.pack $ 'A' : map (nmtoken64 False) nidstr)
+                        52 -> parse32Token32 nidstr
+                        43 -> parseToken32 nidstr
                         _  -> Left "Wrong size of key."
         idbs <- (guard (B.length bs == 32) >> return bs)
-                <|> either fail (return . B.drop 1) enid
+                <|> either fail return enid
         return idbs
 
 readP_key256 :: RP.ReadP ByteString
@@ -112,12 +133,12 @@ readP_key256 = do
         ]
       let failure = fail "Bad key."
       case is64 of
-        32 -> case Base32.decode $ C8.pack $ 'y' : map (fixupDigit32 . toLower) hexhash of
-                  Right bs | B.length bs - 1==32 -> return (BA.drop 1 bs)
-                  _                              -> failure
-        64 -> case Base64.decode $ C8.pack $ 'A' : map (nmtoken64 False) hexhash of
-                  Right bs | B.length bs - 1==32 -> return (BA.drop 1 bs)
-                  _                              -> failure
+        32 -> case parse32Token32 hexhash of
+                  Right bs -> return bs
+                  _        -> failure
+        64 -> case parseToken32 hexhash of
+                  Right bs -> return bs
+                  _        -> failure
         16 -> case Base16.decode $ C8.pack hexhash of
                   (bs,rem) | B.length bs == 32 &&  B.null rem -> return bs
                   _                                           -> failure