Security type system

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE OverloadedStrings #-}

module Main
  ( main
  ) where

import Control.Comonad (extract)
import Control.Comonad.Cofree (Cofree (..))
import Control.Monad (unless)
import Control.Monad.Error.Class (throwError)
import Data.Map (Map)
import Data.Semigroup ((<>))
import Data.Set (Set)
import Data.Text (Text)

import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text as Text
import qualified Data.Text.IO as Text.IO

type Environment = Map Text SecurityClass

data Term a
  = Variable Text
  | Assign Text a a
  | If a a a
  deriving (Functor)

type SecurityClass = Set Text

check :: Environment -> Cofree Term () -> Either Text (Cofree Term SecurityClass)
check env (() :< Variable name)
  | Just sc <- Map.lookup name env =
      pure $ sc :< Variable name
  | otherwise = throwError $ "No such variable: " <> name
check env (() :< Assign variable value continuation)
  | Just variableSC <- Map.lookup variable env = do
      value' <- check env value
      let valueSC = extract value'
      unless (valueSC `Set.isSubsetOf` variableSC) $
        throwError $ "Confidentiality violation in Assign: " <>
                       tshow (variable, valueSC, variableSC)
      continuation' <- check env continuation
      let sc = variableSC `Set.union` extract continuation'
      pure $ sc :< Assign variable value' continuation'
  | otherwise = throwError $ "No such variable: " <> variable
check env (() :< If condition consequent alternative) = do
  condition' <- check env condition
  consequent' <- check env consequent
  alternative' <- check env alternative
  let conditionSC = extract condition'
  let consequentSC = extract consequent'
  let alternativeSC = extract alternative'
  unless (conditionSC `Set.isSubsetOf` consequentSC &&
          conditionSC `Set.isSubsetOf` alternativeSC) $
        throwError $ "Confidentiality violation in If: " <>
                       tshow (conditionSC, consequentSC, alternativeSC)
  let sc = consequentSC `Set.intersection` alternativeSC
  pure $ sc :< If condition' consequent' alternative'

prettyTerm :: (a -> Text) -> Cofree Term a -> Text
prettyTerm prettyAnn (ann :< Variable name) =
  name <> "<" <> prettyAnn ann <> ">"
prettyTerm prettyAnn (ann :< Assign variable value continuation) =
  "assign<" <> prettyAnn ann <> "> " <> variable <>
  " := " <> prettyTerm prettyAnn value <>
  ";\n" <> prettyTerm prettyAnn continuation
prettyTerm prettyAnn (ann :< If condition consequent alternative) =
  "if<" <> prettyAnn ann <> "> " <> prettyTerm prettyAnn condition <>
  " then\n" <> indent (prettyTerm prettyAnn consequent) <>
  "\nelse\n" <> indent (prettyTerm prettyAnn alternative) <>
  "\nend"

prettySecurityClass :: SecurityClass -> Text
prettySecurityClass = Text.intercalate ", " . Set.toList

indent :: Text -> Text
indent = Text.intercalate "\n" . fmap ("  " <>) . Text.splitOn "\n"

tshow :: Show a => a -> Text
tshow = Text.pack . show

main :: IO ()
main = do
  example $ variable "a"
  example $ variable "b"
  example $ variable "c"
  example $ if_ (variable "a") (variable "b") (variable "c")
  example $ if_ (variable "a") (variable "c") (variable "d")
  example $ if_ (variable "b") (variable "c") (variable "d")
  where
    variable name = () :< Variable name
    assign name value continuation = () :< Assign name value continuation
    if_ condition consequent alternative = () :< If condition consequent alternative

    example :: Cofree Term () -> IO ()
    example term = do
      Text.IO.putStrLn $ prettyTerm (const "") term
      Text.IO.putStrLn $ ""
      Text.IO.putStrLn . either id (prettyTerm prettySecurityClass) $
        check environment term
      Text.IO.putStrLn $ "---------------"

    environment :: Environment
    environment = Map.fromList
      [ ("a", Set.fromList [])
      , ("b", Set.fromList ["α"])
      , ("c", Set.fromList ["β", "γ"])
      , ("d", Set.fromList ["γ"])
      ]