Pipes.hs

module Network.Mail.Postie.Pipes
  ( dataChunks,
    attoParser,
    UnexpectedEndOfInputException,
    TooMuchDataException,
  )
where

import Control.Applicative
import Control.Exception (Exception, throw)
import Control.Monad (unless)
import qualified Data.Attoparsec.ByteString as AT
import qualified Data.ByteString.Char8 as BS
import Data.Maybe (fromMaybe)
import Data.Typeable (Typeable)
import Pipes
import Pipes.Parse
import Prelude hiding (lines)

data UnexpectedEndOfInputException = UnexpectedEndOfInputException
  deriving (Show, Typeable)

data TooMuchDataException = TooMuchDataException
  deriving (Show, Typeable)

instance Exception UnexpectedEndOfInputException

instance Exception TooMuchDataException

attoParser :: AT.Parser r -> Parser BS.ByteString IO (Maybe r)
attoParser p = do
  result <- AT.parseWith draw' p ""
  case result of
    AT.Done t r -> do
      unless (BS.null t) (unDraw t)
      return (Just r)
    _ -> return Nothing
  where
    draw' = fromMaybe "" <$> draw

dataChunks :: Int -> Producer BS.ByteString IO () -> Producer BS.ByteString IO ()
dataChunks n p = lines p >-> go n
  where
    go remaining | remaining <= 0 = throw TooMuchDataException
    go remaining = do
      bs <- await
      unless (bs == ".") $ do
        yield (unescape bs)
        yield "\r\n"
        go (remaining - BS.length bs - 2)
    unescape bs
      | BS.null bs = bs
      | BS.head bs == '.' && BS.length bs > 1 = BS.tail bs
      | otherwise = bs

lines :: Producer BS.ByteString IO () -> Producer BS.ByteString IO ()
lines = go
  where
    go p = do
      (line, leftover) <- lift $ runStateT lineParser p
      yield line
      go leftover

lineParser :: Parser BS.ByteString IO BS.ByteString
lineParser = go id
  where
    go f = do
      bs <- maybe (throw UnexpectedEndOfInputException) (return . f) =<< draw
      case BS.elemIndex '\r' bs of
        Nothing -> go (BS.append bs)
        Just n -> do
          let here = killCR $ BS.take n bs
              rest = BS.drop (n + 1) bs
          unDraw rest
          return here
    killCR bs
      | BS.null bs = bs
      | BS.head bs == '\n' || BS.head bs == '\r' = killCR $ BS.tail bs
      | BS.last bs == '\n' || BS.last bs == '\r' = killCR $ BS.init bs
      | otherwise = bs