{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}
{- |
-- borrowed from snap-server. Check there periodically for updates.
-}
module Happstack.Server.Internal.TimeoutSocket where

import           Control.Applicative           (pure)
import           Control.Concurrent            (threadWaitWrite)
import           Control.Exception             as E (catch, throw)
import           Control.Monad                 (liftM, when)
import qualified Data.ByteString.Char8         as B
import qualified Data.ByteString.Lazy.Char8    as L
import qualified Data.ByteString.Lazy.Internal as L
import qualified Data.ByteString               as S
import           Network.Socket                (close)
import qualified Network.Socket.ByteString     as N
import qualified Happstack.Server.Internal.TimeoutManager as TM
import           Happstack.Server.Internal.TimeoutIO (TimeoutIO(..))
import           Network.Socket (Socket, ShutdownCmd(..), shutdown)
import           Network.Socket.SendFile (Iter(..), ByteCount, Offset, sendFileIterWith')
import           Network.Socket.ByteString (sendAll)
import           System.IO.Error (isDoesNotExistError, ioeGetErrorType)
import           System.IO.Unsafe (unsafeInterleaveIO)
import           GHC.IO.Exception (IOErrorType(InvalidArgument))

sPutLazyTickle :: TM.Handle -> Socket -> L.ByteString -> IO ()
sPutLazyTickle :: Handle -> Socket -> ByteString -> IO ()
sPutLazyTickle thandle :: Handle
thandle sock :: Socket
sock cs :: ByteString
cs =
    do (ByteString -> IO () -> IO ()) -> IO () -> ByteString -> IO ()
forall a. (ByteString -> a -> a) -> a -> ByteString -> a
L.foldrChunks (\c :: ByteString
c rest :: IO ()
rest -> Socket -> ByteString -> IO ()
sendAll Socket
sock ByteString
c IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Handle -> IO ()
TM.tickle Handle
thandle IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
rest) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) ByteString
cs
{-# INLINE sPutLazyTickle #-}

sPutTickle :: TM.Handle -> Socket -> B.ByteString -> IO ()
sPutTickle :: Handle -> Socket -> ByteString -> IO ()
sPutTickle thandle :: Handle
thandle sock :: Socket
sock cs :: ByteString
cs =
    do Socket -> ByteString -> IO ()
sendAll Socket
sock ByteString
cs
       Handle -> IO ()
TM.tickle Handle
thandle
       () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE sPutTickle #-}

sGet :: TM.Handle
     -> Socket
     -> IO (Maybe B.ByteString)
sGet :: Handle -> Socket -> IO (Maybe ByteString)
sGet handle :: Handle
handle socket :: Socket
socket =
  do ByteString
s <- Socket -> Int -> IO ByteString
N.recv Socket
socket 65536
     Handle -> IO ()
TM.tickle Handle
handle
     if ByteString -> Bool
S.null ByteString
s
       then Maybe ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing
       else Maybe ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
s)

sGetContents :: TM.Handle
             -> Socket         -- ^ Connected socket
             -> IO L.ByteString  -- ^ Data received
sGetContents :: Handle -> Socket -> IO ByteString
sGetContents handle :: Handle
handle sock :: Socket
sock = IO ByteString
loop where
  loop :: IO ByteString
loop = IO ByteString -> IO ByteString
forall a. IO a -> IO a
unsafeInterleaveIO (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
    ByteString
s <- Socket -> Int -> IO ByteString
N.recv Socket
sock 65536
    Handle -> IO ()
TM.tickle Handle
handle
    if ByteString -> Bool
S.null ByteString
s
      then do
        -- 'InvalidArgument' is GHCs code for eNOTCONN (among other
        -- things). Sometimes the other end of socket is closed first
        -- and this end is already disconnected before we do
        -- 'shutdown'. Ignore this exception.
        Socket -> ShutdownCmd -> IO ()
shutdown Socket
sock ShutdownCmd
ShutdownReceive IO () -> (IOError -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch`
                    (\e :: IOError
e -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (IOError -> Bool
isDoesNotExistError IOError
e Bool -> Bool -> Bool
|| IOError -> IOErrorType
ioeGetErrorType IOError
e IOErrorType -> IOErrorType -> Bool
forall a. Eq a => a -> a -> Bool
== IOErrorType
InvalidArgument)) (IOError -> IO ()
forall a e. Exception e => e -> a
throw IOError
e))
        ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
L.Empty
      else ByteString -> ByteString -> ByteString
L.Chunk ByteString
s (ByteString -> ByteString) -> IO ByteString -> IO ByteString
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` IO ByteString
loop


sendFileTickle :: TM.Handle -> Socket -> FilePath -> Offset -> ByteCount -> IO ()
sendFileTickle :: Handle -> Socket -> FilePath -> Offset -> Offset -> IO ()
sendFileTickle thandle :: Handle
thandle outs :: Socket
outs fp :: FilePath
fp offset :: Offset
offset count :: Offset
count =
    (IO Iter -> IO ())
-> Socket -> FilePath -> Offset -> Offset -> Offset -> IO ()
forall a.
(IO Iter -> IO a)
-> Socket -> FilePath -> Offset -> Offset -> Offset -> IO a
sendFileIterWith' (Handle -> IO Iter -> IO ()
iterTickle Handle
thandle) Socket
outs FilePath
fp 65536 Offset
offset Offset
count

iterTickle :: TM.Handle -> IO Iter -> IO ()
iterTickle :: Handle -> IO Iter -> IO ()
iterTickle thandle :: Handle
thandle =
    IO Iter -> IO ()
iterTickle'
    where
      iterTickle' :: (IO Iter -> IO ())
      iterTickle' :: IO Iter -> IO ()
iterTickle' iter :: IO Iter
iter =
          do Iter
r <- IO Iter
iter
             Handle -> IO ()
TM.tickle Handle
thandle
             case Iter
r of
               (Done _) ->
                      () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
               (WouldBlock _ fd :: Fd
fd cont :: IO Iter
cont) ->
                   do Fd -> IO ()
threadWaitWrite Fd
fd
                      IO Iter -> IO ()
iterTickle' IO Iter
cont
               (Sent _ cont :: IO Iter
cont) ->
                   do IO Iter -> IO ()
iterTickle' IO Iter
cont

timeoutSocketIO :: TM.Handle -> Socket -> TimeoutIO
timeoutSocketIO :: Handle -> Socket -> TimeoutIO
timeoutSocketIO handle :: Handle
handle socket :: Socket
socket =
    TimeoutIO :: Handle
-> (ByteString -> IO ())
-> (ByteString -> IO ())
-> IO (Maybe ByteString)
-> IO ByteString
-> (FilePath -> Offset -> Offset -> IO ())
-> IO ()
-> Bool
-> TimeoutIO
TimeoutIO { toHandle :: Handle
toHandle      = Handle
handle
              , toShutdown :: IO ()
toShutdown    = Socket -> IO ()
close Socket
socket
              , toPutLazy :: ByteString -> IO ()
toPutLazy     = Handle -> Socket -> ByteString -> IO ()
sPutLazyTickle Handle
handle Socket
socket
              , toGet :: IO (Maybe ByteString)
toGet         = Handle -> Socket -> IO (Maybe ByteString)
sGet           Handle
handle Socket
socket
              , toPut :: ByteString -> IO ()
toPut         = Handle -> Socket -> ByteString -> IO ()
sPutTickle     Handle
handle Socket
socket
              , toGetContents :: IO ByteString
toGetContents = Handle -> Socket -> IO ByteString
sGetContents   Handle
handle Socket
socket
              , toSendFile :: FilePath -> Offset -> Offset -> IO ()
toSendFile    = Handle -> Socket -> FilePath -> Offset -> Offset -> IO ()
sendFileTickle Handle
handle Socket
socket
              , toSecure :: Bool
toSecure      = Bool
False
              }