luciexe/host/cas.go

// Copyright 2024 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 host

import (
	"context"
	"fmt"
	"os"
	"path/filepath"

	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/lucictx"
)

// Find CAS Client. It should have been downloaded when bbagent installs
// CIPD packages.
func findCasClient(workDir string, b *bbpb.Build) (string, error) {
	var bbagentUtilityPath string
	for path, purpose := range b.Infra.Buildbucket.Agent.Purposes {
		if purpose == bbpb.BuildInfra_Buildbucket_Agent_PURPOSE_BBAGENT_UTILITY {
			bbagentUtilityPath = path
			break
		}
	}
	if bbagentUtilityPath == "" {
		return "", errors.Reason("Failed to find bbagent utility packages").Err()
	}

	casClient, err := processCmd(filepath.Join(workDir, bbagentUtilityPath), "cas")
	if err != nil {
		return "", err
	}
	return casClient, nil
}

func generateCasCmd(casClient, outputDir string, casRef *bbpb.InputDataRef_CAS) []string {
	return []string{
		casClient,
		"download",
		"-cas-instance",
		casRef.CasInstance,
		"-digest",
		fmt.Sprintf("%s/%d", casRef.Digest.Hash, casRef.Digest.SizeBytes),
		"-dir",
		outputDir,
		"-log-level",
		"info",
	}
}

func execCasCmd(ctx context.Context, args []string) error {
	// Switch to swarming system account to download CAS inputs, it won't
	// work for non-swarming backends.
	sysCtx, err := lucictx.SwitchLocalAccount(ctx, "system")
	if errors.Is(err, lucictx.ErrNoLocalAuthAccount) {
		logging.Infof(ctx, "Failed to find 'system' account; maybe in non-swarming environment?")
		// Continue to execute cas.
		sysCtx = ctx
	} else if err != nil {
		return errors.Annotate(err, "could not switch to 'system' account in LUCI_CONTEXT").Err()
	}

	cmd := execCommandContext(sysCtx, args[0], args[1:]...)
	logging.Infof(ctx, "Running command: %s", cmd.String())
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return errors.Annotate(cmd.Run(), "Failed to run cas download").Err()
}

func downloadCasFiles(ctx context.Context, b *bbpb.Build, workDir string) error {
	var casClient string
	var err error

	inputData := b.GetInfra().GetBuildbucket().GetAgent().GetInput().GetData()
	for dir, ref := range inputData {
		casRef := ref.GetCas()
		if casRef == nil {
			continue
		}

		if casClient == "" {
			if casClient, err = findCasClient(workDir, b); err != nil {
				return errors.Annotate(err, "download cas files").Err()
			}
		}

		outDir := filepath.Join(workDir, dir)
		cmdArgs := generateCasCmd(casClient, outDir, casRef)
		if err = execCasCmd(ctx, cmdArgs); err != nil {
			return errors.Annotate(err, "download cas files").Err()
		}
	}

	// TODO(chanli): populate ResolvedDataRef_CAS as more things use CAS in
	// buildbucket.
	return nil
}