// Command private-modules maintains exact private Go proxy artifacts from an
// already authenticated local module cache. It never downloads dependencies.
package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
)

const bundleRoot = "third_party/go-proxy"
const privatePrefix = "code.fritzlab.net/"
const archiveBoundary = "module code.fritzlab.net/action/image-build/private-module-artifacts\n\ngo 1.27.0\n"

type artifact struct {
	Path   string `json:"path"`
	SHA256 string `json:"sha256"`
	Size   int64  `json:"size_bytes"`
}
type module struct {
	Path    string     `json:"path"`
	Version string     `json:"version"`
	Files   []artifact `json:"files"`
}
type manifest struct {
	Version     int      `json:"version"`
	GoModSHA256 string   `json:"go_mod_sha256"`
	GoSumSHA256 string   `json:"go_sum_sha256"`
	Modules     []module `json:"modules"`
}

func digest(b []byte) string { s := sha256.Sum256(b); return hex.EncodeToString(s[:]) }
func inputDigest(name string) (string, error) {
	b, e := os.ReadFile(name)
	if e != nil {
		return "", e
	}
	return digest(b), nil
}
func nativeGo(args ...string) ([]byte, error) {
	c := exec.Command("go", args...)
	c.Env = append(os.Environ(), "GOPROXY=off", "GONOPROXY=none", "GOTOOLCHAIN=local", "GOSUMDB=off", "GO111MODULE=on", "GOFLAGS=", "GOWORK=off")
	return c.Output()
}
func generate() error {
	m := manifest{Version: 1, Modules: []module{}}
	var err error
	if m.GoModSHA256, err = inputDigest("go.mod"); err != nil {
		return err
	}
	if m.GoSumSHA256, err = inputDigest("go.sum"); err != nil {
		return err
	}
	cache, err := nativeGo("env", "GOMODCACHE")
	if err != nil {
		return err
	}
	cacheRoot := filepath.Join(strings.TrimSpace(string(cache)), "cache", "download")
	raw, err := nativeGo("list", "-m", "-json", "all")
	if err != nil {
		return fmt.Errorf("module metadata must already be cached: %w", err)
	}
	parent := filepath.Dir(bundleRoot)
	if err = os.MkdirAll(parent, 0755); err != nil {
		return err
	}
	temp, err := os.MkdirTemp(parent, ".go-proxy-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(temp)
	if err := os.WriteFile(filepath.Join(temp, "go.mod"), []byte(archiveBoundary), 0444); err != nil {
		return err
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	var total int64
	for {
		var src struct {
			Path, Version, GoMod string
			Main                 bool
			Replace              json.RawMessage
		}
		if err = decoder.Decode(&src); errors.Is(err, io.EOF) {
			break
		} else if err != nil {
			return err
		}
		if src.Main || !strings.HasPrefix(src.Path, privatePrefix) {
			continue
		}
		if len(src.Replace) != 0 || src.Version == "" || len(m.Modules) >= 64 {
			return errors.New("private module must have exact unreplaced version")
		}
		mod := module{Path: src.Path, Version: src.Version, Files: []artifact{}}
		stem, err := cacheStem(cacheRoot, src.Path, src.Version, src.GoMod)
		if err != nil {
			return err
		}
		for _, ext := range []string{".info", ".mod", ".zip"} {
			name := stem + ext
			source := filepath.Join(cacheRoot, name)
			info, err := os.Lstat(source)
			if err != nil || !info.Mode().IsRegular() || info.Size() > 64<<20 {
				return errors.New("private module artifact absent or invalid; populate authenticated cache first")
			}
			total += info.Size()
			if total > 256<<20 {
				return errors.New("private module bundle too large")
			}
			b, err := os.ReadFile(source)
			if err != nil {
				return err
			}
			target := filepath.Join(temp, name)
			if err = os.MkdirAll(filepath.Dir(target), 0755); err != nil {
				return err
			}
			if err = os.WriteFile(target, b, 0444); err != nil {
				return err
			}
			mod.Files = append(mod.Files, artifact{Path: filepath.ToSlash(name), SHA256: digest(b), Size: int64(len(b))})
		}
		m.Modules = append(m.Modules, mod)
	}
	sort.Slice(m.Modules, func(i, j int) bool { return m.Modules[i].Path < m.Modules[j].Path })
	if len(m.Modules) == 0 {
		return errors.New("no private modules")
	}
	encoded, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		return err
	}
	if err = os.WriteFile(filepath.Join(temp, "manifest.json"), append(encoded, '\n'), 0444); err != nil {
		return err
	}
	// This path is exclusively generated. Refuse to replace an unrelated tree.
	if _, err = os.Stat(bundleRoot); err == nil {
		if _, err = os.Stat(filepath.Join(bundleRoot, "manifest.json")); err != nil {
			return errors.New("refuse to replace unrecognized bundle")
		}
		if err = os.RemoveAll(bundleRoot); err != nil {
			return err
		}
	} else if !os.IsNotExist(err) {
		return err
	}
	return os.Rename(temp, bundleRoot)
}

var modulePath = regexp.MustCompile(`^code[.]fritzlab[.]net/[A-Za-z0-9._~-]+(/[A-Za-z0-9._~-]+)+$`)
var moduleVersion = regexp.MustCompile(`^v[0-9]+[.][0-9]+[.][0-9]+(-[0-9A-Za-z.-]+)?([+]incompatible)?$`)

// Standard Go proxy escaping replaces each ASCII capital with !lowercase.
func moduleStem(path, version string) (string, error) {
	if len(path) > 512 || len(version) > 256 || !modulePath.MatchString(path) || !moduleVersion.MatchString(version) {
		return "", errors.New("invalid private module identity")
	}
	for _, part := range strings.Split(path, "/") {
		if part == "." || part == ".." {
			return "", errors.New("invalid private module path")
		}
	}
	escape := func(value string) string {
		var out strings.Builder
		for _, r := range value {
			if r >= 'A' && r <= 'Z' {
				out.WriteByte('!')
				r += 'a' - 'A'
			}
			out.WriteRune(r)
		}
		return out.String()
	}
	return escape(path) + "/@v/" + escape(version), nil
}

func cacheStem(cacheRoot, path, version, goMod string) (string, error) {
	stem, err := moduleStem(path, version)
	if err != nil {
		return "", err
	}
	// Lazy module graphs omit GoMod for selected but unexpanded modules.
	// Their canonical cache path is still determined by exact identity.
	if goMod != "" {
		rel, err := filepath.Rel(cacheRoot, goMod)
		if err != nil || filepath.ToSlash(rel) != stem+".mod" {
			return "", errors.New("native artifact path differs from module identity")
		}
	}
	return stem, nil
}

func safePath(p string) bool {
	return p != "" && filepath.Clean(p) == p && !filepath.IsAbs(p) && p != ".." && !strings.HasPrefix(p, "../") && !strings.ContainsAny(p, "\\\x00")
}
func check() error {
	raw, err := os.ReadFile(filepath.Join(bundleRoot, "manifest.json"))
	if err != nil || len(raw) > 1<<20 {
		return errors.New("private module manifest missing or oversized")
	}
	var m manifest
	d := json.NewDecoder(bytes.NewReader(raw))
	d.DisallowUnknownFields()
	if err = d.Decode(&m); err != nil || m.Version != 1 || len(m.Modules) == 0 || len(m.Modules) > 64 {
		return errors.New("invalid private module manifest")
	}
	var trailing any
	if err = d.Decode(&trailing); !errors.Is(err, io.EOF) {
		return errors.New("trailing private module manifest data")
	}
	mod, err := inputDigest("go.mod")
	if err != nil {
		return err
	}
	sum, err := inputDigest("go.sum")
	if err != nil {
		return err
	}
	if mod != m.GoModSHA256 || sum != m.GoSumSHA256 {
		return errors.New("private module closure stale; regenerate after go.mod/go.sum changes")
	}
	boundary, err := os.ReadFile(filepath.Join(bundleRoot, "go.mod"))
	if err != nil || string(boundary) != archiveBoundary {
		return errors.New("private module archive boundary missing or modified")
	}
	expected := map[string]bool{"manifest.json": true, "go.mod": true}
	last := ""
	var total int64
	for _, m := range m.Modules {
		if !strings.HasPrefix(m.Path, privatePrefix) || m.Path <= last || m.Version == "" || len(m.Files) != 3 {
			return errors.New("invalid private module closure")
		}
		last = m.Path
		stem, err := moduleStem(m.Path, m.Version)
		if err != nil {
			return err
		}
		for i, f := range m.Files {
			if !safePath(f.Path) || expected[f.Path] || f.Path != stem+[]string{".info", ".mod", ".zip"}[i] || f.Size < 0 || f.Size > 64<<20 {
				return errors.New("invalid module artifact")
			}
			total += f.Size
			if total > 256<<20 {
				return errors.New("private module bundle too large")
			}
			full := filepath.Join(bundleRoot, f.Path)
			info, e := os.Lstat(full)
			if e != nil || !info.Mode().IsRegular() || info.Size() != f.Size {
				return errors.New("private module artifact missing or invalid")
			}
			b, e := os.ReadFile(full)
			if e != nil || digest(b) != f.SHA256 {
				return errors.New("private module artifact digest mismatch")
			}
			expected[f.Path] = true
		}
	}
	err = filepath.WalkDir(bundleRoot, func(path string, entry os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if entry.Type()&os.ModeSymlink != 0 {
			return errors.New("symlink in private module bundle")
		}
		if entry.IsDir() {
			return nil
		}
		rel, e := filepath.Rel(bundleRoot, path)
		if e != nil || !expected[filepath.ToSlash(rel)] {
			return errors.New("unexpected private module artifact")
		}
		return nil
	})
	return err
}
func main() {
	write := flag.Bool("write", false, "regenerate from authenticated local module cache (no network)")
	flag.Parse()
	var err error
	if *write {
		err = generate()
	} else {
		err = check()
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	fmt.Println("private module artifacts verified")
}