tools/gcs/gcs.go

// Package gcs provides tools for interacting with GCS buckets and files.
package gcs

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"google.golang.org/api/iterator"
)

const (
	expectedGCSPathFmt     = "gs://<bucket-name>/<directory_path>"
	defaultFilePermissions = 0755
)

// GCSConfig stores the gcs bucket handler and a module gcs directory path.
type GCSConfig struct {
	bucket *storage.BucketHandle
	path   string
}

func NewClient(ctx context.Context) (*storage.Client, error) {
	client, err := storage.NewClient(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to create new GCS client: %v", err)
	}
	return client, err
}

// Init initializes the gcs bucket and module directory path from a
// gsutil uri.
func (c *GCSConfig) Init(ctx context.Context, gsUri string, client *storage.Client) error {
	bucketName, dir, err := getBucketNameAndDir(gsUri)
	if err != nil {
		return fmt.Errorf("failed to retrieve bucket name and module path: %v", err)
	}

	// Creating bucket handle
	bkt := client.Bucket(bucketName)

	// Checking if directory exists and non-empty
	it := bkt.Objects(ctx, &storage.Query{
		Prefix: c.path,
	})
	if _, err = it.Next(); err != nil {
		if err == iterator.Done {
			return fmt.Errorf("failed to validate gcs directory (%s/%s): empty or non-existent directory", bucketName, dir)
		}
		return fmt.Errorf("failed to validate gcs directory (%s/%s) : %v", bucketName, dir, err)
	}

	c.bucket = bkt
	c.path = dir
	return nil
}

// ListModules returns a list of all the modules in a GCS directory.
// If recursive is set, it includes modules from the subdirectories of the gcs path.
// Identifies modules as files with .ko extension.
func (c *GCSConfig) ListModules(ctx context.Context, recursive bool) ([]string, error) {
	// Querying for files in gcs directory.
	var query storage.Query
	if recursive {
		query = storage.Query{
			Prefix: c.path,
		}
	} else {
		query = storage.Query{
			Prefix:    c.path,
			Delimiter: "/",
		}
	}

	it := c.bucket.Objects(ctx, &query)

	var modules []string
	for {
		attrs, err := it.Next()

		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("failed to retrieve modules from the gcs directory (%s/%s): %v", c.bucketName(), c.path, err)
		}
		if filepath.Ext(attrs.Name) == ".ko" {
			moduleName, _ := strings.CutPrefix(attrs.Name, c.path)
			modules = append(modules, moduleName)
		}
	}
	return modules, nil
}

// getModulePath safely joins the localDir and moduleName, ensuring the resulting path is within localDir.
func getModulePath(localDir, moduleName string) (string, error) {
	if moduleName == "" {
		return "", fmt.Errorf("moduleName cannot be empty")
	}
	dst := filepath.Join(localDir, moduleName)

	absLocalDir, err := filepath.Abs(localDir)
	if err != nil {
		return "", fmt.Errorf("failed to get absolute path for local directory (%s): %v", localDir, err)
	}

	absDst, err := filepath.Abs(dst)
	if err != nil {
		return "", fmt.Errorf("failed to get absolute path for destination (%s): %v", dst, err)
	}

	if !strings.HasPrefix(absDst, absLocalDir) {
		return "", fmt.Errorf("invalid module name, path traversal detected: %s", moduleName)
	}

	return absDst, nil
}

// DownloadKernelModule copies a module file from a GCS directory to a local directory.
// It returns the local file path of the downloaded module.
// If a file exists in that directory with the same name as the module, an error
// is returned to the user.
func (c *GCSConfig) DownloadKernelModule(ctx context.Context, moduleName, localDir string) (string, error) {
	dst, err := getModulePath(localDir, moduleName)
	if err != nil {
		return "", fmt.Errorf("failed to get path for module (%s) in directory (%s): %v", moduleName, localDir, err)
	}
	src := filepath.Join(c.path, moduleName)

	// Checking if there exists a file with the same name as module in directory
	// if not, ensuring the path exists.
	if _, err := os.Stat(dst); errors.Is(err, fs.ErrNotExist) {
		if err := os.MkdirAll(filepath.Dir(dst), defaultFilePermissions); err != nil {
			return "", fmt.Errorf("failed to create directory (%s): %v", filepath.Dir(dst), err)
		}
	} else if err == nil {
		return "", fmt.Errorf("failed to download module (%s) to path (%s). File already exists", moduleName, dst)
	} else {
		return "", fmt.Errorf("failed to download module (%s) to path (%s): %v", moduleName, dst, err)
	}

	moduleObj := c.bucket.Object(src)
	rc, err := moduleObj.NewReader(ctx)
	if err != nil {
		return "", fmt.Errorf("failed to read from module (%s) in gcs bucket (%s): %v", src, c.bucketName(), err)
	}
	defer rc.Close()

	// Creating a file with the same name as the module in the local directory
	f, err := os.Create(dst)
	if err != nil {
		return "", fmt.Errorf("failed to create file (%s): %v", dst, err)
	}
	defer f.Close()

	if _, err := io.Copy(f, rc); err != nil {
		return "", fmt.Errorf("failed to copy module (%s) from gcs bucket (%s): %v", src, c.bucketName(), err)
	}
	return dst, nil
}

func getBucketNameAndDir(gcsPath string) (string, string, error) {
	filePath, isValid := strings.CutPrefix(gcsPath, "gs://")
	if !isValid {
		return "", "", fmt.Errorf("invalid syntax for gcs path. Expected gsutil uri format: %s \t got: %v", expectedGCSPathFmt, gcsPath)
	}
	bktName, path, isValid := strings.Cut(filePath, "/")
	if !isValid || strings.TrimSpace(path) == "" {
		return "", "", fmt.Errorf("file path (%s) should have both bucket name and path to module directory", filePath)
	}
	if !strings.HasSuffix(path, "/") {
		path = fmt.Sprintf("%s/", path)
	}
	return bktName, path, nil
}

func (c *GCSConfig) bucketName() string {
	if c.bucket == nil {
		return ""
	}
	return c.bucket.BucketName()
}