{-# LANGUAGE CPP #-}

-- | Bundles: a self-contained executable that is unisonrt and a compiled
-- program (unisonrt/README.md, "Bundles"). ucm's `bundle` command and
-- `unisonrt bundle` write them; unisonrt, finding a program in its own
-- executable, runs it.
--
-- The program goes where the executable's format allows:
--
-- * Mach-O (macOS): in the section @__UNISON,__uc@, inside the image, so
--   that the code signature covers it ("Unison.Runtime.Bundle.MachO"). The
--   section holds 'bundleMagic', the program's length (8 bytes,
--   little-endian) and the program; in a plain unisonrt it holds 16 zero
--   bytes (unisonrt/bundle-section).
-- * Anything else (ELF, PE): after the end of the executable, followed by a
--   16-byte trailer, the program's offset (8 bytes, little-endian) and
--   'bundleMagic'. The loaders ignore bytes after the end of the file's
--   contents.
--
-- unisonrt's main.c tells a bundle apart the same way, to pass every
-- argument to the program.
module Unison.Runtime.Bundle
  ( readBundle,
    bundledProgram,
    writeBundle,
    signExecutable,
    makeExecutable,
  )
where

import Data.Bits (shiftL, shiftR)
import Data.ByteString qualified as B
import Data.Word (Word64)
import Unison.Runtime.Bundle.MachO qualified as MachO
#if defined(darwin_HOST_OS)
import Data.ByteString.Unsafe qualified as BU
import Data.Word (Word8)
import Foreign.C.Types (CSize (..))
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr, castPtr, nullPtr)
import Foreign.Storable (peek)
import System.Exit (ExitCode (..))
import System.Process (readProcessWithExitCode)
#else
import System.IO (IOMode (ReadMode), SeekMode (AbsoluteSeek), hFileSize, hSeek, withBinaryFile)
#endif
#if defined(mingw32_HOST_OS)
import System.FilePath (takeExtension, (<.>))
#else
import System.Posix.Files (setFileMode)
#endif

bundleMagic :: B.ByteString
bundleMagic :: ByteString
bundleMagic = ByteString
"UNISONRT"

-- | The size of the trailer, and of the head of the Mach-O section
trailerSize :: Int
trailerSize :: Int
trailerSize = Int
16

-- | The program in the bytes of an executable, if it is a bundle
bundledProgram :: B.ByteString -> Maybe B.ByteString
bundledProgram :: ByteString -> Maybe ByteString
bundledProgram ByteString
bytes
  | ByteString -> Bool
MachO.isMachO ByteString
bytes = ByteString -> Maybe ByteString
MachO.programSection ByteString
bytes Maybe ByteString
-> (ByteString -> Maybe ByteString) -> Maybe ByteString
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> Maybe ByteString
fromSection
  | Bool
otherwise = (Int, ByteString) -> ByteString
forall a b. (a, b) -> b
snd ((Int, ByteString) -> ByteString)
-> Maybe (Int, ByteString) -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteString -> Maybe (Int, ByteString)
fromTrailer ByteString
bytes

-- | The bytes of a bundle of the program, from the bytes of a unisonrt (a
-- bundle or not). A Mach-O result has to be signed again
-- ('signExecutable').
bundleBytes :: B.ByteString -> B.ByteString -> Either String B.ByteString
bundleBytes :: ByteString -> ByteString -> Either String ByteString
bundleBytes ByteString
runtime ByteString
program
  | ByteString -> Bool
MachO.isMachO ByteString
runtime = ByteString -> ByteString -> Either String ByteString
MachO.setProgramSection ByteString
runtime (ByteString -> ByteString
toSection ByteString
program)
  | Bool
otherwise = ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (ByteString
plain ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
program ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word64 -> ByteString
toLittleEndian (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
plain)) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
bundleMagic)
  where
    -- without the program of its own, if it is a bundle
    plain :: ByteString
plain = ByteString
-> ((Int, ByteString) -> ByteString)
-> Maybe (Int, ByteString)
-> ByteString
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ByteString
runtime ((Int -> ByteString -> ByteString
`B.take` ByteString
runtime) (Int -> ByteString)
-> ((Int, ByteString) -> Int) -> (Int, ByteString) -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, ByteString) -> Int
forall a b. (a, b) -> a
fst) (ByteString -> Maybe (Int, ByteString)
fromTrailer ByteString
runtime)

-- | The Mach-O section's contents for a program
toSection :: B.ByteString -> B.ByteString
toSection :: ByteString -> ByteString
toSection ByteString
program = ByteString
bundleMagic ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Word64 -> ByteString
toLittleEndian (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
program)) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
program

-- | A bundle's program, from the Mach-O section's contents
fromSection :: B.ByteString -> Maybe B.ByteString
fromSection :: ByteString -> Maybe ByteString
fromSection ByteString
section
  | ByteString
magic ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
bundleMagic, Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= ByteString -> Int
B.length ByteString
rest Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
8 = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (Int -> ByteString -> ByteString
B.take Int
len (Int -> ByteString -> ByteString
B.drop Int
8 ByteString
rest))
  | Bool
otherwise = Maybe ByteString
forall a. Maybe a
Nothing
  where
    (ByteString
magic, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
8 ByteString
section
    len :: Int
len = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Word64
fromLittleEndian (Int -> ByteString -> ByteString
B.take Int
8 ByteString
rest))

-- | Where the program starts and the program, from the trailer at the end
-- of the bytes
fromTrailer :: B.ByteString -> Maybe (Int, B.ByteString)
fromTrailer :: ByteString -> Maybe (Int, ByteString)
fromTrailer ByteString
bytes
  | ByteString -> Int
B.length ByteString
bytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
trailerSize,
    (ByteString
off, ByteString
magic) <- Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
8 (Int -> ByteString -> ByteString
B.drop (ByteString -> Int
B.length ByteString
bytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
trailerSize) ByteString
bytes),
    ByteString
magic ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
bundleMagic,
    Int
start <- Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Word64
fromLittleEndian ByteString
off),
    Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= ByteString -> Int
B.length ByteString
bytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
trailerSize =
      (Int, ByteString) -> Maybe (Int, ByteString)
forall a. a -> Maybe a
Just (Int
start, Int -> ByteString -> ByteString
B.take (ByteString -> Int
B.length ByteString
bytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
trailerSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start) (Int -> ByteString -> ByteString
B.drop Int
start ByteString
bytes))
  | Bool
otherwise = Maybe (Int, ByteString)
forall a. Maybe a
Nothing

-- | Writes a bundle of the program, from the file of a unisonrt (a bundle
-- or not), as an executable that anyone can run, and says where it wrote
-- it: on Windows a name without an extension gets @.exe@.
writeBundle :: FilePath -> B.ByteString -> FilePath -> IO (Either String FilePath)
writeBundle :: String -> ByteString -> String -> IO (Either String String)
writeBundle String
runtime ByteString
program String
out0 = do
  ByteString
bytes <- String -> IO ByteString
B.readFile String
runtime
  case ByteString -> ByteString -> Either String ByteString
bundleBytes ByteString
bytes ByteString
program of
    Left String
err -> Either String String -> IO (Either String String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String String
forall a b. a -> Either a b
Left String
err)
    Right ByteString
bundle -> do
      String -> ByteString -> IO ()
B.writeFile String
out ByteString
bundle
      String -> IO ()
makeExecutable String
out
      (() -> String) -> Either String () -> Either String String
forall a b. (a -> b) -> Either String a -> Either String b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (String -> () -> String
forall a b. a -> b -> a
const String
out) (Either String () -> Either String String)
-> IO (Either String ()) -> IO (Either String String)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Either String ())
signExecutable String
out
  where
#if defined(mingw32_HOST_OS)
    out = if null (takeExtension out0) then out0 <.> "exe" else out0
#else
    out :: String
out = String
out0
#endif

-- | Lets everyone run it, as the linker leaves an executable
makeExecutable :: FilePath -> IO ()
#if defined(mingw32_HOST_OS)
makeExecutable _ = pure ()
#else
makeExecutable :: String -> IO ()
makeExecutable String
out = String -> FileMode -> IO ()
setFileMode String
out FileMode
0o755
#endif

-- | Signs a rewritten executable again, ad hoc: macOS on arm64 runs
-- nothing whose signature doesn't match it. It can be signed again with a
-- real identity afterwards. Elsewhere there is no signature.
signExecutable :: FilePath -> IO (Either String ())
#if defined(darwin_HOST_OS)
signExecutable out = do
  (code, _, err) <- readProcessWithExitCode "codesign" ["--sign", "-", "--force", out] ""
  pure $ case code of
    ExitSuccess -> Right ()
    ExitFailure _ -> Left ("codesign failed: " ++ err)
#else
signExecutable :: String -> IO (Either String ())
signExecutable String
_ = Either String () -> IO (Either String ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> Either String ()
forall a b. b -> Either a b
Right ())
#endif

-- | The program this running executable carries, if it is a bundle. On
-- macOS the section is in memory already (cbits/bundle.c); elsewhere only
-- the trailer and the program are read from the file at this path.
readBundle :: FilePath -> IO (Maybe B.ByteString)
#if defined(darwin_HOST_OS)
readBundle _ = alloca \sizePtr -> do
  p <- c_bundleSection sizePtr
  size <- fromIntegral <$> peek sizePtr
  if p == nullPtr || size < trailerSize
    then pure Nothing
    else -- copied, so that nothing points into the image
      fromSection . B.copy <$> BU.unsafePackCStringLen (castPtr p, size)

foreign import ccall unsafe "unisonrt_bundle_section"
  c_bundleSection :: Ptr CSize -> IO (Ptr Word8)
#else
readBundle :: String -> IO (Maybe ByteString)
readBundle String
self = String
-> IOMode
-> (Handle -> IO (Maybe ByteString))
-> IO (Maybe ByteString)
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withBinaryFile String
self IOMode
ReadMode \Handle
h -> do
  Integer
size <- Handle -> IO Integer
hFileSize Handle
h
  if Integer
size Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
trailerSize
    then Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing
    else do
      Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
h SeekMode
AbsoluteSeek (Integer
size Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
trailerSize)
      (ByteString
off, ByteString
magic) <- Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
8 (ByteString -> (ByteString, ByteString))
-> IO ByteString -> IO (ByteString, ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Handle -> Int -> IO ByteString
B.hGet Handle
h Int
trailerSize
      let start :: Integer
start = Word64 -> Integer
forall a. Integral a => a -> Integer
toInteger (ByteString -> Word64
fromLittleEndian ByteString
off)
      if ByteString
magic ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
/= ByteString
bundleMagic Bool -> Bool -> Bool
|| Integer
start Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
size Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
trailerSize
        then Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing
        else do
          Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
h SeekMode
AbsoluteSeek Integer
start
          ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> IO ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Handle -> Int -> IO ByteString
B.hGet Handle
h (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
size Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
trailerSize Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
start))
#endif

toLittleEndian :: Word64 -> B.ByteString
toLittleEndian :: Word64 -> ByteString
toLittleEndian Word64
n = [Word8] -> ByteString
B.pack [Word64 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
n Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k)) | Int
k <- [Int
0 .. Int
7]]

fromLittleEndian :: B.ByteString -> Word64
fromLittleEndian :: ByteString -> Word64
fromLittleEndian = (Word8 -> Word64 -> Word64) -> Word64 -> ByteString -> Word64
forall a. (Word8 -> a -> a) -> a -> ByteString -> a
B.foldr' (\Word8
b Word64
acc -> Word64
acc Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
8 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b) Word64
0