fix(cli): publish binary output atomically and preserve raw input

fun read path =

  let

    val stream = if path = "-" then TextIO.stdIn else TextIO.openIn path

    val text = TextIO.inputN (stream, 1048577)

    val _ = if path = "-" then () else TextIO.closeIn stream

    val fd = if path = "-" then Posix.FileSys.stdin

      else Posix.FileSys.openf (path, Posix.FileSys.O_RDONLY, Posix.FileSys.O.flags [])

    fun close () = if path = "-" then () else Posix.IO.close fd

    fun loop total chunks =

      let

        val bytes = Posix.IO.readVec (fd, Int.min (65536, 1048577 - total))

        val count = Word8Vector.length bytes

      in

        if count = 0 then String.concat (List.rev chunks)

        else if total + count > 1048576 then raise Fail "input exceeds 1 MiB"

        else loop (total + count) (Byte.bytesToString bytes :: chunks)

      end

    val input = loop 0 [] handle exn => ((close () handle _ => ()); raise exn)

    val _ = close ()

  in

    if size text > 1048576 then raise Fail "input exceeds 1 MiB" else text

    input

  end;

fun write path text =

  let

    val stream = BinIO.openOut path

    val _ = BinIO.output (stream, Byte.stringToBytes text)

  in BinIO.closeOut stream end;

    val bytes = Byte.stringToBytes text

    val pid = SysWord.toString (Posix.Process.pidToWord (Posix.ProcEnv.getpid ()))

    val prefix = OS.Path.concat (OS.Path.dir path, "." ^ OS.Path.file path ^ "." ^ pid ^ ".")

    val defaultMode = Posix.FileSys.S.flags [Posix.FileSys.S.irusr, Posix.FileSys.S.iwusr,

      Posix.FileSys.S.irgrp, Posix.FileSys.S.iwgrp, Posix.FileSys.S.iroth, Posix.FileSys.S.iwoth]

    val mode = Posix.FileSys.S.intersect [Posix.FileSys.ST.mode (Posix.FileSys.stat path),

      Posix.FileSys.S.flags [Posix.FileSys.S.irwxu, Posix.FileSys.S.irwxg, Posix.FileSys.S.irwxo]]

      handle exn as OS.SysErr (_, SOME error) =>

        if error = Posix.Error.noent then defaultMode else raise exn

    fun reserve n =

      if n = 128 then raise Fail "cannot reserve output file"

      else

        let

          val temporary = prefix ^ Int.toString n

        in

          (temporary, Posix.FileSys.createf

            (temporary, Posix.FileSys.O_WRONLY, Posix.FileSys.O.excl, mode))

        end handle exn as OS.SysErr (_, SOME error) =>

          if error = Posix.Error.exist then reserve (n + 1) else raise exn

    val (temporary, fd) = reserve 0

    val closed = ref false

    fun close () = (Posix.IO.close fd; closed := true)

    fun discard () =

      ((if !closed then () else close ()) handle _ => ();

       OS.FileSys.remove temporary handle _ => ())

    fun loop offset =

      if offset = Word8Vector.length bytes then ()

      else

        let

          val count = Posix.IO.writeVec (fd, Word8VectorSlice.slice (bytes, offset, NONE))

        in

          if count = 0 then raise Fail "output write made no progress" else loop (offset + count)

        end

  in

    (loop 0; close (); OS.FileSys.rename {old = temporary, new = path})

      handle exn => (discard (); raise exn)

  end;

fun output (ok, message) =

  (TextIO.output (if ok then TextIO.stdOut else TextIO.stdErr, text message);

import itertools

import os

from pathlib import Path

import random

import resource

import signal

import stat

import subprocess

import tempfile

    CHECKS += 1

def cli(*args, text=None, ok=True):

def cli(*args, text=None, ok=True, preexec_fn=None):

    result = subprocess.run(

        [str(ROOT / "span"), *map(str, args)],

        input=text,

        capture_output=True,

        timeout=90,

        cwd=ROOT,

        preexec_fn=preexec_fn,

    )

    check((result.returncode == 0) == ok, (args, result.stdout, result.stderr))

    return result.stdout if ok else result.stderr

def restrict_output():

    signal.signal(signal.SIGXFSZ, signal.SIG_IGN)

    resource.setrlimit(resource.RLIMIT_FSIZE, (1, 1))

def source(program):

    lines = []

    for i, (op, target) in enumerate(program):

        check(out.read_bytes() == b"preserve", "failed assembly clobbered output")

        cli("assemble", "-", "-o", out, text="")

        check(out.read_bytes() == b"", "empty binary output")

        out.write_bytes(b"preserve")

        before = set(Path(directory).iterdir())

        cli("assemble", "examples/countdown.span", "-o", out, ok=False, preexec_fn=restrict_output)

        check(out.read_bytes() == b"preserve", "failed write clobbered output")

        check(set(Path(directory).iterdir()) == before, "failed write leaked temporary file")

        destination = Path(directory) / "destination"

        destination.mkdir()

        before = set(Path(directory).iterdir())

        cli("assemble", "examples/countdown.span", "-o", destination, ok=False)

        check(destination.is_dir(), "failed rename replaced directory")

        check(set(Path(directory).iterdir()) == before, "failed rename leaked temporary file")

        same = Path(directory) / "source"

        same.write_text("inc\nhalt\n")

        cli("assemble", same, "-o", same)

        check(same.read_bytes() == bytes([1, 3]), "assembly cannot replace its source")

        link = Path(directory) / "link"

        link.symlink_to(out)

        cli("assemble", "-", "-o", link, text="halt")

        check(out.read_bytes() == b"preserve", "output followed destination symlink")

        check(not link.is_symlink() and link.read_bytes() == bytes([3]), "symlink not atomically replaced")

        out.write_bytes(b"".join(bytes([18, n, 255 - n]) for n in range(256)) + bytes([3]))

        listing = cli("disassemble", out)

        raw = subprocess.run(

            [str(ROOT / "span"), "disassemble", "-"],

            input=out.read_bytes(), capture_output=True, timeout=90, cwd=ROOT,

        )

        check(raw.returncode == 0 and raw.stdout == listing.encode("ascii"), "binary stdin changed bytes")

        check(len(listing.splitlines()) == 257, "binary file changed bytes")

        out.write_bytes(bytes(1048577))

        check("input exceeds 1 MiB" in cli("disassemble", out, ok=False), "binary size limit missing")

        check("input exceeds 1 MiB" in cli("assemble", "-", text=" " * 1048577, ok=False), "stdin size limit missing")

        out.chmod(0o600)

        cli("assemble", "-", "-o", out, text="halt", preexec_fn=lambda: os.umask(0))

        check(stat.S_IMODE(out.stat().st_mode) == 0o600, "replacement widened output permissions")

        protected = Path(directory) / "protected"

        protected.write_bytes(b"preserve")

        def collide_output():

            out.with_name(f".{out.name}.{os.getpid():X}.0").symlink_to(protected)

        cli("assemble", "-", "-o", out, text="inc\nhalt", preexec_fn=collide_output)

        check(out.read_bytes() == bytes([1, 3]), "temporary collision prevented output")

        check(protected.read_bytes() == b"preserve", "temporary collision followed symlink")

        collisions = list(Path(directory).glob(f".{out.name}.*"))

        check(len(collisions) == 1 and collisions[0].is_symlink(), "temporary collision was removed or output leaked")

    print(f"{CHECKS} checks passed, including 36 exhaustive layout comparisons")