s2sx: Limit max executable size

s2/README.md

@@ -230,7 +230,8 @@ Usage: s2sx [options] file1 file2
 
 Compresses all files supplied as input separately.
 If files have '.s2' extension they are assumed to be compressed already.
-Output files are written as 'filename.s2sfx' and with '.exe' for windows targets.
+Output files are written as 'filename.s2sx' and with '.exe' for windows targets.
+If output is big, an additional file with ".more" is written. This must be included as well.
 By default output files will be overwritten.
 
 Wildcards are accepted: testdir/*.txt will compress all files in testdir ending with .txt
@@ -244,6 +245,8 @@ Options:
         Compress using this amount of threads (default 32)
   -help
         Display help
+  -max string
+        Maximum executable size. Rest will be written to another file. (default "1G")
   -os string
         Destination operating system (default "windows")
   -q    Don't write any output to terminal, except errors
@@ -267,6 +270,17 @@ Available platforms are:
  * windows-386
  * windows-amd64                                                                             
 
+By default, there is a size limit of 1GB for the output executable.
+
+When this is exceeded the remaining file content is written to a file called
+output+`.more`. This file must be included for a successful extraction and 
+placed alongside the executable for a successful extraction.
+
+This file *must* have the same name as the executable, so if the executable is renamed, 
+so must the `.more` file. 
+
+This functionality is disabled with stdin/stdout. 
+
 ### Self-extracting TAR files
 
 If you wrap a TAR file you can specify `-untar` to make it untar on the destination host.

s2/cmd/_s2sx/_unpack/main.go

@@ -56,6 +56,13 @@ func main() {
 	exitErr(err)
 	rd, err := newReader(f, stat.Size())
 	exitErr(err)
+	f2, err := os.Open(me + ".more")
+	if err == nil {
+		rd = io.MultiReader(rd, f2)
+	}
+	if !os.IsNotExist(err) {
+		exitErr(err)
+	}
 	var tmp [1]byte
 	_, err = io.ReadFull(rd, tmp[:])
 	exitErr(err)
@@ -69,8 +76,8 @@ func main() {
 	switch tmp[0] {
 	case opUnpack:
 		outname := me + "-extracted"
-		if idx := strings.Index(me, ".s2sfx"); idx > 0 {
-			// Trim from '.s2sfx'
+		if idx := strings.Index(me, ".s2sx"); idx > 0 {
+			// Trim from '.s2sx'
 			outname = me[:idx]
 		}
 		var out io.Writer

s2/cmd/_s2sx/main.go

@@ -14,9 +14,11 @@ import (
 	"path"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
+	"unicode"
 
 	"github.com/klauspost/compress/s2"
 	"github.com/klauspost/compress/s2/cmd/internal/readahead"
@@ -31,6 +33,7 @@ var (
 	goos   = flag.String("os", runtime.GOOS, "Destination operating system")
 	goarch = flag.String("arch", runtime.GOARCH, "Destination architecture")
 	cpu    = flag.Int("cpu", runtime.GOMAXPROCS(0), "Compress using this amount of threads")
+	max    = flag.String("max", "1G", "Maximum executable size. Rest will be written to another file.")
 	safe   = flag.Bool("safe", false, "Do not overwrite output files")
 	stdout = flag.Bool("c", false, "Write all output to stdout. Multiple input files will be concatenated")
 	remove = flag.Bool("rm", false, "Delete source file(s) after successful compression")
@@ -48,6 +51,8 @@ var embeddedFiles embed.FS
 func main() {
 	flag.Parse()
 	args := flag.Args()
+	sz, err := toSize(*max)
+	exitErr(err)
 	if len(args) == 0 || *help {
 		_, _ = fmt.Fprintf(os.Stderr, "s2sx v%v, built at %v.\n\n", version, date)
 		_, _ = fmt.Fprintf(os.Stderr, "Copyright (c) 2011 The Snappy-Go Authors. All rights reserved.\n"+
@@ -58,7 +63,8 @@ Compresses all files supplied as input separately.
 Use file name - to read from stdin and write to stdout.
 If input is already s2 compressed it will just have the executable wrapped.
 Use s2c commandline tool for advanced option, eg 'cat file.txt||s2c -|s2sx - >out.s2sx'
-Output files are written as 'filename.s2sfx' and with '.exe' for windows targets.
+Output files are written as 'filename.s2sx' and with '.exe' for windows targets.
+If output is big, an additional file with ".more" is written. This must be included as well.
 By default output files will be overwritten.
 
 Wildcards are accepted: testdir/*.txt will compress all files in testdir ending with .txt
@@ -94,6 +100,10 @@ Options:`)
 	exec, err = ioutil.ReadAll(s2.NewReader(bytes.NewBuffer(exec)))
 	exitErr(err)
 
+	written := int64(0)
+	if int64(len(exec))+1 >= sz {
+		exitErr(fmt.Errorf("max size less than unpacker. Max size must be at least %d bytes", len(exec)+1))
+	}
 	mode := byte(opUnpack)
 	if *untar {
 		mode = opUnTar
@@ -107,6 +117,7 @@ Options:`)
 		exitErr(err)
 		_, err = os.Stdout.Write([]byte{mode})
 		exitErr(err)
+		written += int64(len(exec) + 1)
 	}
 
 	if stdIn {
@@ -139,7 +150,7 @@ Options:`)
 	for _, filename := range files {
 		func() {
 			var closeOnce sync.Once
-			dstFilename := fmt.Sprintf("%s%s", strings.TrimPrefix(filename, ".s2"), ".s2sfx")
+			dstFilename := fmt.Sprintf("%s%s", strings.TrimPrefix(filename, ".s2"), ".s2sx")
 			if *goos == "windows" {
 				dstFilename += ".exe"
 			}
@@ -172,7 +183,23 @@ Options:`)
 				dstFile, err := os.OpenFile(dstFilename, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0777)
 				exitErr(err)
 				defer dstFile.Close()
-				bw := bufio.NewWriterSize(dstFile, 4<<20*2)
+				sw := &switchWriter{w: dstFile, left: sz, close: nil}
+				sw.fn = func() {
+					dstFilename := dstFilename + ".more"
+					if *safe {
+						_, err := os.Stat(dstFilename)
+						if !os.IsNotExist(err) {
+							exitErr(fmt.Errorf("destination '%s' file exists", dstFilename))
+						}
+					}
+					dstFile, err := os.OpenFile(dstFilename, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0777)
+					exitErr(err)
+					sw.close = dstFile.Close
+					sw.w = dstFile
+					sw.left = (1 << 63) - 1
+				}
+				defer sw.Close()
+				bw := bufio.NewWriterSize(sw, 4<<20*2)
 				defer bw.Flush()
 				out = bw
 				_, err = out.Write(exec)
@@ -252,3 +279,58 @@ func isS2Input(rd io.Reader) (bool, io.Reader) {
 	exitErr(err)
 	return false, nil
 }
+
+// toSize converts a size indication to bytes.
+func toSize(size string) (int64, error) {
+	size = strings.ToUpper(strings.TrimSpace(size))
+	firstLetter := strings.IndexFunc(size, unicode.IsLetter)
+	if firstLetter == -1 {
+		firstLetter = len(size)
+	}
+
+	bytesString, multiple := size[:firstLetter], size[firstLetter:]
+	bytes, err := strconv.ParseInt(bytesString, 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("unable to parse size: %v", err)
+	}
+
+	switch multiple {
+	case "G", "GB", "GIB":
+		return bytes * 1 << 20, nil
+	case "M", "MB", "MIB":
+		return bytes * 1 << 20, nil
+	case "K", "KB", "KIB":
+		return bytes * 1 << 10, nil
+	case "B", "":
+		return bytes, nil
+	default:
+		return 0, fmt.Errorf("unknown size suffix: %v", multiple)
+	}
+}
+
+type switchWriter struct {
+	w     io.Writer
+	left  int64
+	fn    func()
+	close func() error
+}
+
+func (w *switchWriter) Write(b []byte) (int, error) {
+	if int64(len(b)) <= w.left {
+		w.left -= int64(len(b))
+		return w.w.Write(b)
+	}
+	n, err := w.w.Write(b[:w.left])
+	if err != nil {
+		return n, err
+	}
+	w.fn()
+	n2, err := w.Write(b[n:])
+	return n + n2, err
+}
+func (w *switchWriter) Close() error {
+	if w.close == nil {
+		return nil
+	}
+	return w.close()
+}