{-# LANGUAGE TypeSynonymInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Network.TCP
-- Copyright   :  See LICENSE file
-- License     :  BSD
--
-- Maintainer  :  Ganesh Sittampalam <ganesh@earth.li>
-- Stability   :  experimental
-- Portability :  non-portable (not tested)
--
-- Some utility functions for working with the Haskell @network@ package. Mostly
-- for internal use by the @Network.HTTP@ code.
--
-----------------------------------------------------------------------------
module Network.TCP
   ( Connection
   , EndPoint(..)
   , openTCPPort
   , isConnectedTo

   , openTCPConnection
   , socketConnection
   , isTCPConnectedTo
   
   , HandleStream
   , HStream(..)
   
   , StreamHooks(..)
   , nullHooks
   , setStreamHooks
   , getStreamHooks
   , hstreamToConnection

   ) where

import Network.Socket
   ( Socket, SocketOption(KeepAlive)
   , SocketType(Stream), connect
   , shutdown, ShutdownCmd(..)
   , sClose, setSocketOption, getPeerName
   , socket, Family(AF_UNSPEC), defaultProtocol, getAddrInfo
   , defaultHints, addrFamily, withSocketsDo
   , addrSocketType, addrAddress
   )
import qualified Network.Stream as Stream
   ( Stream(readBlock, readLine, writeBlock, close, closeOnEnd) )
import Network.Stream
   ( ConnError(..)
   , Result
   , failWith
   , failMisc
   )
import Network.BufferType

import Network.HTTP.Base ( catchIO )
import Network.Socket ( socketToHandle )

import Data.Char  ( toLower )
import Data.Word  ( Word8 )
import Control.Concurrent
import Control.Exception ( onException )
import Control.Monad ( liftM, when )
import System.IO ( Handle, hFlush, IOMode(..), hClose )
import System.IO.Error ( isEOFError )

import qualified Data.ByteString      as Strict
import qualified Data.ByteString.Lazy as Lazy

-----------------------------------------------------------------
------------------ TCP Connections ------------------------------
-----------------------------------------------------------------

-- | The 'Connection' newtype is a wrapper that allows us to make
-- connections an instance of the Stream class, without GHC extensions.
-- While this looks sort of like a generic reference to the transport
-- layer it is actually TCP specific, which can be seen in the
-- implementation of the 'Stream Connection' instance.
newtype Connection = Connection (HandleStream String)

newtype HandleStream a = HandleStream {getRef :: MVar (Conn a)}

data EndPoint = EndPoint { epHost :: String, epPort :: Int }

instance Eq EndPoint where
   EndPoint host1 port1 == EndPoint host2 port2 =
     map toLower host1 == map toLower host2 && port1 == port2

data Conn a 
 = MkConn { connSock      :: ! Socket
          , connHandle    :: Handle
          , connBuffer    :: BufferOp a
          , connInput     :: Maybe a
          , connEndPoint  :: EndPoint
          , connHooks     :: Maybe (StreamHooks a)
          , connCloseEOF  :: Bool -- True => close socket upon reaching end-of-stream.
          }
 | ConnClosed
   deriving(Eq)

hstreamToConnection :: HandleStream String -> Connection
hstreamToConnection h = Connection h

connHooks' :: Conn a -> Maybe (StreamHooks a)
connHooks' ConnClosed{} = Nothing
connHooks' x = connHooks x

-- all of these are post-op hooks
data StreamHooks ty
 = StreamHooks
     { hook_readLine   :: (ty -> String) -> Result ty -> IO ()
     , hook_readBlock  :: (ty -> String) -> Int -> Result ty -> IO ()
     , hook_writeBlock :: (ty -> String) -> ty  -> Result () -> IO ()
     , hook_close      :: IO ()
     , hook_name       :: String -- hack alert: name of the hook itself.
     }

instance Eq ty => Eq (StreamHooks ty) where
  (==) _ _ = True

nullHooks :: StreamHooks ty
nullHooks = StreamHooks 
     { hook_readLine   = \ _ _   -> return ()
     , hook_readBlock  = \ _ _ _ -> return ()
     , hook_writeBlock = \ _ _ _ -> return ()
     , hook_close      = return ()
     , hook_name       = ""
     }

setStreamHooks :: HandleStream ty -> StreamHooks ty -> IO ()
setStreamHooks h sh = modifyMVar_ (getRef h) (\ c -> return c{connHooks=Just sh})

getStreamHooks :: HandleStream ty -> IO (Maybe (StreamHooks ty))
getStreamHooks h = readMVar (getRef h) >>= return.connHooks

-- | @HStream@ overloads the use of 'HandleStream's, letting you
-- overload the handle operations over the type that is communicated
-- across the handle. It comes in handy for @Network.HTTP@ 'Request'
-- and 'Response's as the payload representation isn't fixed, but overloaded.
--
-- The library comes with instances for @ByteString@s and @String@, but
-- should you want to plug in your own payload representation, defining
-- your own @HStream@ instance _should_ be all that it takes.
-- 
class BufferType bufType => HStream bufType where
  openStream       :: String -> Int -> IO (HandleStream bufType)
  openSocketStream :: String -> Int -> Socket -> IO (HandleStream bufType)
  readLine         :: HandleStream bufType -> IO (Result bufType)
  readBlock        :: HandleStream bufType -> Int -> IO (Result bufType)
  writeBlock       :: HandleStream bufType -> bufType -> IO (Result ())
  close            :: HandleStream bufType -> IO ()
  closeQuick       :: HandleStream bufType -> IO ()
  closeOnEnd       :: HandleStream bufType -> Bool -> IO ()
  
instance HStream Strict.ByteString where
  openStream       = openTCPConnection
  openSocketStream = socketConnection
  readBlock c n    = readBlockBS c n
  readLine c       = readLineBS c
  writeBlock c str = writeBlockBS c str
  close c          = closeIt c Strict.null True
  closeQuick c     = closeIt c Strict.null False
  closeOnEnd c f   = closeEOF c f

instance HStream Lazy.ByteString where
    openStream       = \ a b -> openTCPConnection_ a b True
    openSocketStream = \ a b c -> socketConnection_ a b c True
    readBlock c n    = readBlockBS c n
    readLine c       = readLineBS c
    writeBlock c str = writeBlockBS c str
    close c          = closeIt c Lazy.null True
    closeQuick c     = closeIt c Lazy.null False
    closeOnEnd c f   = closeEOF c f

instance Stream.Stream Connection where
  readBlock (Connection c)     = Network.TCP.readBlock c
  readLine (Connection c)      = Network.TCP.readLine c
  writeBlock (Connection c)    = Network.TCP.writeBlock c 
  close (Connection c)         = Network.TCP.close c
  closeOnEnd (Connection c) f  = Network.TCP.closeEOF c f
  
instance HStream String where
    openStream      = openTCPConnection
    openSocketStream = socketConnection
    readBlock ref n = readBlockBS ref n

    -- This function uses a buffer, at this time the buffer is just 1000 characters.
    -- (however many bytes this is is left to the user to decypher)
    readLine ref = readLineBS ref
    -- The 'Connection' object allows no outward buffering, 
    -- since in general messages are serialised in their entirety.
    writeBlock ref str = writeBlockBS ref str -- (stringToBuf str)

    -- Closes a Connection.  Connection will no longer
    -- allow any of the other Stream functions.  Notice that a Connection may close
    -- at any time before a call to this function.  This function is idempotent.
    -- (I think the behaviour here is TCP specific)
    close c = closeIt c null True
    
    -- Closes a Connection without munching the rest of the stream.
    closeQuick c = closeIt c null False

    closeOnEnd c f = closeEOF c f
    
-- | @openTCPPort uri port@  establishes a connection to a remote
-- host, using 'getHostByName' which possibly queries the DNS system, hence 
-- may trigger a network connection.
openTCPPort :: String -> Int -> IO Connection
openTCPPort uri port = openTCPConnection uri port >>= return.Connection

-- Add a "persistent" option?  Current persistent is default.
-- Use "Result" type for synchronous exception reporting?
openTCPConnection :: BufferType ty => String -> Int -> IO (HandleStream ty)
openTCPConnection uri port = openTCPConnection_ uri port False

openTCPConnection_ :: BufferType ty => String -> Int -> Bool -> IO (HandleStream ty)
openTCPConnection_ uri port stashInput = do
    -- HACK: uri is sometimes obtained by calling Network.URI.uriRegName, and this includes
    -- the surrounding square brackets for an RFC 2732 host like [::1]. It's not clear whether
    -- it should, or whether all call sites should be using something different instead, but
    -- the simplest short-term fix is to strip any surrounding square brackets here.
    -- It shouldn't affect any as this is the only situation they can occur - see RFC 3986.
    let fixedUri =
         case uri of
            '[':(rest@(c:_)) | last rest == ']'
              -> if c == 'v' || c == 'V'
                     then error $ "Unsupported post-IPv6 address " ++ uri
                     else init rest
            _ -> uri


    -- use withSocketsDo here in case the caller hasn't used it, which would make getAddrInfo fail on Windows
    -- although withSocketsDo is supposed to wrap the entire program, in practice it is safe to use it locally
    -- like this as it just does a once-only installation of a shutdown handler to run at program exit,
    -- rather than actually shutting down after the action
    addrinfos <- withSocketsDo $ getAddrInfo (Just $ defaultHints { addrFamily = AF_UNSPEC, addrSocketType = Stream }) (Just fixedUri) (Just . show $ port)
    case addrinfos of
        [] -> fail "openTCPConnection: getAddrInfo returned no address information"
        (a:_) -> do
                s <- socket (addrFamily a) Stream defaultProtocol
                onException (do
                            setSocketOption s KeepAlive 1
                            connect s (addrAddress a)
                            socketConnection_ fixedUri port s stashInput
                            ) (sClose s)

-- | @socketConnection@, like @openConnection@ but using a pre-existing 'Socket'.
socketConnection :: BufferType ty
                 => String
                 -> Int
                 -> Socket
                 -> IO (HandleStream ty)
socketConnection hst port sock = socketConnection_ hst port sock False

-- Internal function used to control the on-demand streaming of input
-- for /lazy/ streams.
socketConnection_ :: BufferType ty
                  => String
                  -> Int
                  -> Socket
                  -> Bool
                  -> IO (HandleStream ty)
socketConnection_ hst port sock stashInput = do
    h <- socketToHandle sock ReadWriteMode
    mb <- case stashInput of { True -> liftM Just $ buf_hGetContents bufferOps h; _ -> return Nothing }
    let conn = MkConn 
         { connSock     = sock
         , connHandle   = h
         , connBuffer   = bufferOps
         , connInput    = mb
         , connEndPoint = EndPoint hst port
         , connHooks    = Nothing
         , connCloseEOF = False
         }
    v <- newMVar conn
    return (HandleStream v)

closeConnection :: HStream a => HandleStream a -> IO Bool -> IO ()
closeConnection ref readL = do
    -- won't hold onto the lock for the duration
    -- we are draining it...ToDo: have Connection
    -- into a shutting-down state so that other
    -- threads will simply back off if/when attempting
    -- to also close it.
  c <- readMVar (getRef ref)
  closeConn c `catchIO` (\_ -> return ())
  modifyMVar_ (getRef ref) (\ _ -> return ConnClosed)
 where
   -- Be kind to peer & close gracefully.
  closeConn ConnClosed = return ()
  closeConn conn = do
    let sk = connSock conn
    hFlush (connHandle conn)
    shutdown sk ShutdownSend
    suck readL
    hClose (connHandle conn)
    shutdown sk ShutdownReceive
    sClose sk

  suck :: IO Bool -> IO ()
  suck rd = do
    f <- rd
    if f then return () else suck rd

-- | Checks both that the underlying Socket is connected
-- and that the connection peer matches the given
-- host name (which is recorded locally).
isConnectedTo :: Connection -> EndPoint -> IO Bool
isConnectedTo (Connection conn) endPoint = isTCPConnectedTo conn endPoint

isTCPConnectedTo :: HandleStream ty -> EndPoint -> IO Bool
isTCPConnectedTo conn endPoint = do
   v <- readMVar (getRef conn)
   case v of
     ConnClosed -> return False
     _ 
      | connEndPoint v == endPoint ->
          catchIO (getPeerName (connSock v) >> return True) (const $ return False)
      | otherwise -> return False

readBlockBS :: HStream a => HandleStream a -> Int -> IO (Result a)
readBlockBS ref n = onNonClosedDo ref $ \ conn -> do
   x <- bufferGetBlock ref n
   maybe (return ())
         (\ h -> hook_readBlock h (buf_toStr $ connBuffer conn) n x)
         (connHooks' conn)
   return x

-- This function uses a buffer, at this time the buffer is just 1000 characters.
-- (however many bytes this is is left for the user to decipher)
readLineBS :: HStream a => HandleStream a -> IO (Result a)
readLineBS ref = onNonClosedDo ref $ \ conn -> do
   x <- bufferReadLine ref
   maybe (return ())
         (\ h -> hook_readLine h (buf_toStr $ connBuffer conn) x)
         (connHooks' conn)
   return x

-- The 'Connection' object allows no outward buffering, 
-- since in general messages are serialised in their entirety.
writeBlockBS :: HandleStream a -> a -> IO (Result ())
writeBlockBS ref b = onNonClosedDo ref $ \ conn -> do
  x    <- bufferPutBlock (connBuffer conn) (connHandle conn) b
  maybe (return ())
        (\ h -> hook_writeBlock h (buf_toStr $ connBuffer conn) b x)
        (connHooks' conn)
  return x

closeIt :: HStream ty => HandleStream ty -> (ty -> Bool) -> Bool -> IO ()
closeIt c p b = do
   closeConnection c (if b
                      then readLineBS c >>= \ x -> case x of { Right xs -> return (p xs); _ -> return True}
                      else return True)
   conn <- readMVar (getRef c)
   maybe (return ())
         (hook_close)
         (connHooks' conn)

closeEOF :: HandleStream ty -> Bool -> IO ()
closeEOF c flg = modifyMVar_ (getRef c) (\ co -> return co{connCloseEOF=flg})

bufferGetBlock :: HStream a => HandleStream a -> Int -> IO (Result a)
bufferGetBlock ref n = onNonClosedDo ref $ \ conn -> do
   case connInput conn of
    Just c -> do
      let (a,b) = buf_splitAt (connBuffer conn) n c
      modifyMVar_ (getRef ref) (\ co -> return co{connInput=Just b})
      return (return a)
    _ -> do
      catchIO (buf_hGet (connBuffer conn) (connHandle conn) n >>= return.return)
              (\ e ->
                       if isEOFError e
                        then do
                          when (connCloseEOF conn) $ catchIO (closeQuick ref) (\ _ -> return ())
                          return (return (buf_empty (connBuffer conn)))
                        else return (failMisc (show e)))

bufferPutBlock :: BufferOp a -> Handle -> a -> IO (Result ())
bufferPutBlock ops h b = 
  catchIO (buf_hPut ops h b >> hFlush h >> return (return ()))
          (\ e -> return (failMisc (show e)))

bufferReadLine :: HStream a => HandleStream a -> IO (Result a)
bufferReadLine ref = onNonClosedDo ref $ \ conn -> do
  case connInput conn of
   Just c -> do
    let (a,b0)  = buf_span (connBuffer conn) (/='\n') c
    let (newl,b1) = buf_splitAt (connBuffer conn) 1 b0
    modifyMVar_ (getRef ref) (\ co -> return co{connInput=Just b1})
    return (return (buf_append (connBuffer conn) a newl))
   _ -> catchIO
              (buf_hGetLine (connBuffer conn) (connHandle conn) >>= 
                    return . return . appendNL (connBuffer conn))
              (\ e ->
                 if isEOFError e
                  then do
                    when (connCloseEOF conn) $ catchIO (closeQuick ref) (\ _ -> return ())
                    return (return   (buf_empty (connBuffer conn)))
                  else return (failMisc (show e)))
 where
   -- yes, this s**ks.. _may_ have to be addressed if perf
   -- suggests worthiness.
  appendNL ops b = buf_snoc ops b nl
  
  nl :: Word8
  nl = fromIntegral (fromEnum '\n')

onNonClosedDo :: HandleStream a -> (Conn a -> IO (Result b)) -> IO (Result b)
onNonClosedDo h act = do
  x <- readMVar (getRef h)
  case x of
    ConnClosed{} -> return (failWith ErrorClosed)
    _ -> act x