esotericbot-0.0.1: Esotericbot/Safely.hs

{-# OPTIONS -fglasgow-exts #-}
module Esotericbot.Safely where

import Prelude as P

import Data.List.Stream as L

import System.Posix
import System.IO as IO

import Control.Concurrent.STM

import Data.Maybe
import Data.Monoid
import Data.ByteString.Lazy.Char8 as BS
import Data.IORef

import Esotericbot.BSUtils
import Esotericbot.EBTypes
import Esotericbot.IRCCom
import Esotericbot.BSH

data SafeArgs =
   SafeArgs {
      output_limit :: Maybe Integer
      , time_limit :: Maybe Integer
      , mem_limit :: Maybe Integer
      , hard_time_limit :: Maybe Integer
      }

handle_alarm f = 
   installHandler realTimeAlarm
                  ( Catch f )
                  Nothing

term_child dVar term_limit mkill_limit pid = if_not_done dVar $ do
   signalProcess softwareTermination pid
   maybe ( return ( ) )
         ( \ kill_limit ->
            if ( kill_limit ) > term_limit
               then do
                  handle_alarm $ kill_child dVar pid
                  scheduleAlarm_ $ kill_limit - term_limit
               else 
                  return ( )
         )
         mkill_limit
                  

kill_child dVar pid = if_not_done dVar $
   signalProcess killProcess pid

if_not_done d f = do
   done <- atomically $ readTVar d
   if done
      then return ( )
      else f

scheduleAlarm_ l = do
   r <- scheduleAlarm l
   if r < 0
      then error "could not schedule alarm!"
      else return ( )

close_on_error :: forall m . Monoid m => SmallBotM m -> Handle -> SmallBotM m
close_on_error f h = do
   sb <- get
   liftIO $ catch ( evalStateT f sb ) ( const $ do IO.hClose h; return mempty )

replace_newlines :: BS.ByteString -> BS.ByteString
replace_newlines bs = 
   BS.map nlToSp bs
   where
   nlToSp '\n' =
      ' '
   nlToSp '\r' =
      ' '
   nlToSp x =
      x

safely :: IRCCommand cmd => Handle -> cmd -> FilePath -> [ String ] -> ByteString -> Maybe Integer -> SmallBotM ( ) 
safely h irc_command cmd input prog mmem = do
   ( parent_read , child_write ) <- liftIO createPipe
   ( child_read , parent_write ) <- liftIO createPipe
   sb <- get
   io_safe_limits <- safe_limits mmem $ mtime_limit sb
   let run_child = do
          closeFd parent_write
          closeFd parent_read
          io_safe_limits
          dupTo child_read stdInput
          dupTo child_write stdOutput
          closeFd child_read
          closeFd child_write
          executeFile cmd True input Nothing
   pid <- liftIO $ forkProcess run_child
   liftIO $ do
      closeFd child_write
      closeFd child_read
   w <- liftIO $ fdToHandle parent_write
   close_on_error ( liftIO $ do
                     BS.hPut w prog 
                     IO.hClose w
                     ) w
   dVar <- liftIO $ atomically $ newTVar False
   term_ind <- maybe ( return False )
                     ( \ term_limit -> liftIO $ do
                        handle_alarm $ term_child dVar term_limit ( mhard_time_limit sb >>= return . fromEnum ) pid 
                        scheduleAlarm_ term_limit
                        return True
                     )
                     $ mtime_limit sb >>= return . fromEnum
   if term_ind
      then
         return ( ) 
      else
         maybe ( return ( ) )
               ( \ kill_limit -> liftIO $ do
                  handle_alarm $ kill_child dVar pid
                  scheduleAlarm_ kill_limit
                  return ( )
               )
               $ mhard_time_limit sb >>= return . fromEnum
   r <- liftIO $ fdToHandle parent_read
   read <- liftIO $ newIORef False
   close_on_error ( do 
                     maybe ( do c <- liftIO $ BS.hGetContents r
                                res <- null_check $ replace_newlines c
                                priv_msg h irc_command res
                           )
                           ( \ ol -> do
                              bs <- liftIO $ BS.hGet r $ 1 + fromIntegral ol
                              if BS.length bs > ol -- check if we need to add dots ( this looks silly but, the user probably would like to know )
                                 then do 
                                    dots <- liftIO $ ls2bs 3 "..."#
                                    priv_msg h irc_command $ replace_newlines ( BS.take ol bs ) `hAppend` dots
                                  else do
                                     n <- null_check $ replace_newlines bs
                                     priv_msg h irc_command n
                           )
                           $ moutput_limit sb
                     liftIO $ writeIORef read True
                     liftIO $ hClose r
                  ) r
   is_read <- liftIO $ readIORef read
   if is_read
      then return ( )
      else do n <- nullness
              priv_msg h irc_command n
   liftIO $ do
      getProcessStatus True False pid
      atomically $ writeTVar dVar True
   where
   null_check c =
      if BS.null c
         then nullness 
         else return c
   nullness =
      liftIO $ ls2bs 10 "No output."#

-- this is disgusting! XD
safe_limits :: Maybe Integer -> Maybe Integer -> SmallBotM ( IO ( ) )
safe_limits cmmem mtime_limit = do
   mmem <- child_mem_limit cmmem
   return $ do
      maybe ( return ( ) )
            ( \ l ->
               setResourceLimit ResourceCPUTime $ both_limits l 
            )
            mtime_limit
      maybe ( return ( ) )
            ( \ l ->
               setResourceLimit ResourceTotalMemory $ both_limits l
            )
            mmem

child_mem_limit :: Maybe Integer -> SmallBotM ( Maybe Integer )
child_mem_limit mmem = do
   dmmem <- default_mem
   return $ maybe dmmem
                  Just
                  mmem
   where
   default_mem = do
      sb <- get
      return $ default_child_mem_limit sb

both_limits l =
   ResourceLimits {
      softLimit = ResourceLimit l
      , hardLimit = ResourceLimit l
   }