fcmtp-client/src/Network/FCMTP/SendingState.hs

{-# LANGUAGE OverloadedStrings, TypeFamilies, MultiParamTypeClasses, FlexibleInstances, UndecidableInstances, RankNTypes #-}

module Network.FCMTP.SendingState (
  Sending,
  inSendingConnection,
  inSendingSession,
  inSendingCtx,
  peelSendingCtx,
  startSending,
  startSession,
  failures,
  extensions,
  quitCmd,
  rsetCmd,
  mailCmd,
  dataCmd,
  startData,
  dataChunk,
  rcptCmd
  ) where

-- That's another StateT that encapsulates a Stream and sending data
-- Only a bit more complicated, because every command leads to a response
-- that may be treated in batch or individually according to the PIPELINING
-- or XTUNNEL extension.

import qualified Data.FCMTP.Account as Ac
import qualified Data.FCMTP.Extension as Ex
import qualified Data.FCMTP.Mime as Mime
import System.IO (Handle)
import qualified System.IO.Uniform as U
import qualified System.IO.Uniform.Streamline as S
import qualified Network.FCMTP.ClientError as CE
import qualified Data.FCMTP.Response as Resp
import Data.FCMTP.Response (Response, ResponseStatus(..))
import Data.FCMTP.ResponseCode (toResponse, ResponseCode(..))
import Control.Conditional

import Control.Monad (ap, liftM)
import Control.Monad.IO.Class
import Control.Monad.Trans
import Control.Monad.Trans.Interruptible
import Control.Monad.Trans.Control
import Control.Monad.Base

import Control.Exception

import Data.Maybe (isJust)
import Data.Default.Class
import qualified Data.Attoparsec.ByteString as A
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy.Search as Search

import Text.StringConvert
import Debug.Trace

chunkSize :: Int
chunkSize = 100000 -- bytes - about 10 normal ethernet packages

{- |
How to chunk data for sending.
filter text isLast = [(generatesAReply, sentText)]
-}
type SendingDataFilter = ByteString -> Bool -> [(Bool, ByteString)]
data Data = Data {
  hostName :: ByteString, tlsSettings :: U.TlsSettings,
  getExtensions :: [Ex.Extension], activeRcpt :: [Ac.Account],
  failure :: [(Ac.Account, Response)],
  pendingCommands :: [PendingCommand], isPipelining :: Bool,
  dataFilter :: SendingDataFilter}

newtype Sending m a = Sending {withSending :: Data -> S.Streamline m (a, Data)}

-- Commands with pending replies, segregated by what failure means.
data PendingCommand = DummyCommand | TotalyImportantCommand | CommandForRcpt Ac.Account

dummyFilter :: SendingDataFilter
dummyFilter dt _ = [(False, dt)]

startSending :: MonadIO m => Maybe Handle -> Sending m ()
startSending ec = do
  lifts $ S.echoTo ec
  handshake

startSession :: MonadIO m => Sending m ()
startSession = startSendingSession

inSendingConnection :: MonadIO m => ToString host => host -> U.TlsSettings -> Sending m a -> S.Streamline m a
inSendingConnection h set f = do
  (ret, _) <- withSending f' $ Data (s . toString $ h) set [] [] [] [] False dummyFilter
  return ret
  where
    f' = do
      handshake
      ret <- f
      quitCmd
      return ret

startSendingSession :: MonadIO m => Sending m ()
startSendingSession = Sending $ \st -> do
  let s' = st{activeRcpt=[]}{failure=[]}
  return ((), s')

clearSendingSession :: MonadIO m => Sending m ()
clearSendingSession = Sending $ \st -> do
  let s' = st{activeRcpt=[]}{failure=[]}
  return ((), s')

inSendingSession :: MonadIO m => Sending m a -> Sending m (a, [(Ac.Account, Response)])
inSendingSession f = do
  sync
  startSendingSession
  ret <- f
  rsetCmd
  sync
  ff <- failures
  clearSendingSession
  return (ret, ff)

-- | Returns (success address, temporary failures, permanent failures)
failures :: Monad m => Sending m [(Ac.Account, Response)]
failures = Sending $ \st -> return (failure st, st)

extensions :: Monad m => Sending m [Ex.Extension]
extensions = Sending $ \st -> return (getExtensions st, st)

quitCmd :: MonadIO m => Sending m ()
quitCmd = do
  sync
  sendDummyCommand "QUIT"

rsetCmd :: MonadIO m => Sending m ()
rsetCmd = do
  sendDummyCommand "RSET"
  sync

mailCmd :: MonadIO m => Ac.Account -> Sending m ()
mailCmd from = sendTotalyCommand . BS.concat $ ["MAIL FROM:<", Ac.normalize from, ">"]

rcptCmd :: MonadIO m => Ac.Account -> Sending m ()
rcptCmd to = let
  cmd = ["RCPT TO:", s . Ac.fullAccount $ to]
  in sendRcptCommand (BS.concat cmd) to

setDataFilter :: Monad m => SendingDataFilter -> Sending m ()
setDataFilter f = Sending $ \st -> return ((), st{dataFilter=f})

getDataFilter :: Monad m => Sending m SendingDataFilter
getDataFilter = Sending $ \st -> return (dataFilter st, st)

{- |
Starts a data sending command (DATA or BDAT).
-}
startData :: MonadIO m => Mime.BodyEncoding -> Sending m ()
startData enc = do
  sync
  ifM (hasEx Ex.CHUNKING) (
    setDataFilter sendBdatBody
    )(
    do
      ch <- canHandle enc
      traceShow enc $ return ()
      traceShow ch $ return ()
      ifM (canHandle enc) (
        do
          totalyDataIfActive $ case enc of
            Mime.B7BitEncoding -> "DATA\r\n"
            Mime.B8BitEncoding -> "DATA BODY=8BIT\r\n"
            Mime.BBinaryEncoding -> "DATA BODY=BINARY\r\n"
          sync
          setDataFilter $ \m lst ->
            (False, LBS.toStrict . dotStuff . LBS.fromStrict $ m) :
            if lst then [(True, "\r\n.\r\n")] else []
        ) (failAll' . toResponse $ NoConversion)
    )
  where
    canHandle :: Monad m => Mime.BodyEncoding -> Sending m Bool
    canHandle Mime.B7BitEncoding = return True
    canHandle Mime.B8BitEncoding = hasEx Ex.E8BITMIME
    canHandle Mime.BBinaryEncoding = hasEx Ex.BINARYMIME
    -- desiredEncoding :: Mime.BodyEncoding -> Sending m m Mime.BodyEncoding
    -- desiredEncoding o = case o of
    --   Mime.B7BitEncoding -> return Mime.B7BitEncoding
    --   Mime.B8BitEncoding -> ifM (hasEx Ex.E8BITMIME) (return Mime.B8BitEncoding) (
    --     ifM (hasEx Ex.BINARYMIME) (return Mime.BBinaryEncoding) (return Mime.B7BitEncoding)
    --     )
    --   Mime.BBinaryEncoding -> ifM (hasEx Ex.BINARYMIME) (
    --     return Mime.BBinaryEncoding) (return Mime.B7BitEncoding)
    sendBdatBody :: SendingDataFilter
    sendBdatBody dt lst = if lst
                          then [(True, BS.concat ["BDAT ", s . show . BS.length $ dt, " LAST\r\n", dt])]
                          else [(True, BS.concat ["BDAT ", s . show . BS.length $ dt, "\r\n", dt])]

dataChunk :: MonadIO m => ByteString -> Bool -> Sending m ()
dataChunk dt lst = do
  f <- getDataFilter
  let dt' = f dt lst
  mapM_ sendChunk dt'
  where
    --sendChunk :: (Bool, ByteString) -> Sending m ()
    sendChunk (repl, ck) = if repl
                           then totalyCommandNoLn ck
                           else dataIfActive ck

dataCmd :: MonadIO m => Mime.BodyEncoding -> LBS.ByteString -> Sending m ()
dataCmd enc m = ifM (hasEx Ex.CHUNKING) (
  sendBdatBody m -- No reencode is needed.
  ) $ do
    sendTotalyCommand $ case enc of
      Mime.B7BitEncoding -> "DATA"
      Mime.B8BitEncoding -> "DATA BODY=8BIT"
      Mime.BBinaryEncoding -> "DATA BODY=BINARY"
    sync
    totalyDataIfActive' $ LBS.append (dotStuff m) "\r\n.\r\n"
  where
    -- desiredEncoding :: Mime.BodyEncoding -> Sending m m Mime.BodyEncoding
    -- desiredEncoding o = case o of
    --   Mime.B7BitEncoding -> return Mime.B7BitEncoding
    --   Mime.B8BitEncoding -> ifM (hasEx Ex.E8BITMIME) (return Mime.B8BitEncoding) (
    --     ifM (hasEx Ex.BINARYMIME) (return Mime.BBinaryEncoding) (return Mime.B7BitEncoding)
    --     )
    --   Mime.BBinaryEncoding -> ifM (hasEx Ex.BINARYMIME) (
    --     return Mime.BBinaryEncoding) (return Mime.B7BitEncoding)
    sendBdatBody :: MonadIO m => LBS.ByteString -> Sending m ()
    sendBdatBody dt = do
      let (d, r) = LBS.splitAt (fromIntegral chunkSize) dt
      if LBS.null r
        then
        totalyDataIfActive $ BS.concat [
          "BDAT ", s . show . LBS.length $ d, " LAST\r\n", LBS.toStrict d]
        else do
        totalyDataIfActive $ BS.concat [
          "BDAT ", s . show . LBS.length $ d, "\r\n", LBS.toStrict d]
        sendBdatBody r

sendDummyCommand :: MonadIO m => ByteString -> Sending m () 
sendDummyCommand cmd = do
  lifts . S.send . BS.concat $ [cmd, "\r\n"]
  pipeline DummyCommand

sendTotalyCommand :: MonadIO m => ByteString -> Sending m () 
sendTotalyCommand cmd = totalyCommandNoLn $ BS.concat [cmd, "\r\n"]

totalyCommandNoLn :: MonadIO m => ByteString -> Sending m ()
totalyCommandNoLn cmd = do
  lifts . S.send $ cmd
  pipeline TotalyImportantCommand

dataIfActive :: MonadIO m => ByteString -> Sending m ()
dataIfActive dt = whenM hasActiveRcpt (lifts $ S.send dt)

totalyDataIfActive :: MonadIO m => ByteString -> Sending m ()
totalyDataIfActive dt = whenM hasActiveRcpt $ do
  lifts $ S.send dt
  pipeline TotalyImportantCommand
  
totalyDataIfActive' :: MonadIO m => LBS.ByteString -> Sending m ()
totalyDataIfActive' dt = whenM hasActiveRcpt $ do
  lifts $ S.send' dt
  pipeline TotalyImportantCommand

sendRcptCommand :: MonadIO m => ByteString -> Ac.Account -> Sending m () 
sendRcptCommand cmd e = do
  lifts . S.send . BS.concat $ [cmd, "\r\n"]
  Sending (\st -> let
                   a = e : activeRcpt st
                   st' = st{activeRcpt=a}
                   in if  e `elem` activeRcpt st then return ((), st) else return ((), st')
               )
  pipeline . CommandForRcpt $ e

pipeline :: MonadIO m => PendingCommand -> Sending m ()
pipeline cmd = Sending $ \st ->
  if isPipelining st
  then 
    let p = cmd : pendingCommands st
    in return ((), st{pendingCommands=p}) 
  else do
    repl <- receiveReply
    let st' = resolveCommand repl cmd st
    return ((), st')

resolveCommand :: Response -> PendingCommand -> Data -> Data
resolveCommand repl cmd st = case Resp.status repl of
  Preliminary -> st
  Completion -> st
  Intermediate -> st
  TransientError -> case cmd of 
    DummyCommand -> st
    TotalyImportantCommand -> failAll repl st
    CommandForRcpt e -> failAdd repl st e
  PermanentError -> case cmd of
    DummyCommand -> st
    TotalyImportantCommand -> failAll repl st
    CommandForRcpt e -> failAdd repl st e
  DataFollows -> failAll repl st -- Shouldn't appear when sending
  Asynchronous -> st

failAdd :: Response -> Data -> Ac.Account -> Data
failAdd r st e = st{activeRcpt = a}{failure = (e, r): failure st}
  where
    a = removeElem (activeRcpt st) e

removeElem :: Eq a => [a] -> a -> [a]
removeElem l e = filter (/= e) l

failAll :: Response -> Data -> Data
failAll r st = st{activeRcpt=[]}{failure = a  ++ failure st}
  where
    a = map (\x -> (x, r)) $ activeRcpt st

failAll' :: Monad m => Response -> Sending m ()
failAll' r = Sending $ \st -> return ((), failAll r st)

sync :: MonadIO m => Sending m ()
sync = do
  cmds <- queryState pendingCommands
  mapM_ getReply $ reverse cmds
  zeroPendingCommands
  where
    getReply :: MonadIO m => PendingCommand -> Sending m ()
    getReply cmd = Sending $ \st -> do
      repl <- receiveReply
      let st' = resolveCommand repl cmd st
      return ((), st')
    zeroPendingCommands :: MonadIO m => Sending m ()
    zeroPendingCommands = Sending $ \st -> do
      return ((), st{pendingCommands=[]})

queryState :: Monad m => (Data -> a) -> Sending m a
queryState f = Sending $ \st ->
  return (f st, st)

handshake :: MonadIO m => Sending m ()
handshake = Sending (
  \st -> do
    S.recieveLine
    let host = hostName st
    exts <- getEhlo host
    secureLine <- S.isSecure
    let tls = Ex.hasExtension exts Ex.STARTTLS && not secureLine
    if tls
      then do
      S.send "STARTTLS\r\n"
      S.recieveLine
      S.startTls $ tlsSettings st
      exts' <- getEhlo host
      let pipe = Ex.hasExtension exts' Ex.PIPELINING
      let st' = st{getExtensions=exts'}{isPipelining=pipe}
      return ((), st')
      else do
      let pipe = Ex.hasExtension exts Ex.PIPELINING
      let st' = st{getExtensions=exts}{isPipelining=pipe}
      return ((), st')
  )
  where
    getEhlo host = do
      S.send . BS.concat $ ["EHLO ", host, "\r\n"]
      repl <- receiveReply
      if isSuccess . Resp.status $ repl
        then do
        let exts = ehloExts repl
        return exts
        else do
        S.send . BS.concat $ ["HELO ", host, "\r\n"]
        repl' <- receiveReply
        if isSuccess . Resp.status $ repl'
          then return []
          else liftIO . throwIO $ CE.ProtocolError

--Client code needs:
--A parser that reads the server response, and breaks it in:
--Error with:
--  Type and message
--Success with:
--  Message, and possibly contents, where contents are
--  EHLO extension support, MNTR changes, CHSM values, or RTRV data
--
--Because of MNTR, the parser must be lazy, and return
--before command termination.

ehloExts :: Response -> [Ex.Extension]
ehloExts r = map (\msg -> case A.parseOnly Ex.parseExtension msg of
                             Right e -> e         
                             Left e -> Ex.StringExt Ex.UNRECOGNIZED $ BS.concat [s e, ": ", msg]
                         ) $ Resp.respLines r

isSuccess :: ResponseStatus -> Bool
isSuccess TransientError = False
isSuccess PermanentError = False
isSuccess _ = True


receiveReply :: MonadIO m => S.Streamline m Resp.Response
receiveReply = do
  r <- S.runAttoparsec Resp.parseResponse
  case r of
    Left e -> return def{Resp.message=s e}
    Right v -> return v

-- Here Sending becomes a usefull Monad

instance Monad m => Monad (Sending m) where
  --return :: (Monad m) => a -> Sending m a
  return x = Sending  $ \cl -> return (x, cl)
  --(>>=) :: Sending m a -> (a -> Sending m b) -> Sending m b
  a >>= b = Sending $ \cl -> do
    (x, cl') <- withSending a cl
    withSending (b x) cl'

instance Monad m => Functor (Sending m) where
  --fmap :: (a -> b) -> Sending m a -> Sending m b
  fmap f m = Sending $ \cl -> do
    (x, cl') <- withSending m cl
    return (f x, cl')

instance Monad m => Applicative (Sending m) where
    pure = return
    (<*>) = ap


instance MonadTrans Sending where
  lift = lifts . lift

lifts :: Monad m => S.Streamline m a -> Sending m a
lifts x = Sending $ \st -> do
  r <- x
  return (r, st)

instance MonadIO m => MonadIO (Sending m) where
  liftIO = lifts . liftIO

dotStuff :: LBS.ByteString -> LBS.ByteString
dotStuff m = LBS.append
             (LBS.fromStrict leading)
             (Search.replace ("\r\n."::ByteString) ("\r\n.."::ByteString) m)
  where
    leading = if LBS.take 1 m == "." then "a" else ""

hasEx :: Monad m => Ex.ExtName -> Sending m Bool
hasEx x = Sending $ \st -> return (isJust $ Ex.getExtension (getExtensions st) x, st)

hasActiveRcpt :: Monad m => Sending m Bool
hasActiveRcpt = Sending $ \st -> return (not . null . activeRcpt $ st, st)


--------------------------------------------
-- Interruptible and MonadControl support --
--------------------------------------------

instance Interruptible Sending where
  type RSt Sending a = RSt S.Streamline (a, Data)
  resume f st = resume (\(x, dt) -> withSending (f x) dt) st

peelSendingCtx :: RSt Sending a -> (a, U.SomeIO, [(Ac.Account, Response)])
peelSendingCtx st = let
  ((a, dt), io) = S.peelStreamlineCtx st
  in (a, io, failure dt)

inSendingCtx :: (ToString host, U.UniformIO io) => host -> U.TlsSettings -> io -> a -> RSt Sending a
inSendingCtx h set io a = let
  dt = Data (s . toString $ h) set [] [] [] [] False dummyFilter
  in S.inStreamlineCtx io (a, dt)

instance MonadTransControl Sending where
  type StT Sending a = StT S.Streamline (a, Data)
  -- liftWith :: (Run Sending -> m a) -> Sending m a
  liftWith f = Sending $ \st ->
    liftM (\x -> (x, st)) $ liftWith $ \f' -> do
      f $ \(Sending t) -> do
        f' $ t st
  --restoreT :: m (St Sending a) -> Sending m a
  restoreT r = Sending . const . restoreT $ r
    --Sending $ \s -> do
    --(st, dt) <- r
    --restoreT . return $ (\(v, x) -> ((v, dt), x)) st

instance MonadBase b m => MonadBase b (Sending m) where
  liftBase = liftBaseDefault

instance MonadBaseControl b m => MonadBaseControl b (Sending m) where
  type StM (Sending m) a = ComposeSt Sending m a
  liftBaseWith     = defaultLiftBaseWith
  restoreM         = defaultRestoreM