main.go

// The pkgzip command bundles assets into a Go package.
//
// For more information see pkgzip -help
package main // import "code.soquee.net/pkgzip"

import (
	"archive/zip"
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"
	"time"
)

// mtimeDate holds the arbitrary mtime that we assign to files when
// flagNoMtime is set.
var mtimeDate = time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)

const pkgTmpl = `// Package {{.Pkg}} is a virtual filesystem generated by the pkgzip command.
//
// To use this filesystem as an http.FileSystem see
// golang.org/x/tools/godoc/vfs/httpfs
package {{.Pkg}}

import (
	"archive/zip"
	"io"
	"strings"

	"golang.org/x/tools/godoc/vfs"
	"golang.org/x/tools/godoc/vfs/zipfs"
)

var zipData = "{{.ByteLit}}"

// New returns the embedded filesystem.
func New() vfs.FileSystem {
	br := strings.NewReader(zipData)
	/* #nosec */
	r, _ := zip.NewReader(br, int64(br.Len()))
	rc := &zip.ReadCloser{
		Reader: *r,
	}
	return zipfs.New(rc, "{{.Pkg}}")
}

// Reader returns a new reader that reads the mbedded filesystem in zip format.
func Reader() io.Reader {
	return strings.NewReader(zipData)
}`

// Some exit codes. See sysexits(3).
const (
	exCantCreat = 73
	exDataErr   = 65
	exNoPerm    = 77
	exSoftware  = 70
	exUsage     = 64
)

func main() {
	logger := log.New(os.Stderr, "pkgzip ", 0)
	debug := log.New(ioutil.Discard, "", 0)

	var (
		pkgName    = "pkgzip"
		pkgPath    = "internal"
		src        = ""
		force      = false
		noCompress = false
		noMtime    = false
		v          = false
	)
	flags := flag.NewFlagSet("pkgzip", flag.ContinueOnError)
	flags.StringVar(&src, "src", src, "A directory to load files from")
	flags.StringVar(&pkgName, "pkg", pkgName, "The name of the generated package")
	flags.StringVar(&pkgPath, "dest", pkgPath, "The relative or absolute path to the generated package")
	flags.BoolVar(&force, "f", force, "Overwrite the destination tree if it already exists.")
	flags.BoolVar(&noCompress, "Z", noCompress, "Do not use compression to shrink the files.")
	flags.BoolVar(&noMtime, "m", noMtime, "Ignore modification times on files.")
	flags.BoolVar(&v, "v", v, "Use verbose error logging.")
	err := flags.Parse(os.Args[1:])
	if err != nil {
		if err != flag.ErrHelp {
			logger.Println(err)
		}
		os.Exit(exUsage)
	}
	if v {
		debug = logger
	}

	pkgPath = path.Join(pkgPath, pkgName)
	if pkgName == "" || src == "" {
		logger.Println("Package name or src tree must not be empty.")
		os.Exit(exDataErr)
	}
	if pkgPath == "" || pkgPath == "/" {
		logger.Println("Package path should not be root. This could cause very bad things to happen.")
		os.Exit(exDataErr)
	}

	tmp, err := ioutil.TempFile("", "pkgzip")
	if err != nil {
		logger.Printf("Error creating temp file: %q", err)
		os.Exit(exCantCreat)
	}

	b, err := genZIP(debug, logger, noMtime, noCompress, src)
	if err != nil {
		logger.Printf("Error creating filesystem: %q", err)
		os.Exit(exSoftware)
	}

	buf := &bytes.Buffer{}
	err = template.Must(template.New("pkgzip").Parse(pkgTmpl)).Execute(buf, struct {
		Pkg     string
		ByteLit string
	}{
		Pkg:     pkgName,
		ByteLit: b,
	})
	if err != nil {
		logger.Printf("Error rendering generated code: %q", err)
		os.Exit(exDataErr)
	}
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, "", buf.String(), parser.ParseComments)
	if err != nil {
		logger.Printf("Error parsing generated code: %q", err)
		os.Exit(exDataErr)
	}
	err = format.Node(tmp, fset, node)
	if err != nil {
		logger.Printf("Error formatting generated code: %q", err)
		os.Exit(exDataErr)
	}

	err = tmp.Close()
	if err != nil {
		debug.Printf("Error closing temp file")
	}

	if force {
		// If this fails, it either doesn't exist in which case everything is fine,
		// or it does exist but we don't have permission to remove it but we might
		// as well defer error handling until Rename tries to overwrite it and
		// fails.
		err = os.RemoveAll(pkgPath)
		if err != nil {
			debug.Printf("Error removing package at %q: %q", pkgPath, err)
		}
	}
	err = os.MkdirAll(pkgPath, 0750)
	if err != nil {
		debug.Printf("Error creating package at %q: %q", pkgPath, err)
	}
	// Move the temp package to its final destination.
	err = rename(debug, force, tmp.Name(), path.Join(pkgPath, "pkgzip.go"))
	if err != nil {
		logger.Printf("Error creating package at %q: %q", pkgPath, err)
		err = os.RemoveAll(tmp.Name())
		if err != nil {
			logger.Printf("Error cleaning up temp files at %q: %q", tmp.Name(), err)
		}
		os.Exit(exNoPerm) // Probably… this may not be the actual problem.
	}
}

// rename tries to os.Rename, but fall backs to copying from src
// to dest and unlink the source if os.Rename fails.
//
// This function was copied from statik and is used under the terms of the
// Apache license:
//
// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
func rename(debug *log.Logger, force bool, src, dest string) error {
	// Try to rename generated source.
	if err := os.Rename(src, dest); err == nil {
		return nil
	}
	// If the rename failed (might do so due to temporary file residing on a
	// different device), try to copy byte by byte.
	/* #nosec */
	rc, err := os.Open(src)
	if err != nil {
		return err
	}
	defer func() {
		err = rc.Close()
		if err != nil {
			debug.Printf("Error closing %q: %q", src, err)
		}
		err = os.Remove(src)
		if err != nil {
			debug.Printf("Error removing src %q: %q", src, err)
		}
	}()

	if _, err = os.Stat(dest); !os.IsNotExist(err) {
		if force {
			if err = os.Remove(dest); err != nil {
				return fmt.Errorf("file %q could not be deleted", dest)
			}
		} else {
			return fmt.Errorf("file %q already exists; use -f to overwrite", dest)
		}
	}

	wc, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer wc.Close()

	if _, err = io.Copy(wc, rc); err != nil {
		// Delete remains of failed copy attempt.
		err = os.Remove(dest)
		if err != nil {
			debug.Printf("Error removing failed copy attempt: %q", err)
		}
	}
	return err
}

// Chunks of this function are copyright by Google and used under the terms of
// the Apache license. See rename() for details.
func genZIP(debug, logger *log.Logger, noMtime, noCompress bool, src string) (string, error) {
	buf := &bytes.Buffer{}

	w := zip.NewWriter(buf)
	err := filepath.Walk(src, func(path string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		// Ignore directories and hidden files.
		// No entry is needed for directories in a zip file.
		// Each file is represented with a path, no directory
		// entities are required to build the hierarchy.
		if fi.IsDir() || strings.HasPrefix(fi.Name(), ".") {
			return nil
		}
		relPath, err := filepath.Rel(src, path)
		if err != nil {
			return err
		}
		/* #nosec */
		b, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}
		fHeader, err := zip.FileInfoHeader(fi)
		if err != nil {
			return err
		}
		if noMtime {
			// Always use the same modification time so that
			// the output is deterministic with respect to the file contents.
			fHeader.SetModTime(mtimeDate)
		}
		fHeader.Name = filepath.ToSlash(relPath)
		if !noCompress {
			fHeader.Method = zip.Deflate
		}
		f, err := w.CreateHeader(fHeader)
		if err != nil {
			return err
		}
		_, err = f.Write(b)
		return err
	})
	if err != nil {
		return "", err
	}
	err = w.Close()
	s := &strings.Builder{}
	for _, b := range buf.Bytes() {
		if b == '\n' {
			_, err = s.WriteString(`\n`)
			if err != nil {
				logger.Println("Error writing to output string")
				os.Exit(exSoftware)
			}
			continue
		}
		if b == '\\' {
			_, err = s.WriteString(`\\`)
			if err != nil {
				logger.Println("Error writing to output string")
				os.Exit(exSoftware)
			}
			continue
		}
		if b == '"' {
			_, err = s.WriteString(`\"`)
			if err != nil {
				logger.Println("Error writing to output string")
				os.Exit(exSoftware)
			}
			continue
		}
		if (b >= 32 && b <= 126) || b == '\t' {
			err = s.WriteByte(b)
			if err != nil {
				logger.Println("Error writing byte to output string")
				os.Exit(exSoftware)
			}
			continue
		}
		fmt.Fprintf(s, "\\x%02x", b)
	}
	return s.String(), err
}