walrus-backend/src/Walrus/Backend.hs

{-# LANGUAGE OverloadedStrings #-}

module Walrus.Backend (
  Backend(..),
  LazyHandler,
  UIOHandler,
  ConduitHandler,
  module Walrus.Backend.Metadata,
  parseBackend,
  callBackend,
  runLazyOnce,
  runUIOOnce,
  runConduitOnce
  ) where

import qualified Data.Attoparsec.ByteString.Char8 as A
import qualified Data.Attoparsec.ByteString.Lazy as LA
import Data.Attoparsec.ByteString.Char8.Extras
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Char as C
import Control.Monad.IO.Class
import qualified Text.StringConvert as SC
import Control.Lens

import Walrus.Backend.Metadata

import Network
import System.IO
import System.IO.Uniform
import qualified System.IO.Uniform.Streamline as S
--import System.IO.Uniform.HandlePair as HP
--import System.IO.Uniform.Network as Net
import qualified System.Process as P
import qualified System.IO.Uniform.Conduit as CIO
import Data.Conduit
import qualified Data.Conduit.Attoparsec as CA
import qualified Data.Conduit.List as CL

data Backend = TCPBackend String Int |
               ExecBackend String [String] deriving (Show, Read, Ord, Eq)

type LazyHandler = Metadata -> LBS.ByteString -> IO (Either String (Metadata, LBS.ByteString))
type UIOHandler = Metadata -> S.Streamline IO (Either String (Metadata, LBS.ByteString))
type ConduitHandler = Metadata -> ConduitM BS.ByteString BS.ByteString IO (Either String (Metadata, LBS.ByteString))

parseBackend :: A.Parser Backend
parseBackend = A.choice [
  do
    tp "tcp"
    h <- tillSpace
    p <- A.decimal
    return $ TCPBackend h p,
  do
    tp "exec"
    f <- qStr
    skipHorizontalSpace
    pp <- parseParameters
    return $ ExecBackend f pp
  ]
  where
    tp t = do
      skipHorizontalSpace
      A.stringCI t
      skipHorizontalSpace
    tillSpace :: A.Parser String
    tillSpace = SC.s <$> A.takeTill C.isSpace
    parseParameters :: A.Parser [String]
    parseParameters = do
      p <- qStr
      skipHorizontalSpace
      if null p then return [] else do
        pp <- parseParameters
        return $ p : pp
    qStr :: A.Parser String
    qStr = SC.s <$> quotedString '\\' " " ""

callBackend :: Backend -> (Metadata, LBS.ByteString) -> IO (Either String (Metadata, LBS.ByteString))
callBackend b (m, qdt) = do
  let rm = renderMetadata m
  edt' <- intBk b $ LBS.concat [
    SC.s rm,
    qdt]
  case LA.parse repParse edt' of
    LA.Fail _ _ e -> return $ Left e
    LA.Done edt m' -> return $ Right (m', edt)
  where
    intBk :: Backend -> LBS.ByteString -> IO LBS.ByteString
    intBk (TCPBackend h p) = runTcp h p
    intBk (ExecBackend f aa) = runExec f aa
    repParse = parseMetadata

runLazyOnce :: UniformIO u => u -> LazyHandler -> IO (Either String ())
runLazyOnce u f =  S.withTarget u $ do
  m' <- S.runAttoparsec parseMetadata
  case m' of
    Left e -> return . Left . show $ e
    Right m -> do
      dt <- S.recieveN' $ m^.dataSize
      r <- liftIO $ f m dt
      case r of
        Left e -> return . Left $ e
        Right (mo, dto) -> do
          S.send $ renderMetadata mo
          S.send' dto
          return . Right $ ()

runUIOOnce :: UniformIO u => u -> UIOHandler -> IO (Either String ())
runUIOOnce u f =  S.withTarget u $ do
  m' <- S.runAttoparsec parseMetadata
  case m' of
    Left e -> return . Left . show $ e
    Right m -> do
      r <- f m
      case r of
        Left e -> return . Left $ e
        Right (mo, dto) -> do
          S.send $ renderMetadata mo
          S.send' dto
          return . Right $ ()

runConduitOnce :: UniformIO u => u -> ConduitHandler -> IO (Either String ())
runConduitOnce u f = CIO.runConduit u $ do
  m' <- CA.sinkParserEither parseMetadata
  case m' of
    Left e -> return . Left . show $ e
    Right m -> do
      r <- f m
      case r of
        Left e -> return . Left $ e
        Right (mo, dto) -> do
          yield . renderMetadata $ mo
          CL.sourceList . LBS.toChunks $ dto
          return . Right $ ()

runTcp :: String -> Int -> LBS.ByteString -> IO LBS.ByteString
runTcp host port dt = do
  h <- connectTo host (PortNumber . fromIntegral $ port)
  bkgTrans h dt

bkgTrans :: Handle -> LBS.ByteString -> IO LBS.ByteString
bkgTrans h dt = do
  hSetNewlineMode h noNewlineTranslation
  LBS.hPut h dt
  hFlush h
  LBS.hGetContents h

runExec :: String -> [String] -> LBS.ByteString -> IO LBS.ByteString
runExec f args dt = do
  (Just i, Just o, _, _) <- P.createProcess (
    (P.proc f args){P.std_in=P.CreatePipe}{
     P.std_out=P.CreatePipe}{P.std_err=P.Inherit}
    )
  mapM_ (`hSetNewlineMode` noNewlineTranslation) [i, o]
  LBS.hPut i dt
  hFlush i
  LBS.hGetContents o