{-----------------------------------------------------------------------------------------
Module name: BufferedSocket
Made by: Tomas Möre 2015
Usage:
This module is mean to be imported qualified
ex: import Headers qualified BS
Notes:
Buffered sockets are a data type that is a kind of overlay on normal sockets.
All the exported read / write operations are build such that they ALLWAYS read / write the ammount of bytes requested
This package allows some cases of lazy IO. Some people see Lazy io as the devil incarné. However it is excpected that anyone using this module
is cabale of understanding any possible side effects.
WARNINGS:
This module uses a ton of IO and non functional ways of solving problems.
This is because we want to be as spacetime efficient as possible.
This module does NOT contain "beatiful" haskell code
------------------------------------------------------------------------------------------}
{-# LANGUAGE OverloadedStrings #-}
module Core
(
readRaw
, readLazy
, readByte
, readToByte
, sendByteString
, flush
, BufferedSocket
, makeBufferedSocket
, MaxLength
, ReadSize
, inBuffer
, waitForRead
, isReadable
, isWriteable
, closeRead
, closeWrite
, nativeSocket
)
where
import Prelude hiding (getLine, read)
import Control.Monad
import Control.Applicative
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Char8 as BC8
import qualified Data.Text as T
import qualified Data.Text.Lazy as TL
import Data.List
import Data.IORef
import Data.Functor
import Data.Maybe (isJust)
import qualified Network.Socket as NS -- ns for native socket
import Foreign.Storable
import Foreign.Ptr
import Foreign.C.Types
import System.IO.Unsafe -- I'm so sorry for this
import System.Posix.Types
import System.Timeout
import Foreign.ForeignPtr
import Data.Monoid
import Data.Word
import Data.String
import qualified Data.Text.Encoding as ENCSTRICT
import qualified Data.Text.Lazy.Encoding as ENCLAZY
import GHC.Conc.IO
type MaxLength = Int
type BufferSize = Int
type InputBufferSize = BufferSize
type OutputBufferSize = BufferSize
type BytesInBuffer = IORef Int
type ByteOffset = IORef Int
type RemainingBytes = Int
type ReadSize = Int
type Read = Int
type Timeout = Int
type SocketData = (NS.Socket,NS.SockAddr,Fd)
type ByteString = B.ByteString
-- Since newtype is disregarded at compile time we make these newtypes to make sure no errors occur
-- The output buffer is simply a buffer that attemps to store a certain ammount of data before we send the package.
-- This is to make minimalize the TCP overhead. Some operating systems might attemp to do the same.
type OutputBuffer = (ForeignPtr Word8, BytesInBuffer,BufferSize)
type InputBuffer = (ForeignPtr Word8, ByteOffset,BytesInBuffer,BufferSize)
newtype BufferedSocket = BufferedSocket (SocketData, InputBuffer, OutputBuffer)
{-
UTIL
-}
nativeSocket:: BufferedSocket -> NS.Socket
nativeSocket (BufferedSocket ((socket,_,_),_,_)) = socket
isReadable:: BufferedSocket -> IO Bool
isReadable (BufferedSocket ((socket,_,_),_,_)) = NS.sIsReadable socket
isWriteable:: BufferedSocket -> IO Bool
isWriteable (BufferedSocket ((socket,_,_),_,_)) = NS.isWritable socket
closeRead :: BufferedSocket -> IO ()
closeRead (BufferedSocket ((socket,_,_),_,_)) = NS.shutdown socket NS.ShutdownReceive
closeWrite :: BufferedSocket -> IO ()
closeWrite (BufferedSocket ((socket,_,_),_,_)) = NS.shutdown socket NS.ShutdownSend
-- INPUT BUFFER
inBuffer :: BufferedSocket -> InputBuffer
inBuffer (BufferedSocket (_,inBuf,_)) = inBuf
inBufferClear :: InputBuffer -> IO ()
inBufferClear (_, offsetRef, bytesRef, _) = writeIORef offsetRef 0 >> writeIORef bytesRef 0
inBufferReadAll :: InputBuffer -> IO ByteString
inBufferReadAll inBuf@(buf, offsetRef, bytesRef, _) =
do
offset <- readIORef offsetRef
bytesN <- readIORef bytesRef
let unreadBytes = bytesN - offset
if bytesN == 0
then return ""
else (inBufferClear inBuf) >> (return $ BI.fromForeignPtr buf offset unreadBytes)
-- moves the bytes in the buffer to the start
inBufferRealign :: InputBuffer -> IO ()
inBufferRealign (buf, offsetRef, bytesRef, _) =
withForeignPtr buf $ \ptr ->
do putStrLn "realigned"
offset <- readIORef offsetRef
when (offset > 0) $
do bytesN <- readIORef bytesRef
let unreadBytes = bytesN - offset
offsetPtr = plusPtr ptr offset
BI.memcpy ptr offsetPtr unreadBytes
writeIORef offsetRef 0
inBufferFindByteReal :: Ptr Word8 -> Int -> Word8 -> IO (Maybe Int)
inBufferFindByteReal _ 0 _ = return Nothing
inBufferFindByteReal ptr bytesLeft matchByte =
do currentChar <- peek ptr
if isMatch currentChar
then return $ Just bytesLeft
else inBufferFindByteReal (plusPtr ptr 1) (bytesLeft - 1) matchByte
where
isMatch = (==matchByte)
inBufferFindByte :: InputBuffer -> Word8 -> IO (Maybe Int)
inBufferFindByte (buf, offsetRef, bytesRef, _) byte =
do offset <- readIORef offsetRef
bytesN <- readIORef bytesRef
let unreadBytes = bytesN - offset
withForeignPtr buf $
\ ptr -> do let startPtr = plusPtr ptr offset
maybeBytesLeft <- inBufferFindByteReal startPtr unreadBytes byte
case maybeBytesLeft of
Just n -> return $ Just (unreadBytes - n)
Nothing -> return Nothing
-- OUTPUT BUFFER
outBuffer :: BufferedSocket -> OutputBuffer
outBuffer (BufferedSocket (_,_,outBuf)) = outBuf
-- SOCKET DATA
socketData :: BufferedSocket -> SocketData
socketData (BufferedSocket (sockData,_,_)) = sockData
-- Reads up to max the available bytes of data
bSocketRecv :: BufferedSocket -> IO Int
bSocketRecv (BufferedSocket ((sock,_,_),(buf,offsetRef,bytesNRef,bufSize),_)) =
withForeignPtr buf $ \ptr ->
do bytesN <- readIORef bytesNRef
let offsetBuf = plusPtr ptr bytesN
maxRead = bufSize - bytesN
if maxRead > 0
then
do
(recievedBytes, _) <- NS.recvBufFrom sock offsetBuf maxRead
writeIORef bytesNRef (bytesN + recievedBytes)
return recievedBytes
else
return 0
-- Unsafe. Do not use if you don't know what you're doing
bSocketRecvMin :: BufferedSocket -> Int -> IO ()
bSocketRecvMin bSocket n =
do bytesRead <- bSocketRecv bSocket
when (bytesRead < n) $ bSocketRecvMin bSocket (n - bytesRead)
waitForRead :: BufferedSocket -> Timeout -> IO Bool
waitForRead (BufferedSocket ((sock,_,fideDesc),_,_)) timeoutDuration =
do maybeSucess <- timeout timeoutDuration (threadWaitRead fideDesc)
if isJust maybeSucess
then return True
else return False
{-
CREATING A SOCKET
This puts together a new "BufferedSocket". or commonly refered as "bSock"
A buffered socket is a socket which is acommodated by a Ptr. Because of the need to split headers etc up at certain points but still
safet he rest of the data to be whatever the server requires it to be.
A buffered socket should be called with freeBSocket before disregarded.
The server thunk function does so automaticly.
-}
--BufferedSocket ((sock, sockAddr, socketFileDesc), inputBuffer, outPutBuffer)
makeBufferedSocket :: (NS.Socket, NS.SockAddr) -> InputBufferSize -> OutputBufferSize -> IO BufferedSocket
makeBufferedSocket (sock, sockAddr) inBufferSize outBufferSize =
do
inputBuffer <- makeInputBuffer inBufferSize
outPutBuffer <- makeOutputBuffer outBufferSize
let socketFileDesc = NS.fdSocket sock
return $ BufferedSocket ((sock, sockAddr, Fd socketFileDesc), inputBuffer, outPutBuffer)
makeInputBuffer :: InputBufferSize -> IO InputBuffer
makeInputBuffer bufferSize =
do offset <- newIORef 0
bytesCount <- newIORef 0
buffer <- BI.mallocByteString bufferSize
return (buffer, offset, bytesCount, bufferSize)
makeOutputBuffer :: OutputBufferSize -> IO OutputBuffer
makeOutputBuffer bufferSize =
do bytesCount <- newIORef 0
buffer <- BI.mallocByteString bufferSize
return (buffer, bytesCount, bufferSize)
{-
USING THE BUFFERED SOCKETS:
-}
{- Strict reading
If the buffer contains the required ammount of bytes it will simply read it.
If not it will read all the bytes. Clear the buffer and attempt to read more
-}
readRaw :: BufferedSocket -> Int -> IO ByteString
readRaw _ 0 = return ""
readRaw bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) readSize =
do bytesN <- readIORef bytesNRef
offset <- readIORef offsetRef
let unreadBytes = bytesN - offset
availableBytesAfterBytesN = bufSize - bytesN
availableBytesTotal = availableBytesAfterBytesN + offset
missingBytes = readSize - unreadBytes
if missingBytes <= 0
then writeIORef offsetRef (offset + readSize) >> (return $ BI.fromForeignPtr buf offset readSize)
else if missingBytes <= availableBytesAfterBytesN
then bSocketRecvMin bSock missingBytes >>
loop
else if missingBytes <= availableBytesTotal
then inBufferRealign inBuf >>
fillBuffer availableBytesTotal >>
loop
else do
let strFragment = if unreadBytes > 0
then BI.fromForeignPtr buf offset unreadBytes
else ""
inBufferClear inBuf
(strFragment<>) <$> readRaw bSock (readSize - B.length strFragment)
where
fillBuffer = bSocketRecvMin bSock
loop = readRaw bSock readSize
--
lazyReader :: BufferedSocket -> [Int] -> IO [ByteString]
lazyReader _ [] = return []
lazyReader bSock ~(chunkSize:rest) =
do chunk <- unsafeInterleaveIO $ readRaw bSock chunkSize
next <- unsafeInterleaveIO $ lazyReader bSock rest
return (chunk:next)
-- readLazy will read will read in chunks of the same size of the buffered socket
readLazy :: BufferedSocket -> Int -> IO BL.ByteString
readLazy _ 0 = return ""
readLazy bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) readSize =
let chunkSizes = unfoldr (\b -> if b == 0
then Nothing
else if b > bufSize
then Just (bufSize, b - bufSize)
else Just (b,0)) readSize
in BL.fromChunks <$> lazyReader bSock chunkSizes
readByte :: BufferedSocket -> IO Word8
readByte bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) =
do bytesN <- readIORef bytesNRef
offset <- readIORef offsetRef
let unreadBytes = bytesN - offset
if unreadBytes > 0
then do
writeIORef offsetRef (offset + 1)
withForeignPtr buf (`peekByteOff` offset)
else do
inBufferClear inBuf
bSocketRecvMin bSock 1
readByte bSock
readToByte :: BufferedSocket -> Word8 -> IO ByteString
readToByte bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) byte =
do maybeByteIndex <- inBufferFindByte inBuf byte
case maybeByteIndex of
Nothing -> (<>) <$> inBufferReadAll inBuf <*> readToByte bSock byte -- >>= (<>) fmap (<>(readToByte bSock byte)) $ inBufferReadAll inBuf
Just n -> do
offset <- readIORef offsetRef
writeIORef offsetRef (offset + n + 1)
return $ BI.fromForeignPtr buf offset n
readToByteMax :: BufferedSocket -> Word8 -> MaxLength -> IO (Maybe ByteString)
readToByteMax bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) byte maxLen
| hasNoMaxLen = return Nothing
| otherwise =
do maybeByteIndex <- inBufferFindByte inBuf byte
bytesN <- readIORef bytesNRef
offset <- readIORef offsetRef
let unreadBytes = bytesN - offset
case maybeByteIndex of
Nothing -> if unreadBytes >= maxLen
then return Nothing
else do let maxBytesLeft = maxLen - unreadBytes
thisBufData <- inBufferReadAll inBuf -- empties the current buffer data
bSocketRecv bSock -- attempts to rebuffer the socket
maybeData <- readToByteMax bSock byte maxBytesLeft -- repeat
case maybeData of
Nothing -> return Nothing
Just a -> return $ Just (thisBufData <> a)
Just n -> do writeIORef offsetRef (offset + n + 1)
return $ Just $ BI.fromForeignPtr buf offset n
where
hasNoMaxLen = maxLen <= 0
-- Dangerous will read forever if nothing is found
-- You will most probably want to use the readToByteStringMax
readToByteString :: BufferedSocket -> ByteString -> IO ByteString
readToByteString _ "" = error "readToByteString can not take an empty bytestring"
readToByteString bSock searchString =
do dataString <- readToByte bSock firstByte
trail <- readRaw bSock $ B.length restSearchString
if trail == restSearchString
then return dataString
else ((dataString<>trail)<>) <$> readToByteString bSock searchString
where
firstByte = B.head searchString
restSearchString = B.tail searchString
-- Same as above but Will read to a max length. If this is reached It returns "Nothing"
readToByteStringMax :: BufferedSocket -> ByteString -> MaxLength -> IO (Maybe ByteString)
readToByteStringMax _ _ 0 = return Nothing
readToByteStringMax _ "" _ = error "readToByteString max got an empty string"
readToByteStringMax bSock searchString maxLength =
do maybeDataString <- readToByteMax bSock firstByte maxLength
case maybeDataString of
Nothing -> return Nothing
Just dataString -> do trail <- readRaw bSock $ B.length restSearchString
if trail == restSearchString
then return $ Just dataString
else do maybeMoreString <- readToByteStringMax bSock searchString (maxLength - B.length dataString - B.length trail)
case maybeMoreString of
Nothing -> return Nothing
Just a -> return $ Just $ dataString <> trail <> a
where
firstByte = B.head searchString
restSearchString = B.tail searchString
-- Simple and clean. force pushed the data to the actual socket. Afther this
flush :: BufferedSocket -> IO()
flush (BufferedSocket ((socket,sockAddr,_),_,(outBuf, bytesNRef, bufferSize))) =
do bytesInBuffer <- readIORef bytesNRef
when (bytesInBuffer > 0) $
do let sender ptr bytes =
do sent <- NS.sendBufTo socket ptr bytes sockAddr
when (sent < bytes) $ sender (plusPtr ptr sent) (bytes - sent)
str = BI.fromForeignPtr outBuf 0 bytesInBuffer
withForeignPtr outBuf (`sender` bytesInBuffer)
writeIORef bytesNRef 0
-- if the string sent to this function is less or equal to the ammount of bytes left in the buffer
-- we will simply fill the buffer
-- (if equal) we will also flush this buffer.
-- If the string is greater then the ammount of bytes left
-- we will attempt to fill the last bytes in the buffer then
sendByteString:: BufferedSocket -> ByteString -> IO ()
sendByteString _ "" = return ()
sendByteString bSock@(BufferedSocket ((sock,sockAddr,_),_,(outBuf, bytesNRef, bufferSize))) outputStr =
do bytesInBuffer <- readIORef bytesNRef
let bytesLeftInBuffer = bufferSize - bytesInBuffer
if srcLength <= bytesLeftInBuffer
then do
withForeignPtr outBuf $ \bufPtr ->
withForeignPtr sourceFrgnPtr $ \srcPtr ->
BI.memcpy (plusPtr bufPtr bytesInBuffer) (plusPtr srcPtr srcOffet) srcLength
writeIORef bytesNRef (bytesInBuffer + srcLength)
when (srcLength + bytesInBuffer == bufferSize) $ flush bSock
else do
let (current,rest) = B.splitAt bytesLeftInBuffer outputStr
overflowRest = div (B.length rest) bufferSize
directSendChunk = B.take overflowRest rest
directSendLen = B.length directSendChunk
(directSend,_,_)= BI.toForeignPtr directSendChunk
toBufferPart = B.drop overflowRest rest
sendByteString bSock current
when (directSendChunk /= "") $ withForeignPtr sourceFrgnPtr $ \ptr -> void $ NS.sendBufTo sock (plusPtr ptr srcOffet) directSendLen sockAddr
sendByteString bSock toBufferPart
where
(sourceFrgnPtr, srcOffet, srcLength) = BI.toForeignPtr outputStr