content: {-# LANGUAGE BangPatterns, OverloadedStrings #-}
  ------------------------------------------------------------
  -- Copyright : Erik de Castro Lopo <erikd@mega-nerd.com>
  -- License : BSD3
  ------------------------------------------------------------
  
  module Test.Util where
  
  import Blaze.ByteString.Builder
  import Control.Concurrent.Async
  import Control.Exception hiding (assert)
  import Control.Monad (forM_, when, unless)
  import Control.Monad.Trans.Resource (runResourceT)
  import Data.ByteString (ByteString)
  import Data.Conduit (ConduitT, Flush (..), SealedConduitT)
  import Data.Int (Int64)
  import Data.Maybe
  import Data.String (fromString)
  import Network.Socket
  import Network.Connection
  
  import qualified Data.ByteString.Char8 as BS
  import qualified Data.ByteString.Lazy.Char8 as LBS
  import qualified Data.CaseInsensitive as CI
  import qualified Data.Conduit as DC
  import qualified Data.Conduit.Binary as CB
  import qualified Network.HTTP.Conduit as HC
  import qualified Network.HTTP.Types as HT
  import qualified Network.Wai as Wai
  
  import Network.HTTP.Proxy.Request
  
  
  dumpWaiRequest :: Wai.Request -> IO ()
  dumpWaiRequest req =
      mapM_ BS.putStrLn
              [ "------- Wai Request --------------------------------------------------------------"
              , "Method          : " , Wai.requestMethod req
              , "HTTP Version    : " , fromString (show (Wai.httpVersion req))
              , "Path Info       : " , Wai.rawPathInfo req
              , "Query String    : " , Wai.rawQueryString req
              , "Server          : " , waiRequestHost req
              , "Secure (SSL)    : " , fromString (show (Wai.isSecure req))
              , "Remote Host     : " , fromString (show (Wai.remoteHost req))
              , "Request Headers :"
              , headerShow (Wai.requestHeaders req)
              ]
  
  
  dumpHttpConduitRequest :: HC.Request -> IO ()
  dumpHttpConduitRequest req =
      mapM_ BS.putStrLn
              [ "------- HttpConduit Request ------------------------------------------------------"
              , "Method          : " , HC.method req
              , "Secure (SSL)    : " , fromString (show (HC.secure req))
              , "Host Name       : " , HC.host req
              , "Host Port       : " , fromString (show (HC.port req))
              , "Path            : " , HC.path req
              , "Query String    : " , HT.urlDecode False (HC.queryString req)
              , "Request Headers :"
              , headerShow (HC.requestHeaders req)
              ]
  
  
  dumpHttpResponse :: HT.Status -> HT.ResponseHeaders -> IO ()
  dumpHttpResponse s rh = do
      mapM_ BS.putStrLn
          [ "------- Response from upsteam ----------------------------------------------------"
          , "HTTP/1.0 ", BS.pack (show (HT.statusCode s)), " ", HT.statusMessage s
          ]
      BS.putStr . BS.concat $ map (\ (f, v) -> BS.concat [ "    ", CI.original f, ": ", v, "\n" ]) rh
  
  
  headerShow :: [HT.Header] -> ByteString
  headerShow headers =
      BS.concat $ map hdrShow headers
    where
      hdrShow (f, v) = BS.concat [ "    ", CI.original f , ": " , v, "\n" ]
  
  
  --------------------------------------------------------------------------------
  
  simpleResponse :: HT.Status -> ByteString -> Wai.Response
  simpleResponse status text = do
      let respHeaders =
              [ (HT.hContentType, "text/plain")
              , (HT.hContentLength, fromString . show $ BS.length text)
              ]
      responseBS status respHeaders text
  
  
  responseBS :: HT.Status -> HT.ResponseHeaders -> ByteString -> Wai.Response
  responseBS status headers text =
      Wai.responseLBS status headers$ LBS.fromChunks [text]
  
  --------------------------------------------------------------------------------
  
  data Result = Result
      { resultSecure :: Bool
      , resultStatus :: Int
      , resultHeaders :: [HT.Header]
      , resultBS :: ByteString
      }
  
  
  printResult :: Result -> IO ()
  printResult (Result _ status headers body) = do
      putStrLn $ "Response status : " ++ show status
      putStrLn "Response headers :"
      BS.putStr $ headerShow headers
      putStrLn "Response body :"
      BS.putStrLn body
  
  --------------------------------------------------------------------------------
  
  -- | Compare results and error out if they're different.
  compareResult :: Result -> Result -> IO ()
  compareResult (Result secure sa ha ba) (Result _ sb hb bb) = do
      assert (sa == sb) $ "HTTP status codes don't match : " ++ show sa ++ " /= " ++ show sb
      forM_ [ "server", "content-type", "content-length" ] $ \v ->
          assertMaybe (lookup v ha) (lookup v hb) $ \ ja jb ->
              "Header field '" ++ show v ++ "' doesn't match : '" ++ show ja ++ "' /= '" ++ show jb
      assert (ba == bb) $ "HTTP response bodies are different :\n" ++ BS.unpack ba ++ "\n-----------\n" ++ BS.unpack bb
      when (not secure && isJust (lookup "X-Via-Proxy" ha)) $
          error "Error: Direct connection should not contain 'X-Via-Proxy' header."
      when (not secure && isNothing (lookup "X-Via-Proxy" hb)) $
          error "Error: Direct connection should not contain 'X-Via-Proxy' header."
    where
      assert :: Bool -> String -> IO ()
      assert b = unless b . error
  
      assertMaybe :: Eq a => Maybe a -> Maybe a -> (a -> a -> String) -> IO ()
      assertMaybe Nothing _ _ = return ()
      assertMaybe _ Nothing _ = return ()
      assertMaybe (Just a) (Just b) fmsg = unless (a == b) . error $ fmsg a b
  
  
  testSingleUrl :: Bool -> Int -> HC.Request -> IO ()
  testSingleUrl debug testProxyPort request = do
      direct <- httpRun request
      proxy <- httpRun $ addTestProxy testProxyPort request
      when debug $ do
          printResult direct
          printResult proxy
      compareResult direct proxy
  
  
  addTestProxy :: Int -> HC.Request -> HC.Request
  addTestProxy = HC.addProxy "localhost"
  
  
  -- | Use HC.http to fullfil a HC.Request. We need to wrap it because the
  -- Response contains a Source which we need to read to generate our result.
  httpRun :: HC.Request -> IO Result
  httpRun req = do
      mgr <- HC.newManager $ HC.mkManagerSettings (TLSSettingsSimple True False False) Nothing
      runResourceT $ do
          resp <- HC.http (modifyRequest req) mgr
          let contentLen = readInt64 <$> lookup HT.hContentLength (HC.responseHeaders resp)
          bodyText <- checkBodySize (DC.sealConduitT $ HC.responseBody resp) contentLen
          return $ Result (HC.secure req) (HT.statusCode $ HC.responseStatus resp)
                          (HC.responseHeaders resp) bodyText
    where
      modifyRequest r = r { HC.redirectCount = 0  }
  
  
  checkBodySize :: Monad f => SealedConduitT () ByteString f () -> Maybe Int64 -> f ByteString
  checkBodySize bodySrc Nothing = fmap (BS.concat . LBS.toChunks) $ bodySrc DC.$$+- CB.take 1000
  checkBodySize bodySrc (Just len) = do
      let blockSize = 1000
      if len <= blockSize
          then checkBodySize bodySrc Nothing
          else fromMaybe "Success" <$> (bodySrc DC.$$+- byteSink len)
  
  
  byteSink :: Monad m => Int64 -> ConduitT ByteString a m (Maybe ByteString)
  byteSink bytes = sink 0
    where
      sink :: Monad m => Int64 -> ConduitT ByteString a m (Maybe ByteString)
      sink !count = DC.await >>= maybe (closeSink count) (sinkBlock count)
  
      sinkBlock :: Monad m => Int64 -> ByteString -> ConduitT ByteString a m (Maybe ByteString)
      sinkBlock !count bs = sink (count + fromIntegral (BS.length bs))
  
      closeSink :: Monad m => Int64 -> m (Maybe ByteString)
      closeSink !count = return $
              if count == bytes
                  then Nothing
                  else Just . BS.pack $ "Error : Body length " ++ show count
                                      ++ " should have been " ++ show bytes ++ "."
  
  
  builderSource :: Monad m => Int64 -> ConduitT () (Flush Builder) m ()
  builderSource = DC.mapOutput (Chunk . fromByteString) . byteSource
  
  
  byteSource :: Monad m => Int64 -> ConduitT i ByteString m ()
  byteSource bytes = loop 0
    where
      loop :: Monad m => Int64 -> ConduitT i ByteString m ()
      loop !count
          | count >= bytes = return ()
          | count + blockSize64 < bytes = do
              DC.yield bsbytes
              loop $ count + blockSize64
          | otherwise = do
              let n = fromIntegral $ bytes - count
              DC.yield $ BS.take n bsbytes
              pure ()
  
      blockSize = 8192 :: Int
      blockSize64 = fromIntegral blockSize :: Int64
      bsbytes = BS.replicate blockSize '?'
  
  
  readInt64 :: ByteString -> Int64
  readInt64 = read . BS.unpack
  
  
  catchAny :: IO a -> (SomeException -> IO a) -> IO a
  catchAny action onE =
      withAsync action waitCatch >>= either onE return
  
  
  openLocalhostListenSocket :: IO (Socket, Port)
  openLocalhostListenSocket = do
      sock <- socket AF_INET Stream defaultProtocol
      addr:_ <- getAddrInfo Nothing (Just "127.0.0.1") Nothing
      bind sock (addrAddress addr)
      listen sock 10
      port <- fromIntegral <$> socketPort sock
      return (sock, port)
  
  