nmon/internal/vmsup/binary.go

package vmsup

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"

	"github.com/v10networks/nmon/internal/config"
)

// Environment overrides for development and tests built without vmembed.
const (
	envBinary     = "NMON_VM_BINARY"
	envUnverified = "NMON_VM_BINARY_UNVERIFIED" // "1": skip the pin check (tests only)
)

// maxBinarySize bounds decompression so a corrupt .gz cannot fill the disk.
const maxBinarySize = 512 << 20

var (
	// ErrChecksum means the store binary does not match its pinned SHA-256.
	// nmon refuses to start in that case.
	ErrChecksum = errors.New("vmsup: store binary failed its checksum")
	// ErrNoBinary means a store must be started but this build has none.
	ErrNoBinary = errors.New("vmsup: no embedded store binary in this build; set --vm-hot-url and --vm-long-url")
)

// binarySpec says which file to run and which hash it must have ("" = unverified).
type binarySpec struct {
	path string
	want string
}

// verify hashes the file at s.path against s.want. It runs before every start.
func (s binarySpec) verify() error {
	if s.want == "" {
		return nil
	}
	got, err := fileSHA256(s.path)
	if err != nil {
		return fmt.Errorf("vmsup: hash %s: %w", s.path, err)
	}
	if got != s.want {
		return fmt.Errorf("%w: %s has sha256 %s, want %s", ErrChecksum, s.path, got, s.want)
	}
	return nil
}

// prepareBinary decides which binary to run, extracting the embedded one into
// <data-dir>/bin when needed, and verifies it.
func prepareBinary(dataDir string) (binarySpec, error) {
	if p := os.Getenv(envBinary); p != "" {
		return prepareSuppliedBinary(p)
	}
	if len(embeddedGz) == 0 {
		return binarySpec{}, ErrNoBinary
	}
	dst := filepath.Join(config.BinDir(dataDir), "victoria-metrics-prod")
	spec := binarySpec{path: dst, want: embeddedBinSHA256}
	if err := extract(embeddedGz, spec); err != nil {
		return spec, err
	}
	return spec, nil
}

// prepareSuppliedBinary verifies a binary named by NMON_VM_BINARY against the
// pinned hash, unless the test-only override disables the check.
func prepareSuppliedBinary(path string) (binarySpec, error) {
	spec := binarySpec{path: path}
	if os.Getenv(envUnverified) != "1" {
		spec.want = binPin(runtime.GOARCH)
		if runtime.GOOS != "linux" || spec.want == "" {
			return spec, fmt.Errorf("%w: no pinned hash for %s/%s; set %s=1 for tests",
				ErrChecksum, runtime.GOOS, runtime.GOARCH, envUnverified)
		}
	}
	if err := spec.verify(); err != nil {
		return spec, err
	}
	return spec, nil
}

// extract writes the gunzipped data to spec.path (mode 0755) unless a file with
// the right hash is already there. It writes a temp file in the same directory,
// verifies it, and renames it into place, so a bad or partial file is never
// left at the final path.
func extract(gz []byte, spec binarySpec) error {
	if _, err := os.Stat(spec.path); err == nil {
		if spec.verify() == nil {
			return nil
		}
	}
	dir := filepath.Dir(spec.path)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return fmt.Errorf("vmsup: create %s: %w", dir, err)
	}
	tmp, err := os.CreateTemp(dir, ".victoria-metrics-prod-*")
	if err != nil {
		return fmt.Errorf("vmsup: create temp binary: %w", err)
	}
	tmpName := tmp.Name()
	defer func() { _ = os.Remove(tmpName) }() // no-op after a successful rename

	zr, err := gzip.NewReader(bytes.NewReader(gz))
	if err != nil {
		_ = tmp.Close()
		return fmt.Errorf("%w: embedded archive: %w", ErrChecksum, err)
	}
	h := sha256.New()
	n, err := io.Copy(io.MultiWriter(tmp, h), io.LimitReader(zr, maxBinarySize+1))
	if err == nil && n > maxBinarySize {
		err = errors.New("embedded binary is implausibly large")
	}
	if err != nil {
		_ = tmp.Close()
		return fmt.Errorf("%w: unpack embedded binary: %w", ErrChecksum, err)
	}
	if got := hex.EncodeToString(h.Sum(nil)); got != spec.want {
		_ = tmp.Close()
		return fmt.Errorf("%w: embedded binary has sha256 %s, want %s", ErrChecksum, got, spec.want)
	}
	if err := tmp.Chmod(0o755); err != nil {
		_ = tmp.Close()
		return fmt.Errorf("vmsup: chmod binary: %w", err)
	}
	if err := tmp.Sync(); err != nil {
		_ = tmp.Close()
		return fmt.Errorf("vmsup: sync binary: %w", err)
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("vmsup: write binary: %w", err)
	}
	if err := os.Rename(tmpName, spec.path); err != nil {
		return fmt.Errorf("vmsup: install binary: %w", err)
	}
	return nil
}

func fileSHA256(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer func() { _ = f.Close() }()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}