1 files changed, 33 insertions, 20 deletions

diff --git a/Mainline.hs b/Mainline.hs
index 93393c40..ac870a43 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -144,11 +144,17 @@ hexdigit c = ('0' <= c && c <= '9') || ( 'a' <= c && c <= 'f') || ( 'A' <= c &&
 instance Read NodeInfo where
   readsPrec i = RP.readP_to_S $ do
     RP.skipSpaces
-    let n = 40
-    hexhash <- sequence $ replicate n (RP.satisfy hexdigit)
-    RP.char '@' RP.+++ RP.satisfy isSpace
-    addrstr <- RP.between (RP.char '(') (RP.char ')') (RP.munch (/=')'))
-                RP.+++ RP.munch (not . isSpace)
+    let n = 40 -- characters in node id.
+        parseAddr = RP.between (RP.char '(') (RP.char ')') (RP.munch (/=')'))
+                      RP.+++ RP.munch (not . isSpace)
+        nodeidAt = do hexhash <- sequence $ replicate n (RP.satisfy hexdigit)
+                      RP.char '@' RP.+++ RP.satisfy isSpace
+                      addrstr <- parseAddr
+                      nid <- case Base16.decode $ Char8.pack hexhash of
+                              (bs,_) | B.length bs==20 -> return (NodeId bs)
+                              _                        -> fail "Bad node id."
+                      return (nid,addrstr)
+    (nid,addrstr) <- ( nodeidAt RP.+++ ( (zeroID,) <$> parseAddr) )
     let raddr = do
             -- TODO: Support IPv6
             ipv4 <- RP.readS_to_P (readsPrec i)
@@ -159,9 +165,6 @@ instance Read NodeInfo where
     (ip,port) <- case RP.readP_to_S raddr addrstr of
                     [] -> fail "Bad address."
                     ((ip,port),_):_ -> return (ip,port)
-    nid <- case Base16.decode $ Char8.pack hexhash of
-            (bs,_) | B.length bs==20 -> return (NodeId bs)
-            _                        -> fail "Bad node id."
     return $ NodeInfo nid ip port
 
 
@@ -174,12 +177,13 @@ instance Hashable NodeInfo where
 
 
 instance Show NodeInfo where
-    show (NodeInfo (NodeId nid) ip port) =
-        Char8.unpack (Base16.encode nid) ++ "@" ++ show ip' ++ ":" ++ show port
+    showsPrec _ (NodeInfo nid ip port) =
+        shows nid . ('@' :) . showsip . (':' :) . shows port
      where
-        ip' | IPv6 ip6 <- ip
-            , Just ip4 <- un4map ip6 = IPv4 ip4
-            | otherwise              = ip
+        showsip
+            | IPv4 ip4 <- ip                          = shows ip4
+            | IPv6 ip6 <- ip , Just ip4 <- un4map ip6 = shows ip4
+            | otherwise                               = ('[' :) . shows ip . (']' :)
 
 {-
 
@@ -346,13 +350,22 @@ encodeMessage (R origin tid v ip)
         Right (BDict vals) -> encodeResponse tid (BDict $ genericArgs origin False `BE.union` vals) ip
         Left  err          -> encodeError tid err
 
+either4or6 :: SockAddr -> Either SockAddr SockAddr
+either4or6 a4@(SockAddrInet port addr)         = Left a4
+either4or6 a6@(SockAddrInet6 port _ addr _)
+    | Just ip4 <- (fromSockAddr a6 >>= un4map) = Left (setPort port $ toSockAddr ip4)
+    | otherwise                                = Right a6
+
+
 encodeAddr :: SockAddr -> ByteString
-encodeAddr (SockAddrInet port addr)
-    = S.runPut (S.putWord32host addr >> S.putWord16be (fromIntegral port))
-encodeAddr saddr@(SockAddrInet6 port _ addr _)
-    | Just ip4 <- (fromSockAddr saddr >>= un4map) = encodeAddr (setPort port $ toSockAddr ip4)
-    | otherwise                                   = S.runPut (S.put addr >> S.putWord16be (fromIntegral port))
-encodeAddr _  = B.empty
+encodeAddr = either encode4 encode6 . either4or6
+ where
+    encode4 (SockAddrInet port addr)
+              = S.runPut (S.putWord32host addr >> S.putWord16be (fromIntegral port))
+
+    encode6 (SockAddrInet6 port _ addr _)
+              = S.runPut (S.put addr >> S.putWord16be (fromIntegral port))
+    encode6 _ = B.empty
 
 decodeAddr :: ByteString -> Either String SockAddr
 decodeAddr bs = S.runGet g bs
@@ -393,7 +406,7 @@ showParseError bs addr err = L8.unpack $ L8.unlines es
  where
     es = map (L8.append prefix) (L8.pack err : L8.lines pp)
 
-    prefix = L8.pack (show addr) <> " --> "
+    prefix = L8.pack (either show show $ either4or6 addr) <> " --> "
 
     pp = either L8.pack showBEncode $ BE.decode bs