Make "SRV" and "AAAA" lookups; wrap DNS `error' calls

pontarius-xmpp.cabal

@@ -38,6 +38,7 @@ Library
                , data-default      >=0.2
                , dns
                , hslogger          >=1.1.0
+               , iproute
                , lifted-base       >=0.1.0.1
                , mtl               >=2.0.0.0
                , network           >=2.3

source/Network/Xmpp/Stream.hs

@@ -46,6 +46,11 @@ import           Network.Xmpp.Utilities
 
 import           Network.DNS hiding (encode, lookup)
 
+import           Data.Ord
+import           Data.Maybe
+import           Data.List
+import           Data.IP
+
 
 -- import Text.XML.Stream.Elements
 
@@ -254,41 +259,10 @@ streamS expectedTo = do
 -- realm.
 openStream :: Text -> StreamConfiguration -> IO (Either XmppFailure (TMVar Stream))
 openStream realm config = runErrorT $ do
-    (address, port) <- case hardcodedTcpDetails config of
+    stream' <- createStream realm config
-        Nothing -> dnsLookup realm (resolvConf config)
-        Just (address, port) -> return (address, port)
-    stream' <- connectTcp (Text.unpack address) port realm config
     result <- liftIO $ withStream startStream stream'
     return stream'
 
-dnsLookup :: Text -> ResolvConf -> ErrorT XmppFailure IO (Text, PortID)
-dnsLookup realm resolvConf = ErrorT $ do
-    resolvSeed <- makeResolvSeed resolvConf
-    withResolver resolvSeed $ \resolver -> do
-        debugM "Pontarius.Xmpp" "Performing SRV lookup..."
-        srvResult <- lookupSRV resolver (BSC8.pack $ Text.unpack realm)
-        debugM "Pontarius.Xmpp" $ "SRV result: " ++ (show srvResult)
-
-        -- TODO: Use SRV result. Is list always empty?
-
-        -- TODO: Attempt to connect over IPv6 if it is resolvable.
-        -- TODO: Setting field to disable IPv6 lookup.
-        
-        -- aaaaResult <- lookupAAAA resolver (BSC8.pack $ Text.unpack realm)
-        -- debugM "Pontarius.Xmpp" $ "AAAA result: " ++ (show aaaaResult)
-        -- if isJust aaaaResult && (Prelude.length $ fromJust aaaaResult) > 0
-        --     then return $ Right (Text.pack $ show $ Prelude.head $ fromJust aaaaResult, (PortNumber 5222))
-        --    else
-
-        do
-                aResult <- lookupA resolver (BSC8.pack $ Text.unpack realm)
-                debugM "Pontarius.Xmpp" $ "A result: " ++ (show aResult)
-                case aResult of
-                    Nothing -> return $ Left DnsLookupFailed
-                    Just r | Prelude.length r == 0 -> return $ Left DnsLookupFailed
-                           -- Is it safe to ignore tail of A records?
-                           | otherwise -> return $ Right (Text.pack $ show $ Prelude.head r, (PortNumber 5222))
-
 -- | Send "</stream:stream>" and wait for the server to finish processing and to
 -- close the connection. Any remaining elements from the server are returned.
 -- Surpresses StreamEndFailure exceptions, but may throw a StreamCloseError.
@@ -431,45 +405,48 @@ xmppNoStream = Stream {
     zeroSource :: Source IO output
     zeroSource = liftIO . ExL.throwIO $ XmppOtherFailure "zeroSource"
 
-connectTcp :: HostName -> PortID -> Text -> StreamConfiguration -> ErrorT XmppFailure IO (TMVar Stream)
+createStream :: Text -> StreamConfiguration -> ErrorT XmppFailure IO (TMVar Stream)
-connectTcp host port hostname config = ErrorT $ do
+createStream realm config = do
-    let PortNumber portNumber = port
+    result <- connect realm config
-    debugM "Pontarius.Xmpp" $ "Connecting to " ++ host ++ " on port " ++
+    case result of
-        (show portNumber) ++ " through the realm " ++ (Text.unpack hostname) ++ "."
+        Just h -> ErrorT $ do
-    h <- connectTo host port
+            debugM "Pontarius.Xmpp" "Acquired handle."
-    debugM "Pontarius.Xmpp" "Setting NoBuffering mode on handle."
+            debugM "Pontarius.Xmpp" "Setting NoBuffering mode on handle."
-    hSetBuffering h NoBuffering
+            hSetBuffering h NoBuffering
-    let eSource = DCI.ResumableSource
+            let eSource = DCI.ResumableSource
                   ((sourceHandle h $= logConduit) $= XP.parseBytes def)
                   (return ())
-    let hand = StreamHandle { streamSend = \d -> do
+            let hand = StreamHandle { streamSend = \d -> do
-                                     debugM "Pontarius.Xmpp" $
+                                         debugM "Pontarius.Xmpp" $
-                                       "Sending TCP data: " ++ (BSC8.unpack d)
+                                           "Sending TCP data: " ++ (BSC8.unpack d)
-                                       ++ "."
+                                           ++ "."
-                                     catchPush $ BS.hPut h d
+                                         catchPush $ BS.hPut h d
-                                , streamReceive = \n -> do
+                                    , streamReceive = \n -> do
-                                     d <- BS.hGetSome h n
+                                         d <- BS.hGetSome h n
-                                     debugM "Pontarius.Xmpp" $
+                                         debugM "Pontarius.Xmpp" $
-                                       "Received TCP data: " ++
+                                           "Received TCP data: " ++
-                                       (BSC8.unpack d) ++ "."
+                                           (BSC8.unpack d) ++ "."
-                                     return d
+                                         return d
-                                , streamFlush = hFlush h
+                                    , streamFlush = hFlush h
-                                , streamClose = hClose h
+                                    , streamClose = hClose h
-                                }
+                                    }
-    let stream = Stream
+            let stream = Stream
-            { streamState = Plain
+                  { streamState = Plain
-            , streamHandle = hand
+                  , streamHandle = hand
-            , streamEventSource = eSource
+                  , streamEventSource = eSource
-            , streamFeatures = StreamFeatures Nothing [] []
+                  , streamFeatures = StreamFeatures Nothing [] []
-            , streamHostname = (Just hostname)
+                  , streamHostname = (Just realm)
-            , streamFrom = Nothing
+                  , streamFrom = Nothing
-            , streamId = Nothing
+                  , streamId = Nothing
-            , streamLang = Nothing
+                  , streamLang = Nothing
-            , streamJid = Nothing
+                  , streamJid = Nothing
-            , streamConfiguration = config
+                  , streamConfiguration = config
-            }
+                  }
-    stream' <- mkStream stream
+            stream' <- mkStream stream
-    return $ Right stream'
+            return $ Right stream'
+        Nothing -> do
+            lift $ debugM "Pontarius.Xmpp" "Did not acquire handle."
+            throwError TcpConnectionFailure
   where
     logConduit :: Conduit ByteString IO ByteString
     logConduit = CL.mapM $ \d -> do
@@ -478,6 +455,117 @@ connectTcp host port hostname config = ErrorT $ do
             "."
         return d
 
+-- `connect' will perform one or many DNS lookups for the provided
+-- realm host (unless `hardcodedTcpDetails' has been specified, in which case
+-- those details are used instead). Will return the Handle acquired, if any.
+connect :: Text -> StreamConfiguration -> ErrorT XmppFailure IO (Maybe Handle)
+connect realm config = do
+    case hardcodedTcpDetails config of
+        Just (address, port) -> lift $ do
+            debugM "Pontarius.Xmpp" "Connecting to hardcoded TCP settings..."
+            connectTcp' [(address, port)]
+        Nothing -> do
+            resolvSeed <- lift $ makeResolvSeed (resolvConf config)
+            lift $ debugM "Pontarius.Xmpp" "Performing SRV lookup..."
+            srvRecords <- srvLookup realm resolvSeed
+            case srvRecords of
+                -- No SRV records. Try fallback lookup.
+                Nothing -> do
+                    lift $ debugM "Pontarius.Xmpp" "No SRV records, using fallback process..."
+                    lift $ resolvAndConnectTcp resolvSeed (BSC8.pack $ Text.unpack realm) (fromIntegral 5222)
+                Just srvRecords' -> do
+                    lift $ debugM "Pontarius.Xmpp" "SRV records found, performing A/AAAA lookups..."
+                    lift $ resolvSrvsAndConnectTcp resolvSeed srvRecords'
+
+-- Connects to a list of addresses and ports. Surpresses any exceptions from
+-- connectTcp.
+connectTcp' :: [(Text, PortNumber)] -> IO (Maybe Handle)
+connectTcp' [] = return Nothing
+connectTcp' ((address, port):remainder) = do
+    result <- try $ connectTcp address port :: IO (Either IOException Handle)
+    case result of
+        Right handle -> do
+            debugM "Pontarius.Xmpp" "Successfully connected."
+            return $ Just handle
+        Left _ -> do
+            debugM "Pontarius.Xmpp" "Connection could not be established."
+            connectTcp' remainder
+    
+connectTcp :: Text -> PortNumber -> IO Handle
+connectTcp address port = do
+    debugM "Pontarius.Xmpp" $ "Connecting to " ++ (Text.unpack address) ++
+        " on port " ++ (show port) ++ "."
+    connectTo (Text.unpack address) (PortNumber port)
+
+-- Makes an AAAA query to acquire a IPs, and tries to connect to all of them. If
+-- a handle can not be acquired this way, an analogous A query is performed.
+-- Surpresses all IO exceptions.
+resolvAndConnectTcp :: ResolvSeed -> Domain -> PortNumber -> IO (Maybe Handle)
+resolvAndConnectTcp resolvSeed domain port = do
+    aaaaResults <- (try $ rethrowErrorCall $ withResolver resolvSeed $
+                        \resolver -> lookupAAAA resolver domain) :: IO (Either IOException (Maybe [IPv6]))
+    handle <- case aaaaResults of
+        Right Nothing -> return Nothing
+        Right (Just ipv6s) -> connectTcp' $ Data.List.map (\ipv6 -> (Text.pack $ show ipv6, port)) ipv6s
+        Left e -> return Nothing
+    case handle of
+        Nothing -> do
+            aResults <- (try $ rethrowErrorCall $ withResolver resolvSeed $
+                             \resolver -> lookupA resolver domain) :: IO (Either IOException (Maybe [IPv4]))
+            handle' <- case aResults of
+                Right Nothing -> return Nothing
+                Right (Just ipv4s) -> connectTcp' $ Data.List.map (\ipv4 -> (Text.pack $ show ipv4, port)) ipv4s
+            case handle' of
+                Nothing -> return Nothing
+                Just handle'' -> return $ Just handle''
+        Just handle' -> return $ Just handle'
+
+-- Tries `resolvAndConnectTcp' for every SRV record, stopping if a handle is
+-- acquired.
+resolvSrvsAndConnectTcp :: ResolvSeed -> [(Domain, PortNumber)] -> IO (Maybe Handle)
+resolvSrvsAndConnectTcp _ [] = return Nothing
+resolvSrvsAndConnectTcp resolvSeed ((domain, port):remaining) = do
+    result <- resolvAndConnectTcp resolvSeed domain port
+    case result of
+        Just handle -> return $ Just handle
+        Nothing -> resolvSrvsAndConnectTcp resolvSeed remaining
+  
+
+-- The DNS functions may make error calls. This function catches any such
+-- exceptions and rethrows them as IOExceptions.
+rethrowErrorCall :: IO a -> IO a
+rethrowErrorCall action = do
+    result <- try action
+    case result of
+        Right result' -> return result'
+        Left (ErrorCall e) -> ioError $ userError $ "rethrowErrorCall: " ++ e
+        Left e -> throwIO e
+
+-- Provides a list of A(AAA) names and port numbers upon a successful
+-- DNS-SRV request, or `Nothing' if the DNS-SRV request failed.
+srvLookup :: Text -> ResolvSeed -> ErrorT XmppFailure IO (Maybe [(Domain, PortNumber)])
+srvLookup realm resolvSeed = ErrorT $ do
+    result <- try $ rethrowErrorCall $ withResolver resolvSeed $ \resolver -> do
+        srvResult <- lookupSRV resolver $ BSC8.pack $ "_xmpp-client._tcp." ++ (Text.unpack realm) ++ "."
+        case srvResult of
+            Just srvResult -> do
+                debugM "Pontarius.Xmpp" $ "SRV result: " ++ (show srvResult)
+                -- Get [(Domain, PortNumber)] of SRV request, if any.
+                return $ Just $ Prelude.map (\(_, _, port, domain) -> (domain, fromIntegral port)) $
+                    sortBy (comparing $ \(prio, _weight, _, _) -> prio) srvResult
+                -- TODO: Perform the `Weight' probability calculations of
+                -- <http://tools.ietf.org/html/rfc2782>.
+            -- The service is not available at this domain.
+            -- Sorts the records based on the priority value.
+            Just [(_, _, _, ".")] -> do
+                debugM "Pontarius.Xmpp" $ "\".\" SRV result returned."
+                return $ Just []
+            Nothing -> do
+                debugM "Pontarius.Xmpp" "No SRV result returned."
+                return Nothing
+    case result of
+        Right result' -> return $ Right result'
+        Left e -> return $ Left $ XmppIOException e
 
 -- Closes the connection and updates the XmppConMonad Stream state.
 -- killStream :: TMVar Stream -> IO (Either ExL.SomeException ())

source/Network/Xmpp/Types.hs

@@ -656,6 +656,9 @@ data XmppFailure = StreamErrorFailure StreamErrorInfo -- ^ An error XML stream
                                               -- constructor wraps the
                                               -- elements collected so
                                               -- far.
+                 | TcpConnectionFailure -- ^ All attempts to TCP
+                                        -- connect to the server
+                                        -- failed.
                  | DnsLookupFailed -- ^ An IP address to connect to could not be
                                    -- resolved.
                  | TlsError TLS.TLSError -- ^ An error occurred in the
@@ -1024,7 +1027,7 @@ data StreamConfiguration =
                           -- | By specifying these details, Pontarius XMPP will
                           -- connect to the provided address and port, and will
                           -- not perform a DNS look-up
-                        , hardcodedTcpDetails :: Maybe (Text, PortID)
+                        , hardcodedTcpDetails :: Maybe (Text, PortNumber)
                           -- | DNS resolver configuration
                         , resolvConf :: ResolvConf
                         }