import CommandLine(getOptions) import Control.Concurrent(threadDelay) import Control.Lens(ASetter, set, view, elemOf, folded) import Control.Monad(unless, void, foldM) import qualified Data.ByteString.Lazy as L import Data.Int(Int64) import Data.List(sortBy) import Data.List.NonEmpty(NonEmpty(..)) import Data.Maybe(fromMaybe) import Data.Monoid((<>)) import Data.Text(Text, unpack, isPrefixOf) import Network.AWS(Env, AWSRequest, Rs, runResourceT, runAWS, send) import Network.AWS.Data.Body(RqBody(..), toHashed, contentLength) import Network.AWS.Data.Text(toText) import Network.AWS.EC2.DescribeImages(describeImages, deseOwners, deseFilters, diirsImages) import Network.AWS.EC2.DescribeImportSnapshotTasks(DescribeImportSnapshotTasksResponse, distImportTaskIds, describeImportSnapshotTasks, distrsResponseStatus, distrsImportSnapshotTasks) import Network.AWS.EC2.ImportSnapshot(importSnapshot, isDiskContainer, isrsSnapshotTaskDetail, isrsResponseStatus, isrsImportTaskId) import Network.AWS.EC2.RegisterImage(registerImage, riRootDeviceName, riBlockDeviceMappings, riKernelId, riArchitecture, rirsResponseStatus, rirsImageId) import Network.AWS.EC2.Types(istSnapshotTaskDetail, ebsBlockDevice, ebdDeleteOnTermination, ebdVolumeSize, ebdSnapshotId, bdmEBS, blockDeviceMapping, SnapshotTaskDetail, stdStatus, stdSnapshotId, stdProgress, stdStatusMessage, snapshotDiskContainer, sdcURL, sdcFormat, iName, iImageId, filter', fValues, ArchitectureValues(..)) import Network.AWS.S3.AbortMultipartUpload(abortMultipartUpload) import Network.AWS.S3.CompleteMultipartUpload(completeMultipartUpload, cMultipartUpload, crsResponseStatus) import Network.AWS.S3.CreateBucket(createBucket, cbCreateBucketConfiguration) import Network.AWS.S3.CreateMultipartUpload(createMultipartUpload, cmursUploadId, cmursResponseStatus) import Network.AWS.S3.ListBuckets(listBuckets, lbrsBuckets) import Network.AWS.S3.Types(completedMultipartUpload, cmuParts, completedPart, LocationConstraint(..), createBucketConfiguration, cbcLocationConstraint, bName) import Network.AWS.S3.UploadPart(uploadPart, uprsETag, uprsResponseStatus) import Options(Options, optS3Bucket, optTargetKey, optImageName, optS3Bucket, optTargetKey, optKernel, optKernelArgs, optRamdisks, optS3Bucket, optAwsRegion) import System.Directory(copyFile) import System.Environment(getArgs) import System.Exit(ExitCode(..), exitWith) import System.FilePath(takeFileName, ()) import System.IO.Temp(withSystemTempDirectory) import System.Posix.Files(fileSize, getFileStatus) import System.Process(callProcess) import System.Timeout(timeout) main :: IO () main = do args <- getArgs (opts, e) <- getOptions args kernelId <- findPVGrubAKI e makeBucket opts e image <- buildDisk opts uploadFile opts e image snapshot <- makeSnapshot opts e ami <- createAMI opts e kernelId snapshot putStrLn (unpack ami) -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- findPVGrubAKI :: Env -> IO Text findPVGrubAKI e = do res <- awsSend e findRequest let grubs = filter isPVGrub (view diirsImages res) case reverse (sortBy sortGrubs grubs) of [] -> abort "Couldn't find PV-GRUB kernel for your region." (kernel:_) -> do let kernelId = view iImageId kernel Just kernelName = view iName kernel logm "KERNEL" ("Using kernel " ++ unpack kernelId ++ " (" ++ unpack kernelName ++ ")") return kernelId where findRequest = set deseOwners ["amazon"] $ set deseFilters [ set fValues ["x86_64"] (filter' "architecture") , set fValues ["xen"] (filter' "hypervisor") , set fValues ["paravirtual"] (filter' "virtualization-type") ] $ describeImages -- isPVGrub x = case view iName x of Nothing -> False Just n -> "pv-grub-" `isPrefixOf` n -- sortGrubs x y = compare (view iName x) (view iName y) -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- buildDisk :: Options -> IO FilePath buildDisk opts = withSystemTempDirectory "ec2-unikernel" $ \ path -> do logm "DISK" ("Building disk.") sizeb0 <- fileSize `fmap` getFileStatus (view optKernel opts) sizeb <- foldM addSize sizeb0 (view optRamdisks opts) let sizem = ceiling (fromInteger (fromIntegral sizeb) / onemeg) + 1 writeFile (path "menu.lst") (grubMenu opts) writeFile (path "guestfish.scr") (guestFishScript (path "disk.raw") sizem (path "menu.lst") opts) callProcess "guestfish" ["-f", path "guestfish.scr"] let targetFile = unpack (toText (view optTargetKey opts)) copyFile (path "disk.raw") targetFile logm "DISK" ("Built disk " ++ targetFile) return targetFile where addSize inSize fileName = do fsize <- fileSize `fmap` getFileStatus fileName return (fsize + inSize) grubMenu :: Options -> String grubMenu opts = unlines ([ "default 0" , "timeout 1" , "title unikernel_boot" , "\troot (hd0,0)" , "\tkernel /" ++ takeFileName (view optKernel opts) ++ kernelArgs ] ++ map (\ f -> "\tmodule /" ++ takeFileName f) (view optRamdisks opts)) where kargs = view optKernelArgs opts kernelArgs | null kargs = "" | otherwise = " " ++ kargs guestFishScript :: String -> Integer -> FilePath -> Options -> String guestFishScript diskName diskSize menu opts = unlines ([ "disk-create " ++ diskName ++ " raw " ++ show diskSize ++ "M" , "add " ++ diskName , "run" , "part-disk /dev/sda mbr" , "mkfs ext2 /dev/sda1" , "mount /dev/sda1 /" , "mkdir /grub" , "sync" , "copy-in " ++ menu ++ " /grub/" , "rename /grub/" ++ takeFileName menu ++ " /grub/menu.lst" , "copy-in " ++ view optKernel opts ++ " /" ] ++ map (\ f -> "copy-in " ++ f ++ " /") (view optRamdisks opts) ++ [ "sync" , "exit" ]) onemeg :: Double onemeg = 1048576.0 -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- amazonMinimimPartSizeInMB :: Integer amazonMinimimPartSizeInMB = 5 amazonMinimimPartSizeInBytes :: Int64 amazonMinimimPartSizeInBytes = fromIntegral amazonMinimimPartSizeInMB * 1024 * 1024 statusOk :: Int statusOk = 200 makeBucket :: Options -> Env -> IO () makeBucket opts e = do bkts <- awsSend e listBuckets unless (elemOf (lbrsBuckets . folded . bName) bucketName bkts) $ do void (awsSend e createRequest) logm "S3" ("Created bucket " ++ unpack (toText bucketName)) where bucketName = view optS3Bucket opts location = Just (LocationConstraint (view optAwsRegion opts)) config = set cbcLocationConstraint location createBucketConfiguration createRequest = set cbCreateBucketConfiguration (Just config) $ createBucket bucketName uploadFile :: Options -> Env -> FilePath -> IO () uploadFile opts e filename = do logm "UPLOAD" "Creating upload." rsp <- awsSend e (createMultipartUpload bucket key) case (view cmursUploadId rsp, view cmursResponseStatus rsp) of (Nothing, code) -> putStrLn ("ERROR: Upload initialization failed with code: "++show code) (Just upId, _) -> do bytes <- L.readFile filename sizeb <- fileSize `fmap` getFileStatus filename logm "UPLOAD" "Starting upload." runUpload upId 0 (fromIntegral sizeb) bytes 1 [] where bucket = view optS3Bucket opts key = view optTargetKey opts -- abortUploadAndFail upId = do void (awsSend e (abortMultipartUpload bucket key upId)) fail "Execution aborted (bad upload)" -- runUpload upId sentSize totalSize bytes partNo completedTags | L.null bytes = case reverse completedTags of [] -> do putStrLn ("ERROR: Empty upload?") abortUploadAndFail upId (x:rest) -> do let parts = x :| rest baseReq = completeMultipartUpload bucket key upId parts' = set cmuParts (Just parts) completedMultipartUpload rsp <- awsSend e (set cMultipartUpload (Just parts') baseReq) if view crsResponseStatus rsp == statusOk then logm "UPLOAD" "100% complete" else do putStrLn "ERROR: Bad final upload code." abortUploadAndFail upId | otherwise = do let (chunkBS, rest) = L.splitAt amazonMinimimPartSizeInBytes bytes chunk = Hashed (toHashed chunkBS) sentSize' = sentSize + contentLength chunk req = uploadPart bucket key partNo upId chunk rsp <- awsSend e req case (view uprsETag rsp, view uprsResponseStatus rsp) of (Nothing, code) -> do putStrLn ("ERROR: Upload failed with code: " ++ show code) putStrLn (" Attempting abort.") abortUploadAndFail upId (Just etag, _) -> do let percentDble = fromIntegral sentSize / totalSize percentInt = ceiling ((100.0 :: Double) * percentDble) :: Int logm "UPLOAD" (show percentInt ++ "% complete") let partNo' = partNo + 1 completedTags' = completedPart partNo etag : completedTags runUpload upId sentSize' totalSize rest partNo' completedTags' -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- makeSnapshot :: Options -> Env -> IO Text makeSnapshot opts e = do res <- awsSend e importRequest let taskId = view isrsImportTaskId res errCode = view isrsResponseStatus res details = view isrsSnapshotTaskDetail res checkResponse (fromMaybe "" taskId) errCode details where url = "s3://" <> toText (view optS3Bucket opts) <> "/" <> toText (view optTargetKey opts) importRequest = setm isDiskContainer container importSnapshot container = setm sdcFormat "RAW" $ setm sdcURL url $ snapshotDiskContainer loop :: Text -> IO Text loop taskId = do threadDelay 2500000 res <- awsSend e (set distImportTaskIds [taskId] describeImportSnapshotTasks) checkResponse taskId (view distrsResponseStatus res) (getDetails res) getDetails :: DescribeImportSnapshotTasksResponse -> Maybe SnapshotTaskDetail getDetails rsp = case view distrsImportSnapshotTasks rsp of [] -> Nothing (x:_) -> view istSnapshotTaskDetail x checkResponse :: Text -> Int -> Maybe SnapshotTaskDetail -> IO Text checkResponse _ errCode Nothing = abort ("Could not import disk snapshot. Error code: " ++ show errCode) checkResponse taskId err (Just d) | err /= 200 = do warn ("Weird error code from import: " ++ show err) checkResponse taskId 200 (Just d) checkResponse taskId _ (Just d) = processTaskDetail taskId d processTaskDetail :: Text -> SnapshotTaskDetail -> IO Text processTaskDetail taskId detail | view stdStatus detail == Just "completed" = case view stdSnapshotId detail of Nothing -> abort "Snapshot imported, but no id found." Just x -> return x | otherwise = do logm "IMPORT" (show' (view stdProgress detail) ++ "% " ++ "(" ++ show' (view stdStatusMessage detail) ++ ")") loop taskId -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- createAMI :: Options -> Env -> Text -> Text -> IO Text createAMI opts e kernelId snapshot = do res <- awsSend e registerRequest case (view rirsResponseStatus res, view rirsImageId res) of (errCode, Nothing) -> abort ("Could not register AMI. Error code: " ++ show errCode) (200, Just ami) -> do logm "IMPORT" "100%" return ami (err, Just ami) -> do warn ("Weird error code from register AMI: " ++ show err) return ami where ebsVol = setm ebdDeleteOnTermination True $ setm ebdVolumeSize 1 -- FIXME: Compute this $ setm ebdSnapshotId snapshot $ ebsBlockDevice blockDevMap = setm bdmEBS ebsVol $ blockDeviceMapping "/dev/sda1" registerRequest = setm riRootDeviceName "/dev/sda1" $ set riBlockDeviceMappings [blockDevMap] $ setm riArchitecture X86_64 $ setm riKernelId kernelId $ registerImage (view optImageName opts) -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- -- ----------------------------------------------------------------------------- abort :: String -> IO a abort msg = do putStrLn ("ERROR: " ++ msg) exitWith (ExitFailure 1) warn :: String -> IO () warn msg = putStrLn ("WARNING: " ++ msg) setm :: ASetter s t a (Maybe b) -> b -> s -> t setm field value obj = set field (Just value) obj logm :: String -> String -> IO () logm category message = putStrLn (category ++ ": " ++ message) awsSend :: AWSRequest r => Env -> r -> IO (Rs r) awsSend env request = go 5 request where go :: AWSRequest r => Int -> r -> IO (Rs r) go 0 _ = fail "Amazon request failed." go x v = do mrsp <- timeout (15 * 1000000) (runResourceT (runAWS env (send v))) case mrsp of Nothing -> putStrLn "Retrying." >> go (x - 1) v Just rsp -> return rsp show' :: Maybe Text -> String show' Nothing = "" show' (Just x) = unpack x