{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveDataTypeable #-}

{-# OPTIONS -Wall #-}

-- | Main entry point.

module Main where

import           Control.Concurrent
import           Control.Concurrent.Logger
import           Control.Exception
import           Control.Monad
import qualified Data.ByteString as S
import           Data.Tuple.Extra
import           Data.Monoid
import           Network hiding (accept)
import           Network.Socket hiding (recv)
import           Network.Socket.ByteString
import           Prelude hiding (catch)
import           System.Console.CmdArgs
import           System.Posix

-- | Main entry point.
main :: IO ()
main = do
  void $ installHandler sigPIPE Ignore Nothing
  cmdArgs options >>= start

-- | Program configuration.
data Throttle = Throttle
  { listenPort     :: Int
  , connectHost    :: String
  , connectPort    :: Int
  , bytesPerSecond :: Int
  , logging        :: Bool
  } deriving (Show,Data,Typeable)

-- | Opttions.
options :: Throttle
options = Throttle
  { bytesPerSecond = def &= opt (1024 :: Int) &= help "Bytes per second e.g. 1024 (=1KB/s)"
  , connectHost    = "127.0.0.1"
  , connectPort    = 80
  , listenPort     = 8000
  , logging        = def &= help "Log about events on the console."
  }
  &= summary "Throttle v1.0, (C) Chris Done 2010"
  &= help "Listens on port <listen> and proxies a throttled \
          \connection to <host> on <port> at speed <speed>KB/s."

-- | Start the proxy server.
start :: Throttle -> IO ()
start config@Throttle{..} = withSocketsDo $ do
  l <- newLogger logging
  listener <- listenOn (PortNumber . fromIntegral $ listenPort)
  forever $ do
    (client,_) <- accept listener
    echo l $ [show client,": New connection on port ",show listenPort]
    void $ forkIO $ do
      server <- connectToServer config
      echo l $ [show client,": ",show server,": Connected to server at " ,connectHost,":",show connectPort]
      let proxyTo = proxyToWithChan bytesPerSecond l
      void $ client `proxyTo` server
      void $ server `proxyTo` client
    return ()

-- | Connect to the remote server.
connectToServer :: Throttle -> IO Socket
connectToServer Throttle{..} = do
  addrinfos <- getAddrInfo Nothing (Just connectHost) (Just $ show connectPort)
  let serveraddr = head addrinfos
  server <- socket (addrFamily serveraddr) Stream defaultProtocol
  connect server (addrAddress serveraddr)
  return server

-- | Proxy from one direction to another, sending only so many bytes per second.
proxyToWithChan :: Int -> Maybe (Chan String) -> Socket -> Socket -> IO ()
proxyToWithChan bytes c from to = do
  buffer <- newMVar mempty
  fork $ forever $ do
    payload <- recv from 4096
    modifyMVar_ buffer (return . (<> payload))
  fork $ forever $ do
    payload <- modifyMVar buffer (return . swap . S.splitAt (bytes `div` 10))
    sendAll to payload
    threadDelay (1000 * 100)

  where fork = void . forkIO . flip catch (shutdownProxy c from to)

-- | Shut down both connections on a proxy.
shutdownProxy :: Maybe (Chan String) -> Socket -> Socket -> IOException -> IO ()
shutdownProxy c a b (_ :: IOException) = do
  echo c $ [show a,"+",show b,": Shutting down proxy connection."]
  sClose a
  sClose b