{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-
This module contains the executables for the Lambda function.
-}
module Control.Distributed.Fork.AWS.Lambda.Internal.Archive
  ( Archive (..),
    mkArchive,
    archiveSize,
    archiveChecksum,
  )
where

import Codec.Archive.Zip hiding (Archive)
import Control.Distributed.Fork.AWS.Lambda.Internal.Constants
import Control.Distributed.Fork.Backend
import Control.Exception
import Control.Monad
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import Data.Digest.Pure.SHA
import Data.Elf
import Data.Function
import Data.String.Interpolate
import qualified Data.Text as T
import qualified Data.Text.Encoding as T

{-
Since AWS Lambda does not support binary execution by default, our entry point
is a small Python script, whose only purpose is to execute the attached Haskell
binary, provide the input from standard input and return the standard output to
the queue.
-}
handlerPy :: BS.ByteString
handlerPy :: ByteString
handlerPy =
  Text -> ByteString
T.encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$
    String -> Text
T.pack
      [i|
import os
import subprocess
from uuid import uuid4
from json import dumps
from base64 import b64encode, b64decode

import boto3

queue_url = os.environ["#{envAnswerQueueUrl}"]
bucket_url = os.environ["#{envAnswerBucketUrl}"]

sqs = boto3.client('sqs')
s3 = boto3.client('s3')

def send_message(body):
    sqs.send_message(
      QueueUrl=queue_url,
      MessageBody=dumps(body),
    )

def handle(event, context):
    id = event["i"]

    popen = subprocess.Popen(
       ["./#{hsMainName}", "#{argExecutorMode}"],
       stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    (out, _) = popen.communicate(b64decode(event["d"]))

    if popen.returncode:
      print "Subprocess failed, code: ", popen.returncode
      exit(1)

    if len(out) < 200000:
      send_message({
        "id": id,
        "type": "response-inline",
        "payload": b64encode(out)
      })
    else:
      fname = str(uuid4())
      s3.put_object(
        Bucket = bucket_url,
        Key = fname,
        Body = out
      )
      send_message({
        "id": id,
        "type": "response-s3",
        "path": fname
      })
|]

{-
And we read the current executable.

Since it'll run on AWS Lambda, it needs to be a statically linked Linux
executable, so we do a preliminary check here.
-}
mkHsMain :: IO BS.ByteString
mkHsMain :: IO ByteString
mkHsMain = do
  String
path <- IO String
getExecutablePath
  ByteString
contents <- String -> IO ByteString
BS.readFile String
path
  ByteString -> IO ()
assertBinary ByteString
contents
  ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
contents

assertBinary :: BS.ByteString -> IO ()
assertBinary :: ByteString -> IO ()
assertBinary ByteString
contents = do
  Elf
elf <-
    (Elf -> IO Elf
forall (m :: * -> *) a. Monad m => a -> m a
return (Elf -> IO Elf) -> Elf -> IO Elf
forall a b. (a -> b) -> a -> b
$! ByteString -> Elf
parseElf ByteString
contents)
      IO Elf -> (SomeException -> IO Elf) -> IO Elf
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(SomeException
_ :: SomeException) -> FileException -> IO Elf
forall e a. Exception e => e -> IO a
throwIO FileException
FileExceptionNotElf)
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Elf -> ElfClass
elfClass Elf
elf ElfClass -> ElfClass -> Bool
forall a. Eq a => a -> a -> Bool
== ElfClass
ELFCLASS64) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    FileException -> IO ()
forall e a. Exception e => e -> IO a
throwIO FileException
FileExceptionNot64Bit
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((ElfSegment -> Bool) -> [ElfSegment] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\ElfSegment
s -> ElfSegment -> ElfSegmentType
elfSegmentType ElfSegment
s ElfSegmentType -> ElfSegmentType -> Bool
forall a. Eq a => a -> a -> Bool
== ElfSegmentType
PT_DYNAMIC) (Elf -> [ElfSegment]
elfSegments Elf
elf)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    FileException -> IO ()
forall e a. Exception e => e -> IO a
throwIO FileException
FileExceptionNotStatic

data FileException
  = FileExceptionNotElf
  | FileExceptionNot64Bit
  | FileExceptionNotStatic

instance Exception FileException

instance Show FileException where
  show :: FileException -> String
show FileException
FileExceptionNotElf =
    [i|
    Error: I am not an ELF (Linux) binary.

    The executable will run on AWS environment, because of that
    this library currently only supports Linux.
    |]
  show FileException
FileExceptionNot64Bit =
    [i|
    Error: I am not a 64bit executable.

    AWS Lambda currently only runs 64 bit executables.
    |]
  show FileException
FileExceptionNotStatic =
    [i|
    Error: I am not a dynamic executable.

    Since the executable will run on AWS environment, it needs
    to be statically linked.

    You can give GHC "-optl-static -optl-pthread -fPIC" flags
    to statically compile executables.
    |]

{-
And we're going to put all of them in a zip archive.
-}
newtype Archive
  = Archive {Archive -> ByteString
archiveToByteString :: BS.ByteString}

mkArchive :: IO Archive
mkArchive :: IO Archive
mkArchive = do
  ByteString
hsMain <- IO ByteString
mkHsMain
  Archive -> IO Archive
forall (m :: * -> *) a. Monad m => a -> m a
return (Archive -> IO Archive)
-> (Archive -> Archive) -> Archive -> IO Archive
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Archive
Archive (ByteString -> Archive)
-> (Archive -> ByteString) -> Archive -> Archive
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BL.toStrict (ByteString -> ByteString)
-> (Archive -> ByteString) -> Archive -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Archive -> ByteString
fromArchive (Archive -> IO Archive) -> Archive -> IO Archive
forall a b. (a -> b) -> a -> b
$
    Archive
emptyArchive
      Archive -> (Archive -> Archive) -> Archive
forall a b. a -> (a -> b) -> b
& Entry -> Archive -> Archive
addEntryToArchive
        (String -> Integer -> ByteString -> Entry
toEntry String
handlerPyName Integer
0 (ByteString -> Entry) -> ByteString -> Entry
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.fromStrict ByteString
handlerPy)
      Archive -> (Archive -> Archive) -> Archive
forall a b. a -> (a -> b) -> b
& Entry -> Archive -> Archive
addEntryToArchive
        (String -> Integer -> ByteString -> Entry
toEntry String
hsMainName Integer
0 (ByteString -> Entry) -> ByteString -> Entry
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.fromStrict ByteString
hsMain)
          { eExternalFileAttributes :: Word32
eExternalFileAttributes = Word32
0b10000 {- rwx -}
          }

archiveSize :: Archive -> Integer
archiveSize :: Archive -> Integer
archiveSize = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> (Archive -> Int) -> Archive -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Int
BS.length (ByteString -> Int) -> (Archive -> ByteString) -> Archive -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Archive -> ByteString
archiveToByteString

archiveChecksum :: Archive -> T.Text
archiveChecksum :: Archive -> Text
archiveChecksum =
  String -> Text
T.pack (String -> Text) -> (Archive -> String) -> Archive -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Digest SHA1State -> String
forall t. Digest t -> String
showDigest (Digest SHA1State -> String)
-> (Archive -> Digest SHA1State) -> Archive -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Digest SHA1State
sha1 (ByteString -> Digest SHA1State)
-> (Archive -> ByteString) -> Archive -> Digest SHA1State
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BL.fromStrict (ByteString -> ByteString)
-> (Archive -> ByteString) -> Archive -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Archive -> ByteString
archiveToByteString