Conduit proxy server

proxy.hs

{-# LANGUAGE OverloadedStrings #-}

import Data.Conduit
import Data.Conduit.Network
import Data.Conduit.Text (encode, decode, utf8)
import qualified Data.Conduit.List as CL
import qualified Data.Conduit.Binary as CB

import Data.Text (toUpper)
import qualified Data.ByteString.Char8 as S8
import Data.ByteString (ByteString)
import Control.Monad.IO.Class (liftIO)
import Control.Concurrent (forkIO)

main :: IO ()
main = do
    _ <- forkIO $ runTCPServer (ServerSettings 4000 HostAny) server
    _ <- forkIO $ runTCPServer (ServerSettings 5000 HostAny) proxy
    runTCPClient (ClientSettings 5000 "localhost") client

server :: Application IO
server src sink = src $$ decode utf8 =$ CL.map toUpper =$ encode utf8 =$ sink

takeLine :: Sink ByteString IO ByteString
takeLine = do
    let linefeed = 10
    bss <- CB.takeWhile (/= linefeed) =$ CL.consume
    CB.drop 1 -- drop the newline
    return $ S8.takeWhile (/= '\r') $ S8.concat bss

getPortHost :: Sink ByteString IO ClientSettings
getPortHost = do
    portBS <- takeLine
    hostBS <- takeLine
    return $ ClientSettings (read $ S8.unpack portBS) (S8.unpack hostBS)

proxy :: Application IO
proxy fromClient0 toClient = do
    (fromClient, clientSettings) <- fromClient0 $$+ getPortHost
    runTCPClient clientSettings (proxyLoop fromClient toClient)

proxyLoop :: ResumableSource IO ByteString
          -> Sink ByteString IO ()
          -> Source IO ByteString
          -> Sink ByteString IO ()
          -> IO ()
proxyLoop fromClient0 toClient fromServer0 toServer = do
    yield "Connected to server" $$ toClient
    -- convert fromServer0 from a normal Source to a ResumableSource
    (fromServer, ()) <- fromServer0 $$+ return ()
    loop fromClient0 fromServer
  where
    loop fromClient fromServer = do
        (fromClient', mbs1) <- fromClient $$++ await
        case mbs1 of
            Nothing -> close fromClient' fromServer
            Just bs1 -> do
                yield bs1 $$ toServer
                (fromServer', mbs2) <- fromServer $$++ await
                case mbs2 of
                    Nothing -> do
                        yield "Server closed connection" $$ toClient
                        close fromClient' fromServer'
                    Just bs2 -> do
                        yield bs2 $$ toClient
                        loop fromClient' fromServer'
    close x y = do
        x $$+- return ()
        y $$+- return ()

client :: Application IO
client src sink =
    src $$ conduit =$ sink
  where
    conduit = do
        yield "4000\r\n"
        yield "localhost\r\n"
        await >>= liftIO . print

        yield "hello"
        await >>= liftIO . print

        yield "world"
        await >>= liftIO . print

        yield "goodbye"
        await >>= liftIO . print

It would be nice to see a proxyLoop example where client --> server and server --> client were in independent threads, avoiding introducing deadlocks. I wonder how much more complicated it would become...

I think it's just a bit simpler actually:

proxyLoop fromClient0 toClient fromServer0 toServer = do
    yield "Connected to server" $$ toClient
    m <- M.newEmptyMVar
    tid1 <- forkIO $ do
        fromServer0 $$ toClient
        M.putMVar m True
    tid2 <- forkIO $ do
        fromClient0 $$+- toServer
        M.putMVar m False
    x <- M.takeMVar m
    if x
        then killThread tid2
        else killThread tid1

I wrote the code the way I did to demonstrate connect-and-resume better. Also, in an actual proxy, you really would want to force one request to complete before starting the next (at least, I think you'd want that... maybe not).