tricu

Resolver.hs (3214B)

---

 {-# LANGUAGE ScopedTypeVariables #-}
 -- Explicit host name resolution, not an implicit part of connection acquisition.
 -- getaddrinfo may buffer more results and delay cancellation; only the published
 -- endpoint prefix is bounded here. No worker is abandoned on interruption.
 module IOV2.Resolver (Family(..), Resolution(..), resolveTCP, resolveWith, hostLimit, addressLimit, validHost) where
 import qualified IOV2.Resources as R
 import qualified Network.Socket as N
 import Control.Exception
 
 hostLimit, addressLimit :: Int
 hostLimit = 253
 addressLimit = 64
 data Family = IPv4Only | IPv6Only | AnyFamily deriving (Eq, Show)
 data Resolution = Resolution [R.Endpoint] Bool deriving (Eq, Show)
 validHost :: String -> Bool
 validHost [] = False
 validHost host = walk hostLimit host
   where walk _ [] = True
         walk 0 _ = False
         walk n (c:rest) = ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
           (c >= '0' && c <= '9') || c == '-' || c == '.') && walk (n-1) rest
 
 resolveTCP :: String -> Family -> Integer -> Int -> IO (Either R.Problem Resolution)
 resolveTCP host family port limit = resolveWith host port limit $ do
   let af = case family of IPv4Only -> N.AF_INET; IPv6Only -> N.AF_INET6; AnyFamily -> N.AF_UNSPEC
       hints = N.defaultHints {N.addrFamily=af, N.addrSocketType=N.Stream,
                              N.addrFlags=[N.AI_NUMERICSERV]}
   infos <- N.getAddrInfo (Just hints) (Just host) (Just (show port))
   pure [endpoint | info <- infos, Just endpoint <- [convert (N.addrAddress info)]]
   where convert (N.SockAddrInet p h) = Just (R.IPv4 (N.hostAddressToTuple h) (fromIntegral p))
         convert (N.SockAddrInet6 p _ h scope) = Just (R.IPv6 (N.hostAddress6ToTuple h) scope (fromIntegral p))
         convert _ = Nothing
 
 -- Injected query supports deterministic failure/bound tests without public DNS.
 resolveWith :: String -> Integer -> Int -> IO [R.Endpoint] -> IO (Either R.Problem Resolution)
 resolveWith host port limit query
   | not (validHost host) = pure (Left (R.Problem "invalid-argument" "resolve" "expected bounded ASCII host name" "none"))
   | port < 0 || port > 65535 || limit <= 0 || limit > addressLimit =
       pure (Left (R.Problem "invalid-bound" "resolve" "" "none"))
   | otherwise = (query >>= bounded limit [] >>= pure . Right) `catch` \(e :: IOException) ->
       pure (Left (R.Problem "resolution-failed" "resolve" (take 512 (displayException e)) "none"))
   where
     bounded 0 acc rest = do
       truncated <- evaluate (not (null rest))
       finish acc truncated
     bounded _ acc [] = finish acc False
     bounded n acc (first:rest) = forceEndpoint first `seq` bounded (n-1) (first:acc) rest
     finish acc truncated = do
       let prefix = reverse acc
       -- Neither a lazy split nor endpoint projections may retain the discarded
       -- host-result tail/AddrInfo records after returning the bounded reply.
       _ <- evaluate (foldr seq () prefix)
       pure (Resolution prefix truncated)
 
 forceEndpoint :: R.Endpoint -> ()
 forceEndpoint (R.IPv4 (a,b,c,d) p) = a `seq` b `seq` c `seq` d `seq` p `seq` ()
 forceEndpoint (R.IPv6 (a,b,c,d,e,f,g,h) scope p) =
   a `seq` b `seq` c `seq` d `seq` e `seq` f `seq` g `seq` h `seq` scope `seq` p `seq` ()