2 files changed, 72 insertions, 1 deletions
diff --git a/Utility/SRV.hs b/Utility/SRV.hs
new file mode 100644
index 000000000..51d4360e2
--- /dev/null
+++ b/Utility/SRV.hs
@@ -0,0 +1,69 @@
+{- SRV record lookup
+ -
+ - Uses either the ADNS Haskell library, or if it's not installed,
+ - the host command.
+ -
+ - Copyright 2012 Joey Hess <joey@kitenet.net>
+ -
+ - Licensed under the GNU GPL version 3 or higher.
+ -}
+
+{-# LANGUAGE CPP #-}
+
+module Utility.SRV (
+	lookupSRV,
+) where
+
+import qualified Build.SysConfig
+import Utility.Process
+import Utility.Exception
+import Utility.PartialPrelude
+
+import Network
+import Data.Function
+import Data.List
+import Control.Applicative
+import Data.Maybe
+
+#ifdef WITH_ADNS
+import ADNS.Resolver
+import Data.Either
+#endif
+
+type HostPort = (HostName, PortID)
+
+{- Returns an ordered list, with highest priority hosts first.
+ -
+ - On error, returns an empty list. -}
+lookupSRV :: String -> IO [HostPort]
+#ifdef WITH_ADNS
+lookupSRV srv = initResolver [] $ \resolver -> do
+	r <- catchDefaultIO (Right []) $
+		resolveSRV resolver srv
+	return $ either (\_ -> []) id r
+#else
+lookupSRV = lookupSRVHost
+#endif
+
+lookupSRVHost :: String -> IO [HostPort]
+lookupSRVHost srv
+	| Build.SysConfig.host = catchDefaultIO [] $ 
+		parseSrvHost <$> readProcessEnv "host" ["-t", "SRV", "--", srv]
+			-- clear environment, to avoid LANG affecting output
+			(Just [])
+	| otherwise = return []
+
+parseSrvHost :: String -> [HostPort]
+parseSrvHost = map snd . reverse . sortBy priority . catMaybes . map parse . lines
+	where
+		priority = compare `on` fst
+		parse l = case words l of
+			[_, _, _, _, priority, weight, sport, hostname] ->
+				case PortNumber . fromIntegral <$> readish sport of
+					Nothing -> Nothing
+					Just port -> Just
+						( (priority, weight)
+						, (hostname, port)
+						)
+			_ -> Nothing
+
diff --git a/Utility/State.hs b/Utility/State.hs
index c27f3c261..7f8919082 100644
--- a/Utility/State.hs
+++ b/Utility/State.hs
@@ -5,9 +5,11 @@
  - Licensed under the GNU GPL version 3 or higher.
  -}
 
+{-# LANGUAGE PackageImports #-}
+
 module Utility.State where
 
-import Control.Monad.State.Strict
+import "mtl" Control.Monad.State.Strict
 
 {- Modifies Control.Monad.State's state, forcing a strict update.
  - This avoids building thunks in the state and leaking.