-- | This module implements secure multi-execution.
module SME.ME where

import Control.Concurrent (MVar, forkIO,putMVar,takeMVar,newEmptyMVar)
import Control.Concurrent.Chan
import Prelude hiding (readFile,writeFile)
import qualified System.IO as IO (readFile,writeFile,putStrLn)
import SME.Lattice


-- | Type class to specify security policies for programs run under secure multi-execution. 
class FiniteLattice l => Policy l a b | a -> l b where 
      -- | It defines the security level assigned to an input or output of type @a@.
      level    :: a -> l
      -- | It assigns a default value (of type @b@) to an input of type @a@.
      defvalue :: a -> b 


-- | The multi-execution monad.
data ME a   =  Return a 
            |  Write FilePath String (ME a) 
            |  Read FilePath (String -> ME a)

instance Monad ME where
    return x                = Return x
    (Return x)       >>= f  = f x
    (Write file s p) >>= f  = Write file s (p >>= f)
    (Read file g)    >>= f  = Read file (\i -> g i >>= f)

-- | Secure operation for writing files.
writeFile        :: FilePath -> String -> ME ()
writeFile file s = Write file s (return ())

-- | Secure operation for reading files.
readFile         :: FilePath -> ME String
readFile file    = Read file return

-- | Interpreter for the 'ME' monad
run :: Policy l FilePath String => l -> ChanMatrix l -> ME a -> IO a
run _ _ (Return a)          = return a
run l c (Write file o t)  
    | level file == l       = do  IO.writeFile file o
                                  run l c t 
    | otherwise             = run l c t 
run l c (Read file f) 
    | level file ==  l      = do  x <- IO.readFile file
                                  broadcast c l file x 
                                  run l c (f x) 
    | sless (level file) l  = do  x <- reuseInput c l file
                                  run l c (f x)
    | otherwise             = run l c (f (defvalue file))

type ChanMatrix l = [(l,[(l,Chan (FilePath,String))])]  


-- | Data type to set the security lattice to be used by function 'sme'.
data SetLevel l = SetLevel

-- | Function to perform secure multi-execution. The first argument is only there for type-checking purposes. 
sme :: Policy l FilePath String => SetLevel l -> ME a -> IO ()
sme _ t = do
  c     <- newChanMatrix
  sync  <- newSyncVars

  sequence_ $ do  (l,s) <- sync
                  return (forkIO (do _ <- run l c t ; putMVar s ()))
            
  sequence_ $ do  (_,s) <- sync
                  return (takeMVar s)

reuseInput :: Policy l FilePath String => ChanMatrix l -> l -> FilePath -> IO String
reuseInput cm l f = case lookup (level f) cm of
                      Nothing -> error "Not possible to reuse input!"
                      Just xs -> case lookup l xs of
                                   Nothing -> error "Not possible to reuse an input"
                                   Just c  -> reuseContents c f

reuseContents :: Chan (FilePath, String) -> FilePath -> IO String
reuseContents c f = do p@(f',s) <- readChan c
                       if f==f' then return s
                                else do s' <- reuseContents c f
                                        unGetChan c p
                                        return s'

broadcast :: FiniteLattice l => ChanMatrix l -> l -> FilePath -> String -> IO ()
broadcast cm l f str = case lookup l cm of
                         Nothing -> return ()
                         Just xs -> mapM_ (\c -> writeChan (snd c) (f,str)) xs

newChanMatrix :: FiniteLattice l => IO (ChanMatrix l)
newChanMatrix = mapM newChanLevel universe
                
newChanLevel :: FiniteLattice l => l -> IO (l,[(l, Chan (FilePath,String))])
newChanLevel l = do ls <- mapM (\e -> newChan >>= \c -> return (e,c)) (filter (/=l) (upset l)) 
                    return (l,ls)

newSyncVars :: FiniteLattice l => IO [(l, MVar ())]
newSyncVars = mapM (\x -> newEmptyMVar >>= \v -> return (x,v)) universe