ProxySpec.hs

-- SPDX-License-Identifier: Apache-2.0
--
-- Copyright (C) 2026 Bin Jin. All Rights Reserved.

module Network.HProx.ProxySpec
  ( spec
  ) where

import Control.Concurrent              (forkIO)
import Control.Concurrent.Async        (withAsync)
import Control.Concurrent.MVar         (newEmptyMVar, putMVar, takeMVar)
import Control.Exception
    (AsyncException(..), IOException, SomeException, bracket, bracketOnError, finally, throwIO, try)
import Control.Monad                   (unless)
import Data.ByteString                 qualified as BS
import Data.ByteString.Base64          qualified as B64
import Data.ByteString.Char8           qualified as BS8
import Data.ByteString.Lazy            qualified as LBS
import Data.ByteString.Lazy.Char8      qualified as LBS8
import Data.CaseInsensitive            qualified as CI
import Data.IORef
import Data.Maybe                      (fromMaybe, mapMaybe)
import Data.Streaming.Network          qualified as CN
import Data.Streaming.Network.Internal (AppData(..))
import Network.HTTP.Client             qualified as HC
import Network.HTTP.Types              qualified as HT
import Network.HTTP.Types.Header       qualified as HT
import Network.Socket
import Network.Socket.ByteString       qualified as SocketBS
import Network.Wai
import Network.Wai.Internal            (Response(..), ResponseReceived(..))
import Network.Wai.Test
import System.Log.FastLogger           qualified as FL
import System.Timeout                  (timeout)

import Network.HProx.Impl

import Test.Hspec

spec :: Spec
spec = do
    describe "HTTP proxy auth decisions" $ do
        it "challenges unauthorized HTTP proxy requests" $ do
            response <- runProxyApp (httpProxy authRequiredSettings) rawProxyRequest
            simpleStatus response `shouldBe` HT.status407
            lookup "Proxy-Authenticate" (simpleHeaders response) `shouldBe` Just "Basic realm=\"hprox\""

        it "accepts valid Basic credentials" $
            withRawHttpServer (const "ok") $ \port -> do
                response <- runProxyApp (httpProxy authorizedSettings) (proxyRequestForPort port basicAliceSecret)
                simpleStatus response `shouldBe` HT.status200

        it "rejects valid Basic credentials with the wrong password" $ do
            response <- runProxyApp (httpProxy authorizedSettings) (rawProxyRequestWithAuth $ basicCredential "alice:wrong")
            simpleStatus response `shouldBe` HT.status407

        it "rejects malformed base64 credentials even if lenient decoding would authenticate" $ do
            response <- runProxyApp (httpProxy authorizedSettings) (rawProxyRequestWithAuth "Basic !!!YWxpY2U6c2VjcmV0!!!")
            simpleStatus response `shouldBe` HT.status407

        it "rejects non-Basic credentials even when the payload is otherwise valid" $ do
            response <- runProxyApp (httpProxy authorizedSettings) (rawProxyRequestWithAuth $ "Bearer " <> B64.encode "alice:secret")
            simpleStatus response `shouldBe` HT.status407

        it "redacts proxy auth passwords in trace logs" $
            withRawHttpServer (const "ok") $ \port -> do
                logs <- newIORef []
                let loggingSettings = authorizedSettings
                            { logger = \_ msg -> modifyIORef' logs (LBS.fromStrict (FL.fromLogStr msg) :)
                            }
                response <- runProxyApp (httpProxy loggingSettings) (proxyRequestForPort port basicAliceSecret)
                simpleStatus response `shouldBe` HT.status200
                renderedLogs <- LBS.toStrict . LBS.concat . reverse <$> readIORef logs
                renderedLogs `shouldSatisfy` BS8.isInfixOf "authorized request"
                renderedLogs `shouldSatisfy` BS8.isInfixOf "alice:<redacted>"
                renderedLogs `shouldNotSatisfy` BS8.isInfixOf "secret"
                renderedLogs `shouldNotSatisfy` BS8.isInfixOf "alice:secret"

        it "falls back instead of challenging when hidden auth is enabled" $ do
            response <- runProxyApp (httpProxy hiddenAuthSettings) rawProxyRequest
            simpleStatus response `shouldBe` fallbackStatus
            simpleBody response `shouldBe` fallbackBody

    describe "CONNECT proxy auth decisions" $ do
        it "challenges unauthorized CONNECT requests" $ do
            response <- runConnectApp authRequiredSettings connectRequest
            simpleStatus response `shouldBe` HT.status407
            lookup "Proxy-Authenticate" (simpleHeaders response) `shouldBe` Just "Basic realm=\"hprox\""

        it "falls back instead of challenging CONNECT requests when hidden auth is enabled" $ do
            response <- runConnectApp hiddenAuthSettings connectRequest
            simpleStatus response `shouldBe` fallbackStatus
            simpleBody response `shouldBe` fallbackBody

        it "establishes an HTTP/1 CONNECT tunnel after upstream connection succeeds" $ do
            response <- rawConnectRoundTrip
            response `shouldBe` "ping"

        it "returns 502 when upstream CONNECT fails before tunnel establishment" $ do
            port <- getFreePort
            response <- runConnectApp authorizedSettings (connectRequestFor "127.0.0.1" port HT.http11)
            simpleStatus response `shouldBe` HT.status502

        it "returns 502 for HTTP/2 CONNECT when upstream connection fails" $ do
            port <- getFreePort
            response <- runConnectApp authorizedSettings (connectRequestFor "127.0.0.1" port HT.http20)
            simpleStatus response `shouldBe` HT.status502

        it "does not log upstream failure after an HTTP/2 tunnel response starts" $
            withTcpEchoServer $ \targetPort -> do
                logs <- newIORef []
                let loggingSettings = authorizedSettings
                            { logger = \_ msg -> modifyIORef' logs (LBS.fromStrict (FL.fromLogStr msg) :)
                            }
                    tunnelRequest      = connectRequestFor "127.0.0.1" targetPort HT.http20
                    failAfterResponse _ = ioError (userError "client stream closed after connect")
                result <- try (httpConnectProxy loggingSettings fallback tunnelRequest failAfterResponse) :: IO (Either SomeException ResponseReceived)
                case result of
                    Left _  -> return ()
                    Right _ -> expectationFailure "response failure was swallowed"
                renderedLogs <- LBS.toStrict . LBS.concat . reverse <$> readIORef logs
                renderedLogs `shouldNotSatisfy` BS8.isInfixOf "CONNECT upstream failure"

        it "does not log upstream failure when HTTP/2 cancels before the tunnel response starts" $ do
            logs <- newIORef []
            let loggingSettings = authorizedSettings
                        { logger = \_ msg -> modifyIORef' logs (LBS.fromStrict (FL.fromLogStr msg) :)
                        }
                cancelledBeforeResponse _ _ = throwIO ThreadKilled
                respondOk _                 = ioError (userError "respond should not be called")
            result <- try (httpConnectProxyWith cancelledBeforeResponse loggingSettings fallback (connectRequestFor "127.0.0.1" 443 HT.http20) respondOk) :: IO (Either SomeException ResponseReceived)
            case result of
                Left _  -> return ()
                Right _ -> expectationFailure "pre-response cancellation was swallowed"
            renderedLogs <- LBS.toStrict . LBS.concat . reverse <$> readIORef logs
            renderedLogs `shouldNotSatisfy` BS8.isInfixOf "CONNECT upstream failure"

    describe "HTTP proxy fallback decisions" $
        it "falls through for non-proxy GET requests" $ do
            response <- runProxyApp (httpProxy authRequiredSettings) defaultRequest
            simpleStatus response `shouldBe` fallbackStatus
            simpleBody response `shouldBe` fallbackBody

    describe "HTTP proxy target parsing" $ do
        it "selects raw absolute-URI proxy targets with default ports" $
            selectHttpProxyTarget rawProxyRequest
                `shouldBe` Just HttpProxyTarget
                    { httpProxyTargetHost = "example.com"
                    , httpProxyTargetPort = 80
                    , httpProxyTargetRawPath = "/resource"
                    }

        it "preserves bracketed IPv6 proxy targets without explicit ports" $
            selectHttpProxyTarget rawProxyRequest { rawPathInfo = "http://[::1]/resource" }
                `shouldBe` Just HttpProxyTarget
                    { httpProxyTargetHost = "[::1]"
                    , httpProxyTargetPort = 80
                    , httpProxyTargetRawPath = "/resource"
                    }

        it "treats HTTP/2 proxy scheme values case-insensitively" $
            selectHttpProxyTarget http2ProxyRequest { requestHeaders = [("X-Scheme", "HTTP")] }
                `shouldBe` Just HttpProxyTarget
                    { httpProxyTargetHost = "example.com"
                    , httpProxyTargetPort = 80
                    , httpProxyTargetRawPath = "/resource"
                    }

        it "rejects malformed explicit ports in raw absolute-URI requests" $
            selectHttpProxyTarget rawProxyRequest { rawPathInfo = "http://example.com:not-a-port/resource" }
                `shouldBe` Nothing

        it "rejects malformed explicit ports in Host-header proxy requests" $
            selectHttpProxyTarget hostHeaderProxyRequest { requestHeaderHost = Just "example.com:not-a-port" }
                `shouldBe` Nothing

        it "rejects HTTP/2 proxy requests without a Host header" $
            selectHttpProxyTarget http2ProxyRequest { requestHeaderHost = Nothing }
                `shouldBe` Nothing

        it "renders HTTP proxy authorities without the default port" $ do
            renderHttpProxyAuthority (HttpProxyTarget "example.com" 80 "/")
                `shouldBe` "example.com"
            renderHttpProxyAuthority (HttpProxyTarget "example.com" 443 "/")
                `shouldBe` "example.com:443"
            renderHttpProxyAuthority (HttpProxyTarget "example.com" 8080 "/")
                `shouldBe` "example.com:8080"

    describe "HTTP proxy upstream requests" $
        it "uses absolute URI authority for Host and strips proxy boundary headers" $
            withRawHttpServer observeRawProxyRequestBody $ \port -> do
                let authority = "127.0.0.1:" <> BS8.pack (show port)
                    responseBody = LBS8.lines . simpleBody
                    proxiedReq =
                        defaultRequest
                            { requestMethod     = "GET"
                            , rawPathInfo       = "http://" <> authority <> "/resource"
                            , requestHeaderHost = Just "evil.example"
                            , requestHeaders    =
                                [ (HT.hHost, "evil.example")
                                , (HT.hProxyAuthorization, "Basic " <> B64.encode "alice:secret")
                                , ("Forwarded", "for=203.0.113.7")
                                , ("X-Forwarded-For", "203.0.113.7")
                                , ("Proxy-Connection", "keep-alive")
                                , ("cf-ray", "abc123")
                                , ("X-Real-IP", "203.0.113.7")
                                , ("User-Agent", "hprox-test")
                                ]
                            }
                response <- runProxyApp (httpProxy authorizedSettings) proxiedReq
                simpleStatus response `shouldBe` HT.status200
                responseBody response `shouldBe`
                    [ LBS.fromStrict authority
                    , LBS.fromStrict authority
                    , "False"
                    , "False"
                    , "False"
                    , "False"
                    , "False"
                    , "True"
                    ]

-- Keep proxy upstream tests independent of Warp's Windows accept loop.
withRawHttpServer :: (BS.ByteString -> LBS.ByteString) -> (Int -> IO a) -> IO a
withRawHttpServer responseBody action =
    bracket open cleanup run
  where
    open = do
        accepted <- newIORef Nothing
        sock <- openTcpServerSocket
        return (sock, accepted)

    cleanup (sock, accepted) = do
        closeIgnoringErrors sock
        readIORef accepted >>= mapM_ closeIgnoringErrors

    run server@(sock, _) = do
        done <- newEmptyMVar
        _ <- forkIO $ try (acceptAndRespond server) >>= putMVar done
        result <- action . fromIntegral =<< socketPort sock
        serverResult <- timeout 2000000 (takeMVar done)
        case serverResult of
            Just (Left ex) -> throwIO (ex :: SomeException)
            _              -> return result

    acceptAndRespond (sock, accepted) =
        bracket (acceptClient sock accepted) closeIgnoringErrors $ \conn -> do
            requestBytes <- recvHttpRequest conn
            sendHttpResponse conn $ responseBody requestBytes

    acceptClient sock accepted = do
        (conn, _) <- accept sock
        writeIORef accepted (Just conn)
        return conn

openTcpServerSocket :: IO Socket
openTcpServerSocket = bracketOnError open closeIgnoringErrors $ \sock -> do
    setSocketOption sock ReuseAddr 1
    bind sock (SockAddrInet 0 loopbackAddress)
    listen sock (max 2048 maxListenQueue)
    return sock
  where
    open = socket AF_INET Stream defaultProtocol

recvHttpRequest :: Socket -> IO BS.ByteString
recvHttpRequest conn = go BS.empty
  where
    go received
        | "\r\n\r\n" `BS.isInfixOf` received = return received
        | BS.length received > 65536         = ioError $ userError "upstream HTTP request headers exceeded test limit"
        | otherwise                          = do
            chunk <- recvWithTimeout conn
            if BS.null chunk
            then ioError $ userError "upstream HTTP request ended before headers completed"
            else go (received <> chunk)

sendHttpResponse :: Socket -> LBS.ByteString -> IO ()
sendHttpResponse conn body =
    let strictBody = LBS.toStrict body
    in SocketBS.sendAll conn $ BS8.concat
        [ "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: "
        , BS8.pack (show $ BS.length strictBody)
        , "\r\nConnection: close\r\n\r\n"
        , strictBody
        ]

closeIgnoringErrors :: Socket -> IO ()
closeIgnoringErrors sock = do
    closed <- try (close sock) :: IO (Either IOException ())
    case closed of
        Left _   -> return ()
        Right () -> return ()

observeRawProxyRequestBody :: BS.ByteString -> LBS.ByteString
observeRawProxyRequestBody requestBytes =
    LBS8.unlines
        [ LBS.fromStrict $ headerValue HT.hHost
        , LBS.fromStrict $ headerValue HT.hHost
        , LBS8.pack $ show $ hasHeader HT.hProxyAuthorization
        , LBS8.pack $ show $ hasHeader "Forwarded"
        , LBS8.pack $ show $ hasHeader "X-Forwarded-For"
        , LBS8.pack $ show $ hasHeader "Proxy-Connection"
        , LBS8.pack $ show $ hasHeader "cf-ray"
        , LBS8.pack $ show $ hasHeader "User-Agent"
        ]
  where
    headers = mapMaybe parseHeader $ takeWhile (not . BS.null) $ drop 1 $ map trimCR $ BS8.lines requestBytes
    normalizedHeaders = [(CI.mk name, value) | (name, value) <- headers]

    headerValue name = fromMaybe "" $ lookup name normalizedHeaders
    hasHeader name = lookup name normalizedHeaders /= Nothing

    parseHeader line =
        let (name, rest) = BS8.break (== ':') line
        in if BS.null rest
           then Nothing
           else Just (name, BS8.dropWhile (== ' ') $ BS.drop 1 rest)

    trimCR line
        | "\r" `BS.isSuffixOf` line = BS.init line
        | otherwise                 = line

runProxyApp :: (HC.Manager -> Middleware) -> Request -> IO SResponse
runProxyApp middleware req = do
    manager <- HC.newManager $
        HC.managerSetProxy HC.noProxy HC.defaultManagerSettings
            { HC.managerIdleConnectionCount = 0
            , HC.managerResponseTimeout     = HC.responseTimeoutMicro 2000000
            }
    runSession (srequest $ SRequest req "") (middleware manager fallback)

runConnectApp :: ProxySettings -> Request -> IO SResponse
runConnectApp pset req = runSession (srequest $ SRequest req "") (httpConnectProxy pset fallback)

rawProxyRequest :: Request
rawProxyRequest = defaultRequest
    { requestMethod     = "GET"
    , rawPathInfo       = "http://example.com/resource"
    , requestHeaderHost = Just "example.com"
    }

rawProxyRequestWithAuth :: BS8.ByteString -> Request
rawProxyRequestWithAuth auth = rawProxyRequest
    { requestHeaders = [(HT.hProxyAuthorization, auth)]
    }

proxyRequestForPort :: Int -> BS8.ByteString -> Request
proxyRequestForPort port auth =
    let authority = "127.0.0.1:" <> BS8.pack (show port)
    in (rawProxyRequestWithAuth auth)
        { rawPathInfo       = "http://" <> authority <> "/resource"
        , requestHeaderHost = Just authority
        }

basicAliceSecret :: BS8.ByteString
basicAliceSecret = basicCredential "alice:secret"

basicCredential :: BS8.ByteString -> BS8.ByteString
basicCredential credential = "Basic " <> B64.encode credential

hostHeaderProxyRequest :: Request
hostHeaderProxyRequest = defaultRequest
    { requestMethod     = "GET"
    , rawPathInfo       = "/resource"
    , requestHeaderHost = Just "example.com"
    , requestHeaders    = [("Proxy-Connection", "keep-alive")]
    }

http2ProxyRequest :: Request
http2ProxyRequest = defaultRequest
    { requestMethod     = "GET"
    , rawPathInfo       = "/resource"
    , requestHeaderHost = Just "example.com"
    , requestHeaders    = [("X-Scheme", "http")]
    , httpVersion       = HT.http20
    , isSecure          = True
    }

connectRequest :: Request
connectRequest = defaultRequest
    { requestMethod     = "CONNECT"
    , rawPathInfo       = "example.com:443"
    , requestHeaderHost = Just "example.com:443"
    }

connectRequestFor :: BS8.ByteString -> Int -> HT.HttpVersion -> Request
connectRequestFor host port version =
    let authority = host <> ":" <> BS8.pack (show port)
    in connectRequest
        { rawPathInfo       = authority
        , requestHeaderHost = Just authority
        , requestHeaders    = [(HT.hProxyAuthorization, basicAliceSecret)]
        , httpVersion       = version
        }

getFreePort :: IO Int
getFreePort = bracket open close socketPortInt
  where
    open = do
        sock <- socket AF_INET Stream defaultProtocol
        setSocketOption sock ReuseAddr 1
        bind sock (SockAddrInet 0 loopbackAddress)
        return sock

    socketPortInt sock = fromIntegral <$> socketPort sock

withTcpEchoServer :: (Int -> IO a) -> IO a
withTcpEchoServer action = do
    accepted <- newIORef Nothing
    bracket (open accepted) cleanup run
  where
    open accepted = do
        sock <- openTcpServerSocket
        return (sock, accepted)

    cleanup (sock, accepted) = do
        closeIgnoringErrors sock
        readIORef accepted >>= mapM_ closeIgnoringErrors

    run server@(sock, _) =
        withAsync (acceptAndEcho server) $ \_ ->
            (action . fromIntegral =<< socketPort sock) `finally` cleanup server

    acceptAndEcho (sock, accepted) =
        bracket (acceptClient sock accepted) closeIgnoringErrors echoLoop

    acceptClient sock accepted = do
        (conn, _) <- accept sock
        writeIORef accepted (Just conn)
        return conn

    echoLoop conn = do
        bs <- SocketBS.recv conn 4096
        unless (BS8.null bs) $ do
            SocketBS.sendAll conn bs
            echoLoop conn

-- Exercise the HTTP/1 raw response path without running a nested Warp server.
rawConnectRoundTrip :: IO BS8.ByteString
rawConnectRoundTrip = do
    outbound <- newIORef []
    clientChunks <- newIORef ["ping", ""]
    let targetHost = "127.0.0.1"
        targetPort = 18443
        runTCPClient settings app = do
            CN.getHost settings `shouldBe` targetHost
            CN.getPort settings `shouldBe` targetPort
            inbound <- newEmptyMVar
            app AppData
                { appRead'            = takeMVar inbound
                , appWrite'           = \bytes -> do
                    putMVar inbound bytes
                    putMVar inbound BS.empty
                , appSockAddr'        = SockAddrInet 0 loopbackAddress
                , appLocalAddr'       = Nothing
                , appCloseConnection' = return ()
                , appRawSocket'       = Nothing
                }
        nextClientChunk =
            atomicModifyIORef' clientChunks $ \chunks ->
                case chunks of
                    []     -> ([], BS.empty)
                    x : xs -> (xs, x)
        respondRaw (ResponseRaw raw _) = do
            raw nextClientChunk (\bytes -> modifyIORef' outbound (bytes :))
            return ResponseReceived
        respondRaw _ = expectationFailure "expected HTTP/1 raw CONNECT response" >> return ResponseReceived
    _ <- httpConnectProxyWith runTCPClient authorizedSettings fallback (connectRequestFor targetHost targetPort HT.http11) respondRaw
    sent <- BS.concat . reverse <$> readIORef outbound
    sent `shouldSatisfy` BS8.isPrefixOf "HTTP/1.1 200"
    case BS.breakSubstring "\r\n\r\n" sent of
        (_, body) | not (BS.null body) -> return $ BS.drop 4 body
        _                             -> expectationFailure "raw CONNECT response did not contain an HTTP header terminator" >> return BS.empty

recvWithTimeout :: Socket -> IO BS8.ByteString
recvWithTimeout sock = do
    received <- timeout 2000000 (SocketBS.recv sock 4096)
    case received of
        Just bytes -> return bytes
        Nothing    -> ioError $ userError "timed out waiting for proxy socket response"

loopbackAddress :: HostAddress
loopbackAddress = tupleToHostAddress (127, 0, 0, 1)

fallback :: Application
fallback _req respond = respond $ responseLBS fallbackStatus [("Content-Type", "text/plain")] fallbackBody

fallbackStatus :: HT.Status
fallbackStatus = HT.mkStatus 418 "fallback"

fallbackBody :: LBS.ByteString
fallbackBody = "fallback"

authRequiredSettings :: ProxySettings
authRequiredSettings = baseSettings
    { proxyAuth = Just (const False)
    }

hiddenAuthSettings :: ProxySettings
hiddenAuthSettings = authRequiredSettings
    { hideProxyAuth = True
    }

authorizedSettings :: ProxySettings
authorizedSettings = baseSettings
    { proxyAuth = Just (== "alice:secret")
    }

baseSettings :: ProxySettings
baseSettings = ProxySettings
    { proxyAuth      = Nothing
    , passPrompt     = Just "hprox"
    , wsRemote       = Nothing
    , revRemoteMap   = []
    , hideProxyAuth  = False
    , naivePadding   = False
    , acmeThumbprint = Nothing
    , logger         = \_ _ -> return ()
    }