go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/tools/cmd/assets/main.go

// Copyright 2015 The LUCI Authors.
//
// 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.

// Package main hosts the utility that converts binary assets into assets.gen.go
// file, so that they can be baked directly into the executable. Intended to
// be used only for small files, like HTML templates.
//
// This utility is used via `go generate`. Corresponding incantation:
//
//	//go:generate assets
package main

import (
	"bytes"
	"crypto/sha256"
	"flag"
	"fmt"
	"go/build"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"text/template"
	"time"

	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/flag/fixflagpos"
	"go.chromium.org/luci/common/flag/stringlistflag"
)

// defaultExts lists glob patterns for files to put into generated
// *.go file.
var defaultExts = stringset.NewFromSlice(
	"*.css",
	"*.html",
	"*.js",
	"*.tmpl",
)

// funcMap contains functions used when rendering assets.gen.go template.
var funcMap = template.FuncMap{
	"asByteArray": asByteArray,
}

// assetsGenGoTmpl is template for generated assets.gen.go file. Result of
// the execution will also be passed through gofmt.
var assetsGenGoTmpl = template.Must(template.New("tmpl").Funcs(funcMap).Parse(strings.TrimSpace(`
// Copyright {{.Year}} The LUCI Authors.
//
// 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.

// AUTOGENERATED. DO NOT EDIT.

// Package {{.PackageName}} is generated by go.chromium.org/luci/tools/cmd/assets.
//
// It contains all {{.Patterns}} files found in the package as byte arrays.
package {{.PackageName}}

// GetAsset returns an asset by its name. Returns nil if no such asset exists.
func GetAsset(name string) []byte {
	return []byte(files[name])
}

// GetAssetString is version of GetAsset that returns string instead of byte
// slice. Returns empty string if no such asset exists.
func GetAssetString(name string) string {
	return files[name]
}

// GetAssetSHA256 returns the asset checksum. Returns nil if no such asset
// exists.
func GetAssetSHA256(name string) []byte {
	data := fileSha256s[name]
	if data == nil {
		return nil
	}
	return append([]byte(nil), data...)
}

// Assets returns a map of all assets.
func Assets() map[string]string {
	cpy := make(map[string]string, len(files))
	for k, v := range files {
		cpy[k] = v
	}
	return cpy
}

var files = map[string]string{
{{range .Assets}}{{.Path | printf "%q"}}: string({{.Body | asByteArray }}),
{{end}}
}

var fileSha256s = map[string][]byte{
{{range .Assets}}{{.Path | printf "%q"}}: {{.SHA256 | asByteArray }},
{{end}}
}
`)))

// assetsTestTmpl is template to assets_test.go file.
var assetsTestTmpl = template.Must(template.New("tmpl").Funcs(funcMap).Parse(strings.TrimSpace(`
// Copyright {{.Year}} The LUCI Authors.
//
// 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.

// AUTOGENERATED. DO NOT EDIT.

// This file is generated by go.chromium.org/luci/tools/cmd/assets.
//
// It contains tests that ensure that assets embedded into the binary are
// identical to files on disk.

package {{.PackageName}}

import (
	"go/build"
	"os"
	"path/filepath"
	"testing"
)

func TestAssets(t *testing.T) {
	t.Parallel()

	pkg, err := build.ImportDir(".", build.FindOnly)
	if err != nil {
		t.Fatalf("can't load package: %s", err)
	}

	fail := false
	for name := range Assets() {
		GetAsset(name) // for code coverage
		path := filepath.Join(pkg.Dir, filepath.FromSlash(name))
		blob, err := os.ReadFile(path)
		if err != nil {
			t.Errorf("can't read file with assets %q (%s) - %s", name, path, err)
			fail = true
		} else if string(blob) != GetAssetString(name) {
			t.Errorf("embedded asset %q is out of date", name)
			fail = true
		}
	}

	if fail {
		t.Fatalf("run 'go generate' to update assets.gen.go")
	}
}
`)))

// templateData is passed to tmpl when rendering it.
type templateData struct {
	Year        int
	Patterns    []string
	PackageName string
	Assets      []asset
}

// asset is single file to be embedded into assets.gen.go.
type asset struct {
	Path string // path relative to package directory
	Body []byte // body of the file
}

func (a asset) SHA256() []byte {
	h := sha256.Sum256(a.Body)
	return h[:]
}

type assetMap map[string]asset

func main() {
	destPkg := ""
	flag.StringVar(&destPkg, "dest-pkg", "",
		`Path to a package to write assets.gen.go to (default is the same as input dir). `+
			`If it's different from the input dir, no *_test.go will be written, since `+
			`it wouldn't know how to discover the original files.`)

	exts := stringlistflag.Flag{}
	flag.Var(&exts, "ext", fmt.Sprintf(
		`(repeatable) Additional extensions to pack up. `+
			`Should be in the form of a glob (e.g. '*.foo'). `+
			`By default this recognizes %q.`, defaultExts.ToSlice()))

	flag.CommandLine.Parse(fixflagpos.Fix(os.Args[1:]))

	var dir string
	switch len(flag.Args()) {
	case 0:
		dir = "."
	case 1:
		dir = flag.Args()[0]
	default:
		fmt.Fprintf(os.Stderr, "usage: assets [dir] [-ext .ext]+\n")
		os.Exit(2)
	}

	if destPkg == "" {
		destPkg = dir
	}

	if err := run(dir, destPkg, exts); err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		os.Exit(1)
	}
}

// run generates assets.gen.go file with all assets discovered in the directory.
func run(inDir, destPkg string, extraExts []string) error {
	exts := defaultExts.Union(stringset.NewFromSlice(extraExts...)).ToSlice()
	sort.Strings(exts)

	assets, err := findAssets(inDir, exts)
	if err != nil {
		return fmt.Errorf("can't find assets in %s - %s", inDir, err)
	}

	pkg, err := build.ImportDir(destPkg, build.ImportComment)
	if err != nil {
		return fmt.Errorf("can't find destination package %q - %s", destPkg, err)
	}

	err = generate(assetsGenGoTmpl, pkg.Name, assets, exts, filepath.Join(pkg.Dir, "assets.gen.go"))
	if err != nil {
		return fmt.Errorf("can't generate assets.gen.go - %s", err)
	}

	if samePaths(inDir, pkg.Dir) {
		err = generate(assetsTestTmpl, pkg.Name, assets, exts, filepath.Join(pkg.Dir, "assets_test.go"))
		if err != nil {
			return fmt.Errorf("can't generate assets_test.go - %s", err)
		}
	}

	return nil
}

// samePaths is true if two paths are identical when converted to absolutes.
//
// Panics if some path can't be converted to absolute.
func samePaths(a, b string) bool {
	var err error
	if a, err = filepath.Abs(a); err != nil {
		panic(err)
	}
	if b, err = filepath.Abs(b); err != nil {
		panic(err)
	}
	return a == b
}

// findAssets recursively scans pkgDir for asset files.
func findAssets(pkgDir string, exts []string) (assetMap, error) {
	assets := assetMap{}

	err := filepath.Walk(pkgDir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() || !isAssetFile(path, exts) {
			return err
		}
		rel, err := filepath.Rel(pkgDir, path)
		if err != nil {
			return err
		}
		blob, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		assets[filepath.ToSlash(rel)] = asset{
			Path: filepath.ToSlash(rel),
			Body: blob,
		}
		return nil
	})

	if err != nil {
		return nil, err
	}
	return assets, nil
}

// isAssetFile returns true if `path` base name matches some of
// `assetExts` glob.
func isAssetFile(path string, assetExts []string) (ok bool) {
	base := filepath.Base(path)
	for _, pattern := range assetExts {
		if match, _ := filepath.Match(pattern, base); match {
			return true
		}
	}
	return false
}

// generate executes the template, runs output through gofmt and dumps it to disk.
func generate(t *template.Template, pkgName string, assets assetMap, assetExts []string, path string) error {
	keys := make([]string, 0, len(assets))
	for k := range assets {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	data := templateData{
		Year:        time.Now().Year(),
		Patterns:    assetExts,
		PackageName: pkgName,
	}
	for _, key := range keys {
		data.Assets = append(data.Assets, assets[key])
	}

	out := bytes.Buffer{}
	if err := t.Execute(&out, data); err != nil {
		return err
	}

	formatted, err := gofmt(out.Bytes())
	if err != nil {
		return fmt.Errorf("can't gofmt %s - %s", path, err)
	}

	return os.WriteFile(path, formatted, 0666)
}

// gofmt applies "gofmt -s" to the content of the buffer.
func gofmt(blob []byte) ([]byte, error) {
	out := bytes.Buffer{}
	cmd := exec.Command("gofmt", "-s")
	cmd.Stdin = bytes.NewReader(blob)
	cmd.Stdout = &out
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	return out.Bytes(), nil
}

func asByteArray(blob []byte) string {
	out := &bytes.Buffer{}
	fmt.Fprintf(out, "[]byte{")
	for i := 0; i < len(blob); i++ {
		fmt.Fprintf(out, "%d, ", blob[i])
		if i%14 == 1 {
			fmt.Fprintln(out)
		}
	}
	fmt.Fprintf(out, "}")
	return out.String()
}