// Copyright 2023 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"

	"github.com/maruel/subcommands"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/logging/gologger"

	"go.fuchsia.dev/infra/cmd/gcs-util/types"
)

// Canonical GCS namespace pattern for versioned blobs.
const blobNamespacePattern = "blobs/[0-9]*/*"

// This command exists to verify delivery blobs per security considerations
// noted in RFC-0207. See fxbug.dev/124944 for more details.
func cmdVerifyBlobs() *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "verify-blobs -blobfs-compression-path <blobfs-compression-path> -manifest-path <manifest-path>",
		ShortDesc: "Verify blobs in an upload manifest.",
		LongDesc:  "Verify blobs in an upload manifest.",
		CommandRun: func() subcommands.CommandRun {
			c := &verifyBlobsCmd{}
			c.Init()
			return c
		},
	}
}

type verifyBlobsCmd struct {
	subcommands.CommandRunBase
	blobfsCompressionPath string
	manifestPath          string
	j                     int
	logLevel              logging.Level
}

func (c *verifyBlobsCmd) Init() {
	c.Flags.StringVar(&c.blobfsCompressionPath, "blobfs-compression-path", "", "Path to blobfs-compression tool.")
	c.Flags.StringVar(&c.manifestPath, "manifest-path", "", "Path to upload manifest.")
	c.Flags.IntVar(&c.j, "j", 32, "Maximum number of concurrent uploading processes.")
	c.Flags.Var(&c.logLevel, "log-level", "Logging level. Can be debug, info, warning, or error.")
}

func (c *verifyBlobsCmd) parseArgs() error {
	if c.blobfsCompressionPath == "" {
		return errors.New("-blobfs-compression-path is required")
	}
	if c.manifestPath == "" {
		return errors.New("-manifest-path is required")
	}
	return nil
}

func (c *verifyBlobsCmd) Run(a subcommands.Application, _ []string, _ subcommands.Env) int {
	if err := c.parseArgs(); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	if err := c.main(); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	return 0
}

// getBlobs returns any Uploads in the input manifest whose destination lies in
// the canonical namespaces for delivery blobs.
func getBlobs(ctx context.Context, manifest []types.Upload) ([]types.Upload, error) {
	var blobs []types.Upload
	for _, upload := range manifest {
		matched, err := filepath.Match(blobNamespacePattern, upload.Destination)
		if err != nil {
			return nil, err
		}
		if matched {
			blobs = append(blobs, upload)
		}
	}
	logging.Debugf(ctx, "got %d blobs", len(blobs))
	return blobs, nil
}

// merkleRoot returns the merkle root of the given blob.
type merkleRoot func(blob types.Upload, tool string) (string, error)

// runBlobfsCompression runs the blobfs-compression tool to get the merkle root
// of the given blob.
func runBlobfsCompression(blob types.Upload, tool string) (string, error) {
	var stdout, stderr bytes.Buffer
	cmd := exec.Command(tool, fmt.Sprintf("--calculate_digest=%s", blob.Source))
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("could not compute merkle root of blob %s: %w\n%s", blob.Source, err, stderr.String())
	}
	return strings.TrimSuffix(stdout.String(), "\n"), nil
}

// verifyBlobs checks that the merkle root of each blob is equal to its
// filename.
func verifyBlobs(ctx context.Context, blobs []types.Upload, tool string, f merkleRoot, j int) error {
	if j <= 0 {
		return fmt.Errorf("concurrency factor j must be a positive number")
	}
	toVerify := make(chan types.Upload, j)
	errs := make(chan error, j)

	queueBlobs := func() {
		defer close(toVerify)
		for _, blob := range blobs {
			toVerify <- blob
		}
	}
	var wg sync.WaitGroup
	wg.Add(j)
	verify := func() {
		defer wg.Done()
		for blob := range toVerify {
			err := verifyBlob(ctx, blob, tool, f)
			if err != nil {
				errs <- err
			}
		}
	}

	go queueBlobs()
	for i := 0; i < j; i++ {
		go verify()
	}
	go func() {
		wg.Wait()
		close(errs)
	}()

	if err := <-errs; err != nil {
		return err
	}
	return nil
}

// verifyBlob checks that the merkle root of the blob is equal to the blob's
// filename.
func verifyBlob(ctx context.Context, blob types.Upload, tool string, f merkleRoot) error {
	logging.Debugf(ctx, "verifying blob %s", blob.Source)
	m, err := f(blob, tool)
	if err != nil {
		return err
	}
	if m != filepath.Base(blob.Source) {
		return fmt.Errorf("blob source %s does not match merkle root %s", blob.Source, m)
	}
	if m != filepath.Base(blob.Destination) {
		return fmt.Errorf("blob destination %s does not match merkle root %s", blob.Destination, m)
	}
	return nil
}

func (c *verifyBlobsCmd) main() error {
	ctx := context.Background()
	ctx = logging.SetLevel(ctx, c.logLevel)
	ctx = gologger.StdConfig.Use(ctx)

	jsonInput, err := os.ReadFile(c.manifestPath)
	if err != nil {
		return err
	}
	var manifest []types.Upload
	if err := json.Unmarshal(jsonInput, &manifest); err != nil {
		return err
	}
	blobs, err := getBlobs(ctx, manifest)
	if err != nil {
		return err
	}
	return verifyBlobs(ctx, blobs, c.blobfsCompressionPath, runBlobfsCompression, c.j)
}