redirect-server.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Prelude hiding (log)

import Control.Concurrent
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.Reader
import Debug.Trace (traceIO)
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8
import Network
import System.IO

class Monad m => Serve m where
    recv :: Int -> m ByteString
    send :: ByteString -> m ()
    getClientNumber :: m Int
    log :: String -> m ()
    quit :: m a

sl :: Serve m => String -> m ()
sl s = send $ B8.pack $ s ++ "\r\n"

serve :: Serve m => m ()
serve = do
    qs <- recvQueryString
    case qs of
        "redirect" -> do
            sl "HTTP/1.1 301 Moved Permanently"
            sl "Location: 200"
            sl ""

        "200" -> do
            n <- getClientNumber
            sl "HTTP/1.1 200 OK"
            sl "Content-Type: text/html"
            sl ""
            sl $ "Client " ++ show n

        _ -> do
            log $ "No handler for /" ++ qs
            sl "HTTP/1.1 404 Not Found"
            sl "Content-Type: text/html"
            sl ""
            sl $ "<h3>No handler for /" ++ qs ++ "</h3>"

recvQueryString :: Serve m => m String
recvQueryString = do
    s <- recv 4096
    case words (B8.unpack s) of
        ("GET" : ('/' : qs) : "HTTP/1.1" : _) -> do
            log $ "GET /" ++ qs
            return qs
        _ -> do
            log "Invalid request; closing connection"
            sl "HTTP/1.1 400 Bad Request"
            sl "Content-Type: text/html"
            sl ""
            sl "<h3>Invalid request</h3>"
            quit

------------------------------------------------------------------------

data Env = Env
    { envHandle       :: Handle
    , envClientNumber :: Int
    }

instance Serve (ReaderT Env IO) where
    recv n = do
        Env{..} <- ask
        liftIO $ B.hGetSome envHandle n
    send s = do
        Env{..} <- ask
        liftIO $ B.hPutStr envHandle s
    getClientNumber = asks envClientNumber
    log s = do
        n <- getClientNumber
        liftIO $ traceIO $ "Client " ++ show n ++ ": " ++ s
    quit = liftIO $ throwIO ThreadKilled

serverLoop :: Socket -> Int -> IO loop
serverLoop sock n = do
    (h, host, port) <- accept sock
    traceIO $ "server: received connection from " ++ host ++ ":" ++ show port
    _ <- forkIO $ do
        hSetBuffering h NoBuffering
        res <- try $ runReaderT serve Env{envHandle = h, envClientNumber = n}
        case res of
            Left (ex :: SomeException) ->
                traceIO $ "Client " ++ show n ++ ": " ++ show ex
            Right () -> return ()
        hClose h
    serverLoop sock $! n + 1

main :: IO ()
main = do
    let portno = 1234
    sock <- listenOn $ PortNumber portno
    traceIO $ "Listening on port " ++ show portno
    serverLoop sock 1