cos-gpu-installer-v2: Add cos package

Change-Id: I77a029ae87affd29961acb722f06c869b72ac70d
diff --git a/src/pkg/cos/artifacts.go b/src/pkg/cos/artifacts.go
new file mode 100644
index 0000000..48b3594
--- /dev/null
+++ b/src/pkg/cos/artifacts.go
@@ -0,0 +1,143 @@
+package cos
+
+import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+
+	log "github.com/golang/glog"
+	"github.com/pkg/errors"
+
+	"pkg/utils"
+)
+
+const (
+	// TODO(mikewu): consider making GCS buckets as flags.
+	cosToolsGCS      = "cos-tools"
+	internalGCS      = "container-vm-image-staging"
+	chromiumOSSDKGCS = "chromiumos-sdk"
+	kernelInfo       = "kernel_info"
+	kernelSrcArchive = "kernel-src.tar.gz"
+	kernelHeaders    = "kernel-headers.tgz"
+	toolchainURL     = "toolchain_url"
+	toolchainArchive = "toolchain.tar.xz"
+	toolchainEnv     = "toolchain_env"
+	crosKernelRepo   = "https://chromium.googlesource.com/chromiumos/third_party/kernel"
+)
+
+// ArtifactsDownloader defines the interface to download COS artifacts.
+type ArtifactsDownloader interface {
+	DownloadKernelSrc(destDir string) error
+	DownloadToolchainEnv(destDir string) error
+	DownloadToolchain(destDir string) error
+	DownloadKernelHeaders(destDir string) error
+	DownloadArtifact(destDir, artifact string) error
+	GetArtifact(artifact string) ([]byte, error)
+}
+
+// GCSDownloader is the struct downloading COS artifacts from GCS bucket.
+type GCSDownloader struct {
+	envReader *EnvReader
+	Internal  bool
+}
+
+// DownloadKernelSrc downloads COS kernel sources to destination directory.
+func (d *GCSDownloader) DownloadKernelSrc(destDir string) error {
+	return d.DownloadArtifact(destDir, kernelSrcArchive)
+}
+
+// DownloadToolchainEnv downloads toolchain compilation environment variables to destination directory.
+func (d *GCSDownloader) DownloadToolchainEnv(destDir string) error {
+	return d.DownloadArtifact(destDir, toolchainEnv)
+}
+
+// DownloadToolchain downloads toolchain package to destination directory.
+func (d *GCSDownloader) DownloadToolchain(destDir string) error {
+	downloadURL, err := d.getToolchainURL()
+	if err != nil {
+		return errors.Wrap(err, "failed to download toolchain")
+	}
+	outputPath := filepath.Join(destDir, toolchainArchive)
+	if err := utils.DownloadContentFromURL(downloadURL, outputPath, toolchainArchive); err != nil {
+		return errors.Wrap(err, "failed to download toolchain")
+	}
+	return nil
+}
+
+// DownloadKernelHeaders downloads COS kernel headers to destination directory.
+func (d *GCSDownloader) DownloadKernelHeaders(destDir string) error {
+	return d.DownloadArtifact(destDir, kernelHeaders)
+}
+
+// GetArtifact gets an artifact from GCS buckets and returns its content.
+func (d *GCSDownloader) GetArtifact(artifactPath string) ([]byte, error) {
+	tmpDir, err := ioutil.TempDir("", "tmp")
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to create temp dir")
+	}
+	defer os.RemoveAll(tmpDir)
+
+	if err = d.DownloadArtifact(tmpDir, artifactPath); err != nil {
+		return nil, errors.Wrapf(err, "failed to download artifact %s", artifactPath)
+	}
+
+	content, err := ioutil.ReadFile(filepath.Join(tmpDir, filepath.Base(artifactPath)))
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to read file %s", filepath.Join(tmpDir, artifactPath))
+	}
+
+	return content, nil
+}
+
+// DownloadArtifact downloads an artifact from GCS buckets, including public bucket and internal bucket.
+// TODO(mikewu): consider allow users to pass in GCS directories in arguments.
+func (d *GCSDownloader) DownloadArtifact(destDir, artifactPath string) error {
+	var err error
+
+	if err = utils.DownloadFromGCS(destDir, cosToolsGCS, d.artifactPublicPath(artifactPath)); err == nil {
+		return nil
+	}
+	log.Errorf("Failed to download %s from public GCS: %v", artifactPath, err)
+
+	if d.Internal {
+		if err = utils.DownloadFromGCS(destDir, internalGCS, d.artifactInternalPath(artifactPath)); err == nil {
+			return nil
+		}
+		log.Errorf("Failed to download %s from internal GCS: %v", artifactPath, err)
+	}
+
+	return errors.Errorf("failed to download %s", artifactPath)
+}
+
+func (d *GCSDownloader) artifactPublicPath(artifactPath string) string {
+	return fmt.Sprintf("%s/%s", d.envReader.BuildNumber(), artifactPath)
+}
+
+func (d *GCSDownloader) artifactInternalPath(artifactPath string) string {
+	return fmt.Sprintf("lakitu-release/R%s-%s/%s", d.envReader.Milestone(), d.envReader.BuildNumber(), artifactPath)
+}
+
+func (d *GCSDownloader) getToolchainURL() (string, error) {
+	// First, check if the toolchain path is available locally
+	tcPath := d.envReader.ToolchainPath()
+	if tcPath != "" {
+		log.V(2).Info("Found toolchain path file locally")
+		return fmt.Sprintf("https://storage.googleapis.com/%s/%s", chromiumOSSDKGCS, tcPath), nil
+	}
+
+	// Next, check if the toolchain path is available in GCS.
+	tmpDir, err := ioutil.TempDir("", "temp")
+	if err != nil {
+		return "", errors.Wrap(err, "failed to create tmp dir")
+	}
+	defer os.RemoveAll(tmpDir)
+	if err := d.DownloadArtifact(tmpDir, toolchainURL); err != nil {
+		return "", err
+	}
+	toolchainURLContent, err := ioutil.ReadFile(filepath.Join(tmpDir, toolchainURL))
+	if err != nil {
+		return "", errors.Wrap(err, "failed to read toolchain URL file")
+	}
+	return string(toolchainURLContent), nil
+}
diff --git a/src/pkg/cos/cos.go b/src/pkg/cos/cos.go
new file mode 100644
index 0000000..549dd57
--- /dev/null
+++ b/src/pkg/cos/cos.go
@@ -0,0 +1,280 @@
+// Package cos provides functionality to read and configure system configs that are specific to COS images.
+package cos
+
+import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"syscall"
+
+	"pkg/utils"
+
+	log "github.com/golang/glog"
+	"github.com/pkg/errors"
+)
+
+const (
+	espPartition = "/dev/sda12"
+	utsFilepath  = "include/generated/utsrelease.h"
+)
+
+var (
+	execCommand = exec.Command
+)
+
+// DisableKernelModuleLocking disables kernel modules signing enforcement and loadpin so that unsigned kernel modules
+// can be loaded to COS kernel.
+func DisableKernelModuleLocking() error {
+	log.Info("Checking if third party kernel modules can be installed")
+
+	mountDir, err := ioutil.TempDir("", "mountdir")
+	if err != nil {
+		return errors.Wrap(err, "failed to create mount dir")
+	}
+
+	if err := syscall.Mount(espPartition, mountDir, "vfat", 0, ""); err != nil {
+		return errors.Wrap(err, "failed to mount path")
+	}
+
+	grubCfgPath := filepath.Join(mountDir, "esp/efi/boot/grub.cfg")
+	grubCfg, err := ioutil.ReadFile(grubCfgPath)
+	if err != nil {
+		return errors.Wrapf(err, "failed to read grub config from %s", grubCfgPath)
+	}
+
+	grubCfgStr := string(grubCfg)
+	needReboot := false
+	for _, kernelOption := range []string{
+		"module.sig_enforce",
+		"loadpin.enforce",
+		"loadpin.enabled",
+	} {
+		if newGrubCfgStr, needRebootOption := disableKernelOptionFromGrubCfg(kernelOption, grubCfgStr); needRebootOption {
+			needReboot = true
+			grubCfgStr = newGrubCfgStr
+		}
+	}
+
+	if needReboot {
+		log.Info("Modifying grub config to disable module locking.")
+		if err := os.Rename(grubCfgPath, grubCfgPath+".orig"); err != nil {
+			return errors.Wrapf(err, "failed to rename file %s", grubCfgPath)
+		}
+		if err := ioutil.WriteFile(grubCfgPath, []byte(grubCfgStr), 0644); err != nil {
+			return errors.Wrapf(err, "failed to write to file %s", grubCfgPath)
+		}
+	} else {
+		log.Info("Module locking has been disabled.")
+	}
+
+	syscall.Sync()
+	if err := syscall.Unmount(mountDir, 0); err != nil {
+		return err
+	}
+
+	if needReboot {
+		log.Warning("Rebooting")
+		if err := syscall.Reboot(syscall.LINUX_REBOOT_CMD_RESTART); err != nil {
+			return errors.Wrap(err, "failed to reboot")
+		}
+	}
+	return nil
+}
+
+// SetCompilationEnv sets compilation environment variables (e.g. CC, CXX) for third-party kernel module compilation.
+// TODO(mikewu): pass environment variables to the *exec.Cmd that runs the installer.
+func SetCompilationEnv(downloader ArtifactsDownloader) error {
+	log.Info("Downloading compilation environment variables")
+
+	compilationEnvs := make(map[string]string)
+
+	if err := downloader.DownloadToolchainEnv(os.TempDir()); err != nil {
+		// Required to support COS builds not having toolchain_env file
+		log.Info("Using default compilation environment variables")
+		compilationEnvs["CC"] = "x86_64-cros-linux-gnu-gcc"
+		compilationEnvs["CXX"] = "x86_64-cros-linux-gnu-g++"
+	} else {
+		if compilationEnvs, err = utils.LoadEnvFromFile(os.TempDir(), toolchainEnv); err != nil {
+			return errors.Wrap(err, "failed to parse toolchain_env file")
+		}
+	}
+
+	log.Info("Setting compilation environment variables")
+	for key, value := range compilationEnvs {
+		log.Infof("%s=%s", key, value)
+		os.Setenv(key, value)
+	}
+	return nil
+}
+
+// InstallCrossToolchain installs COS toolchain to destination directory.
+func InstallCrossToolchain(downloader ArtifactsDownloader, destDir string) error {
+	log.Info("Installing the toolchain")
+
+	if err := os.MkdirAll(destDir, 0755); err != nil {
+		return errors.Wrapf(err, "failed to create dir %s", destDir)
+	}
+	if empty, _ := utils.IsDirEmpty(destDir); !empty {
+		log.Info("Found existing toolchain. Skipping download and installation")
+		return nil
+	}
+
+	if err := downloader.DownloadToolchain(destDir); err != nil {
+		return errors.Wrap(err, "failed to download toolchain")
+	}
+
+	if err := exec.Command("tar", "xf", filepath.Join(destDir, toolchainArchive), "-C", destDir).Run(); err != nil {
+		return errors.Wrap(err, "failed to extract toolchain archive tarball")
+	}
+
+	log.Info("Configuring environment variables for cross-compilation")
+	os.Setenv("PATH", fmt.Sprintf("%s/bin:%s", destDir, os.Getenv("PATH")))
+	os.Setenv("SYSROOT", filepath.Join(destDir, "usr/x86_64-cros-linux-gnu"))
+	return nil
+}
+
+// InstallKernelSrcPkg installs COS kernel source package to destination directory.
+func InstallKernelSrcPkg(downloader ArtifactsDownloader, destDir string) error {
+	log.Info("Installing the kernel source package")
+
+	if err := downloadKernelSrc(downloader, destDir); err != nil {
+		return errors.Wrap(err, "failed to download kernel source")
+	}
+
+	if err := configureKernel(destDir); err != nil {
+		return errors.Wrap(err, "failed to configure kernel source")
+	}
+
+	if err := correctKernelMagicVersionIfNeeded(destDir); err != nil {
+		return errors.Wrap(err, "failed to run correctKernelMagicVersionIfNeeded")
+	}
+
+	return nil
+}
+
+// InstallKernelHeaderPkg installs kernel header package to destination directory.
+func InstallKernelHeaderPkg(downloader ArtifactsDownloader, destDir string) error {
+	log.Info("Installing the kernel header package")
+
+	if err := os.MkdirAll(destDir, 0755); err != nil {
+		return errors.Wrapf(err, "failed to create dir %s", destDir)
+	}
+	if empty, _ := utils.IsDirEmpty(destDir); !empty {
+		return nil
+	}
+
+	log.Info("Kernel headers not found locally, downloading")
+	if err := downloader.DownloadKernelHeaders(destDir); err != nil {
+		return errors.Wrap(err, "failed to download kernel headers")
+	}
+	if err := exec.Command("tar", "xf", filepath.Join(destDir, kernelHeaders), "-C", destDir).Run(); err != nil {
+		return errors.Wrap(err, "failed to extract kernel header tarball")
+	}
+
+	return nil
+}
+
+// ConfigureModuleSymvers copys Module.symvers file from kernel header dir to kernel source dir.
+func ConfigureModuleSymvers(kernelHeaderDir, kernelSrcDir string) error {
+	log.Info("Configuring Module.symvers file")
+	if err := utils.CopyFile(filepath.Join(kernelHeaderDir, "Module.symvers"),
+		filepath.Join(kernelSrcDir, "Module.symvers")); err != nil {
+		return errors.Wrap(err, "failed to copy Module.symvers file")
+	}
+	return nil
+}
+
+func disableKernelOptionFromGrubCfg(kernelOption, grubCfg string) (newGrubCfg string, needReboot bool) {
+	newGrubCfg = grubCfg
+	needReboot = false
+	kernelOptionEnabled := fmt.Sprintf("%v=1", kernelOption)
+	kernelOptionDisabled := fmt.Sprintf("%v=0", kernelOption)
+
+	if strings.Contains(grubCfg, kernelOption) {
+		if strings.Contains(grubCfg, kernelOptionEnabled) {
+			newGrubCfg = strings.ReplaceAll(grubCfg, kernelOptionEnabled, kernelOptionDisabled)
+			needReboot = true
+		}
+	} else {
+		newGrubCfg = strings.ReplaceAll(grubCfg, "cros_efi", fmt.Sprintf("cros_efi %v", kernelOptionDisabled))
+		needReboot = true
+	}
+	return newGrubCfg, needReboot
+}
+
+func downloadKernelSrc(downloader ArtifactsDownloader, destDir string) error {
+	if err := os.MkdirAll(destDir, 0755); err != nil {
+		return errors.Wrapf(err, "failed to create dir %s", destDir)
+	}
+
+	if empty, _ := utils.IsDirEmpty(destDir); !empty {
+		return nil
+	}
+
+	log.Info("Kernel sources not found locally, downloading")
+	if err := downloader.DownloadKernelSrc(destDir); err != nil {
+		return errors.Wrap(err, "failed to download kernel sources")
+	}
+	if err := exec.Command("tar", "xf", filepath.Join(destDir, kernelSrcArchive), "-C", destDir).Run(); err != nil {
+		return errors.Wrap(err, "failed to extract kernel source tarball")
+	}
+
+	return nil
+}
+
+func configureKernel(kernelSrcDir string) error {
+	log.Info("Configuring kernel")
+	// TODO(mikewu): consider getting kernel configs from kernel headers.
+	kConfig, err := exec.Command("zcat", "/proc/config.gz").Output()
+	if err != nil {
+		return errors.Wrap(err, "failed to read kernel config")
+	}
+	if err := ioutil.WriteFile(filepath.Join(kernelSrcDir, ".config"), kConfig, 0644); err != nil {
+		return errors.Wrap(err, "failed to write kernel config file")
+	}
+	cmd := exec.Command("make", "olddefconfig")
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	cmd.Dir = kernelSrcDir
+	if err := cmd.Run(); err != nil {
+		return errors.Wrap(err, "failed to run `make olddefconfig`")
+	}
+
+	cmd = exec.Command("make", "modules_prepare")
+	cmd.Dir = kernelSrcDir
+	if err := cmd.Run(); err != nil {
+		return errors.Wrap(err, "failed to run `make modules_prepare`")
+	}
+
+	// COS doesn't enable module versioning, disable Module.symvers file check.
+	os.Setenv("IGNORE_MISSING_MODULE_SYMVERS", "1")
+	return nil
+}
+
+func correctKernelMagicVersionIfNeeded(kernelSrcDir string) error {
+	// Normally COS kernel release version has a "+" in the end, e.g. "4.19.102+". But
+	// the utsrelease file generated here doesn't have it, e.g. "4.19.102". Thus we need
+	// to correct the utsrelease file to make it match the real COS kernel release version.
+	utsCmd, err := execCommand("uname", "-r").Output()
+	if err != nil {
+		return errors.Wrap(err, "failed to run `uname -r`")
+	}
+	kernelVersionCmd := strings.TrimSpace(string(utsCmd))
+	utsFile, err := ioutil.ReadFile(filepath.Join(kernelSrcDir, utsFilepath))
+	if err != nil {
+		return errors.Wrap(err, "failed to read utsrelease file")
+	}
+
+	kernelVersionFile := strings.Trim(strings.Fields(string(utsFile))[2], `"`)
+	if kernelVersionCmd != kernelVersionFile {
+		newUtsFile := strings.ReplaceAll(string(utsFile), kernelVersionFile, kernelVersionCmd)
+		log.Info("Modifying kernel release version magic string in source files")
+		if err := ioutil.WriteFile(filepath.Join(kernelSrcDir, utsFilepath), []byte(newUtsFile), 0644); err != nil {
+			return errors.Wrap(err, "failed to write to utsrelease file")
+		}
+	}
+	return nil
+}
diff --git a/src/pkg/cos/cos_test.go b/src/pkg/cos/cos_test.go
new file mode 100644
index 0000000..c6eda3e
--- /dev/null
+++ b/src/pkg/cos/cos_test.go
@@ -0,0 +1,307 @@
+package cos
+
+import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"testing"
+
+	"pkg/utils"
+
+	log "github.com/golang/glog"
+)
+
+var (
+	mockCmdStdout     string
+	mockCmdExitStatus = 0
+)
+
+func fakeExecCommand(command string, args ...string) *exec.Cmd {
+	cs := []string{"-test.run=TestHelperProcess", "--", command}
+	cs = append(cs, args...)
+	cmd := exec.Command(os.Args[0], cs...)
+	es := strconv.Itoa(mockCmdExitStatus)
+	cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1",
+		"STDOUT=" + mockCmdStdout,
+		"EXIT_STATUS=" + es}
+	return cmd
+}
+
+// TestHelperProcess is not a real test. It is a helper process for faking exec.Command.
+func TestHelperProcess(t *testing.T) {
+	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
+		return
+	}
+	fmt.Fprintf(os.Stdout, os.Getenv("STDOUT"))
+	es, err := strconv.Atoi(os.Getenv("EXIT_STATUS"))
+	if err != nil {
+		t.Fatalf("Failed to convert EXIT_STATUS to int: %v", err)
+	}
+	os.Exit(es)
+}
+
+func TestCorrectKernelMagicVersionIfNeeded(t *testing.T) {
+	execCommand = fakeExecCommand
+	defer func() {
+		execCommand = exec.Command
+		mockCmdExitStatus = 0
+	}()
+	for _, tc := range []struct {
+		testName              string
+		kernelVersionUname    string
+		utsRelease            string
+		expectedNewUTSRelease string
+	}{
+		{
+			"NeedHack",
+			"4.19.101+",
+			`#define UTS_RELEASE "4.19.100+"`,
+			`#define UTS_RELEASE "4.19.101+"`,
+		},
+		{
+			"NoNeedHack",
+			"4.19.101+",
+			`#define UTS_RELEASE "4.19.101+"`,
+			`#define UTS_RELEASE "4.19.101+"`,
+		},
+	} {
+
+		tmpDir, err := ioutil.TempDir("", "testing")
+		if err != nil {
+			t.Fatalf("Failed to create temp dir: %v", err)
+		}
+		defer os.RemoveAll(tmpDir)
+		utsFile := filepath.Join(tmpDir, utsFilepath)
+		if err := os.MkdirAll(filepath.Dir(utsFile), 0755); err != nil {
+			t.Fatalf("Failed to create dir: %v", err)
+		}
+		if err := ioutil.WriteFile(utsFile, []byte(tc.utsRelease), 0644); err != nil {
+			t.Fatalf("Failed to write to utsfile: %v", err)
+		}
+		mockCmdStdout = tc.kernelVersionUname
+
+		if err := correctKernelMagicVersionIfNeeded(tmpDir); err != nil {
+			t.Fatalf("Failed to run correctKernelMagicVersionIfNeeded: %v", err)
+		}
+
+		gotUTSRelease, err := ioutil.ReadFile(utsFile)
+		if err != nil {
+			t.Fatalf("Failed to read utsfile: %v", err)
+		}
+		if string(gotUTSRelease) != tc.expectedNewUTSRelease {
+			t.Errorf("%s: Unexpected newUtsRelease, want: %s, got: %s", tc.testName, tc.expectedNewUTSRelease, gotUTSRelease)
+		}
+	}
+}
+
+func TestDownloadKernelSrc(t *testing.T) {
+	tmpDir, err := ioutil.TempDir("", "testing")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	downloder := fakeDownloader{}
+	if err := downloadKernelSrc(&downloder, tmpDir); err != nil {
+		t.Fatalf("Failed to run downloadKernelSrc: %v", err)
+	}
+
+	if _, err := os.Stat(filepath.Join(tmpDir, "kernel-source")); err != nil {
+		t.Errorf("Failed to get kernel source file: %v", err)
+	}
+}
+
+func TestInstallKernelHeaderPkg(t *testing.T) {
+	tmpDir, err := ioutil.TempDir("", "testing")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	downloder := fakeDownloader{}
+	if err := InstallKernelHeaderPkg(&downloder, tmpDir); err != nil {
+		t.Fatalf("Failed to run InstallKernelHeaderPkg: %v", err)
+	}
+
+	if _, err := os.Stat(filepath.Join(tmpDir, "kernel-header")); err != nil {
+		t.Errorf("Failed to get kernel headers file: %v", err)
+	}
+}
+
+func TestSetCompilationEnv(t *testing.T) {
+	origEnvs := os.Environ()
+	defer func() {
+		os.Clearenv()
+		for _, env := range origEnvs {
+			log.Info(env)
+			fields := strings.SplitN(env, "=", 2)
+			os.Setenv(fields[0], fields[1])
+		}
+	}()
+	tmpDir, err := ioutil.TempDir("", "testing")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+	os.Setenv("TMPDIR", tmpDir)
+
+	downloder := fakeDownloader{}
+	if err := SetCompilationEnv(&downloder); err != nil {
+		t.Fatalf("Failed to run SetCompilationEnv: %v", err)
+	}
+
+	for _, tc := range []struct {
+		envKey           string
+		expectedEnvValue string
+	}{
+		{"CC", "x86_64-cros-linux-gnu-clang"},
+		{"CXX", "x86_64-cros-linux-gnu-clang++"},
+	} {
+		if os.Getenv(tc.envKey) != tc.expectedEnvValue {
+			t.Errorf("Unexpected env %s value: want: %s, got: %s", tc.envKey, tc.expectedEnvValue, os.Getenv(tc.envKey))
+		}
+	}
+}
+
+func TestInstallCrossToolchain(t *testing.T) {
+	origEnvs := os.Environ()
+	defer func() {
+		os.Clearenv()
+		for _, env := range origEnvs {
+			log.Info(env)
+			fields := strings.SplitN(env, "=", 2)
+			os.Setenv(fields[0], fields[1])
+		}
+	}()
+	tmpDir, err := ioutil.TempDir("", "testing")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+	origPath := os.Getenv("PATH")
+
+	downloder := fakeDownloader{}
+	if err := InstallCrossToolchain(&downloder, tmpDir); err != nil {
+		t.Fatalf("Failed to run InstallCrossToolchain: %v", err)
+	}
+
+	if _, err := os.Stat(filepath.Join(tmpDir, "x86_64-cros-linux-gnu-clang")); err != nil {
+		t.Errorf("Failed to check file in toolchain: %v", err)
+	}
+	for _, tc := range []struct {
+		envKey           string
+		expectedEnvValue string
+	}{
+		{"PATH", tmpDir + "/bin:" + origPath},
+		{"SYSROOT", filepath.Join(tmpDir, "usr/x86_64-cros-linux-gnu")},
+	} {
+		if os.Getenv(tc.envKey) != tc.expectedEnvValue {
+			t.Errorf("Unexpected env %s value: want: %s, got: %s", tc.envKey, tc.expectedEnvValue, os.Getenv(tc.envKey))
+		}
+	}
+}
+
+func TestDisableKernelOptionFromGrubCfg(t *testing.T) {
+	for _, tc := range []struct {
+		testName           string
+		kernelOption       string
+		grubCfg            string
+		expectedNewGrubCfg string
+		expectedNeedReboot bool
+	}{
+		{
+			"LoadPin",
+			"loadpin.enabled",
+
+			`BOOT_IMAGE=/syslinux/vmlinuz.A init=/usr/lib/systemd/systemd boot=local rootwait ro noresume noswap ` +
+				`loglevel=7 noinitrd console=ttyS0 security=apparmor virtio_net.napi_tx=1 ` +
+				`systemd.unified_cgroup_hierarchy=false systemd.legacy_systemd_cgroup_controller=false csm.disabled=1 ` +
+				`dm_verity.error_behavior=3 dm_verity.max_bios=-1 dm_verity.dev_wait=1 i915.modeset=1 cros_efi root=/dev/dm-0 ` +
+				`"dm=1 vroot none ro 1,0 2539520 verity payload=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 ` +
+				`hashtree=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 hashstart=2539520 alg=sha256 ` +
+				`root_hexdigest=0ff80250bd97ad47a65e7cd330ab70bcf5013d7a86817dca59fcac77f0ba1a8f ` +
+				`salt=414038a6ed9b1f528c327aff4eac16ad5ca4a6699d142ae096e90374af907c34`,
+
+			`BOOT_IMAGE=/syslinux/vmlinuz.A init=/usr/lib/systemd/systemd boot=local rootwait ro noresume noswap ` +
+				`loglevel=7 noinitrd console=ttyS0 security=apparmor virtio_net.napi_tx=1 ` +
+				`systemd.unified_cgroup_hierarchy=false systemd.legacy_systemd_cgroup_controller=false csm.disabled=1 ` +
+				`dm_verity.error_behavior=3 dm_verity.max_bios=-1 dm_verity.dev_wait=1 i915.modeset=1 cros_efi loadpin.enabled=0 root=/dev/dm-0 ` +
+				`"dm=1 vroot none ro 1,0 2539520 verity payload=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 ` +
+				`hashtree=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 hashstart=2539520 alg=sha256 ` +
+				`root_hexdigest=0ff80250bd97ad47a65e7cd330ab70bcf5013d7a86817dca59fcac77f0ba1a8f ` +
+				`salt=414038a6ed9b1f528c327aff4eac16ad5ca4a6699d142ae096e90374af907c34`,
+			true,
+		},
+		{
+			"LoadPinEnabled",
+			"loadpin.enabled",
+			"cros_efi loadpin.enabled=1",
+			"cros_efi loadpin.enabled=0",
+			true,
+		},
+		{
+			"LoadPinDisabled",
+			"loadpin.enabled",
+			"cros_efi loadpin.enabled=0",
+			"cros_efi loadpin.enabled=0",
+			false,
+		},
+	} {
+		newGrubCfg, needReboot := disableKernelOptionFromGrubCfg(tc.kernelOption, tc.grubCfg)
+		if newGrubCfg != tc.expectedNewGrubCfg || needReboot != tc.expectedNeedReboot {
+			t.Errorf("%v: Unexpected output:\nexpect grubcfg: %v\ngot grubcfg: %v\nexpect needReboot: %v, got needReboot: %v",
+				tc.testName, tc.expectedNewGrubCfg, newGrubCfg, tc.expectedNeedReboot, needReboot)
+		}
+	}
+}
+
+type fakeDownloader struct {
+}
+
+func (*fakeDownloader) DownloadKernelSrc(destDir string) error {
+	var archive = map[string][]byte{
+		"kernel-source": []byte("foo"),
+	}
+	if err := utils.CreateTarFile(filepath.Join(destDir, kernelSrcArchive), archive); err != nil {
+		return fmt.Errorf("Failed to download kernel source: %v", err)
+	}
+	return nil
+}
+
+func (*fakeDownloader) DownloadToolchainEnv(destDir string) error {
+	toolchainEnvStr := `CC=x86_64-cros-linux-gnu-clang
+CXX=x86_64-cros-linux-gnu-clang++
+`
+	if err := ioutil.WriteFile(filepath.Join(destDir, toolchainEnv), []byte(toolchainEnvStr), 0644); err != nil {
+		return fmt.Errorf("Failed to download toolchain env file: %v", err)
+	}
+	return nil
+}
+
+func (*fakeDownloader) DownloadToolchain(destDir string) error {
+	var archive = map[string][]byte{
+		"x86_64-cros-linux-gnu-clang": []byte("foo"),
+	}
+	if err := utils.CreateTarFile(filepath.Join(destDir, toolchainArchive), archive); err != nil {
+		return fmt.Errorf("Failed to download toolchain archive: %v", err)
+	}
+	return nil
+}
+
+func (*fakeDownloader) DownloadKernelHeaders(destDir string) error {
+	var archive = map[string][]byte{
+		"kernel-header": []byte("bar"),
+	}
+	if err := utils.CreateTarFile(filepath.Join(destDir, kernelHeaders), archive); err != nil {
+		return fmt.Errorf("Failed to download kernel headers: %v", err)
+	}
+	return nil
+}
+
+func (*fakeDownloader) DownloadArtifact(string, string) error { return nil }
+
+func (*fakeDownloader) GetArtifact(string) ([]byte, error) { return nil, nil }
diff --git a/src/pkg/cos/env_reader.go b/src/pkg/cos/env_reader.go
new file mode 100644
index 0000000..340b516
--- /dev/null
+++ b/src/pkg/cos/env_reader.go
@@ -0,0 +1,77 @@
+package cos
+
+import (
+	"io/ioutil"
+	"path/filepath"
+	"syscall"
+
+	"pkg/utils"
+
+	"github.com/pkg/errors"
+)
+
+const (
+	osReleasePath     = "/etc/os-release"
+	toolchainPathFile = "/etc/toolchain-path"
+
+	buildID        = "BUILD_ID"
+	version        = "VERSION"
+	kernelCommitID = "KERNEL_COMMIT_ID"
+)
+
+// EnvReader is to read system configurations of COS.
+// TODO(mikewu): rename EnvReader to a better name.
+type EnvReader struct {
+	osRelease     map[string]string
+	toolchainPath string
+	uname         syscall.Utsname
+}
+
+// NewEnvReader returns an instance of EnvReader.
+func NewEnvReader(hostRootPath string) (reader *EnvReader, err error) {
+	reader = &EnvReader{}
+	reader.osRelease, err = utils.LoadEnvFromFile(hostRootPath, osReleasePath)
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to read OsRelease file from %s", osReleasePath)
+	}
+
+	if toolchainPath, err := ioutil.ReadFile(filepath.Join(hostRootPath, toolchainPathFile)); err == nil {
+		reader.toolchainPath = string(toolchainPath)
+	}
+
+	if err := syscall.Uname(&reader.uname); err != nil {
+		return nil, errors.Wrap(err, "failed to get uname")
+	}
+	return reader, nil
+}
+
+// OsRelease returns configs of /etc/os-release as a map.
+func (c *EnvReader) OsRelease() map[string]string { return c.osRelease }
+
+// BuildNumber returns COS build number.
+func (c *EnvReader) BuildNumber() string { return c.osRelease[buildID] }
+
+// Milestone returns COS milestone.
+func (c *EnvReader) Milestone() string { return c.osRelease[version] }
+
+// KernelCommit returns commit hash of the COS kernel.
+func (c *EnvReader) KernelCommit() string { return c.osRelease[kernelCommitID] }
+
+// ToolchainPath returns the toolchain path of the COS version.
+// It may return an empty string if the COS version doesn't support the feature.
+func (c *EnvReader) ToolchainPath() string { return c.toolchainPath }
+
+// KernelRelease return COS kernel release, i.e. `uname -r`
+func (c *EnvReader) KernelRelease() string { return charsToString(c.uname.Release[:]) }
+
+// charsToString converts a c-style byte array (null-terminated string) to string.
+func charsToString(chars []int8) string {
+	s := make([]byte, 0, len(chars))
+	for _, ch := range chars {
+		if ch == 0 {
+			break
+		}
+		s = append(s, byte(ch))
+	}
+	return string(s)
+}
diff --git a/src/pkg/cos/env_reader_test.go b/src/pkg/cos/env_reader_test.go
new file mode 100644
index 0000000..a6ceaf3
--- /dev/null
+++ b/src/pkg/cos/env_reader_test.go
@@ -0,0 +1,94 @@
+package cos
+
+import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"reflect"
+	"testing"
+)
+
+func TestEnvReader(t *testing.T) {
+	testDir, err := ioutil.TempDir("", "testing")
+	if err != nil {
+		t.Fatalf("Failed to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(testDir)
+
+	osReleaseString := `BUILD_ID=12688.0.0
+NAME="Container-Optimized OS"
+KERNEL_COMMIT_ID=5d8615d1e135275cbfdf9522517a3b198e7199ee
+GOOGLE_CRASH_ID=Lakitu
+VERSION_ID=80
+BUG_REPORT_URL="https://cloud.google.com/container-optimized-os/docs/resources/support-policy#contact_us"
+PRETTY_NAME="Container-Optimized OS from Google"
+VERSION=80
+GOOGLE_METRICS_PRODUCT_ID=26
+HOME_URL="https://cloud.google.com/container-optimized-os/docs"
+ID=cos`
+	if err := createConfigFile(osReleaseString, osReleasePath, testDir); err != nil {
+		t.Fatalf("Failed to create osRelease file: %v", err)
+	}
+	toolchainPathString := `2019/11/x86_64-cros-linux-gnu-2019.11.16.041937.tar.xz`
+	if err := createConfigFile(toolchainPathString, toolchainPathFile, testDir); err != nil {
+		t.Fatalf("Failed to create toolchain path file: %v", err)
+	}
+
+	envReader, err := NewEnvReader(testDir)
+	if err != nil {
+		t.Fatalf("Failed to create EnvReader: %v", err)
+	}
+
+	for _, tc := range []struct {
+		testName string
+		got      interface{}
+		expect   interface{}
+	}{
+		{
+			"OsRelease",
+			envReader.OsRelease(),
+			map[string]string{
+				"BUILD_ID":                  "12688.0.0",
+				"NAME":                      "Container-Optimized OS",
+				"KERNEL_COMMIT_ID":          "5d8615d1e135275cbfdf9522517a3b198e7199ee",
+				"GOOGLE_CRASH_ID":           "Lakitu",
+				"VERSION_ID":                "80",
+				"BUG_REPORT_URL":            "https://cloud.google.com/container-optimized-os/docs/resources/support-policy#contact_us",
+				"PRETTY_NAME":               "Container-Optimized OS from Google",
+				"VERSION":                   "80",
+				"GOOGLE_METRICS_PRODUCT_ID": "26",
+				"HOME_URL":                  "https://cloud.google.com/container-optimized-os/docs",
+				"ID":                        "cos",
+			},
+		},
+		{"BuildNumber", envReader.BuildNumber(), "12688.0.0"},
+		{"Milestone", envReader.Milestone(), "80"},
+		{"Milestone", envReader.KernelCommit(), "5d8615d1e135275cbfdf9522517a3b198e7199ee"},
+		{"ToolchainPath", envReader.ToolchainPath(), "2019/11/x86_64-cros-linux-gnu-2019.11.16.041937.tar.xz"},
+	} {
+		if !reflect.DeepEqual(tc.expect, tc.got) {
+			t.Errorf("Unexpected %s,\nwant: %v\n got: %v", tc.testName, tc.testName, tc.expect)
+		}
+	}
+}
+
+func createConfigFile(configStr, configFileName, testDir string) error {
+	path := filepath.Join(testDir, configFileName)
+	if err := os.MkdirAll(filepath.Dir(path), 0744); err != nil {
+		return fmt.Errorf("Failed to create dir: %v", err)
+	}
+	configFile, err := os.Create(path)
+	if err != nil {
+		return fmt.Errorf("Failed to create file: %v", err)
+	}
+	defer configFile.Close()
+
+	if _, err = configFile.WriteString(configStr); err != nil {
+		return fmt.Errorf("Failed to write to file %s: %v", configFile.Name(), err)
+	}
+	if err = configFile.Close(); err != nil {
+		return fmt.Errorf("Failed to close file %s: %v", configFile.Name(), err)
+	}
+	return nil
+}
diff --git a/src/pkg/cos/extensions.go b/src/pkg/cos/extensions.go
new file mode 100644
index 0000000..4d77ad8
--- /dev/null
+++ b/src/pkg/cos/extensions.go
@@ -0,0 +1,92 @@
+package cos
+
+import (
+	"fmt"
+	"path/filepath"
+	"regexp"
+
+	"pkg/utils"
+
+	log "github.com/golang/glog"
+	"github.com/pkg/errors"
+)
+
+const (
+	// GPUExtension is the name of GPU extension.
+	GPUExtension = "gpu"
+)
+
+// ExtensionsDownloader is the struct downloading COS extensions from GCS bucket.
+type ExtensionsDownloader interface {
+	ListExtensions() ([]string, error)
+	ListExtensionArtifacts(extension string) ([]string, error)
+	DownloadExtensionArtifact(destDir, extension, artifact string) error
+	GetExtensionArtifact(extension, artifact string) ([]byte, error)
+}
+
+// ListExtensions lists all supported extensions.
+func (d *GCSDownloader) ListExtensions() ([]string, error) {
+	var objects []string
+	var err error
+	if objects, err = utils.ListGCSBucket(cosToolsGCS, d.artifactPublicPath("extensions")); err != nil || len(objects) == 0 {
+		log.Errorf("Failed to list extensions from public GCS: %v", err)
+		if d.Internal {
+			if objects, err = utils.ListGCSBucket(internalGCS, d.artifactInternalPath("extensions")); err != nil {
+				log.Errorf("Failed to list extensions from internal GCS: %v", err)
+			}
+		}
+	}
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to list extensions")
+	}
+
+	var extensions []string
+	re := regexp.MustCompile(`extensions/(\w+)$`)
+	for _, object := range objects {
+		if match := re.FindStringSubmatch(object); match != nil {
+			extensions = append(extensions, match[1])
+		}
+	}
+	return extensions, nil
+}
+
+// ListExtensionArtifacts lists all artifacts of a given extension.
+// TODO(mikewu): make this extension specific.
+func (d *GCSDownloader) ListExtensionArtifacts(extension string) ([]string, error) {
+	var objects []string
+	var err error
+	extensionPath := filepath.Join("extensions", extension)
+	if objects, err = utils.ListGCSBucket(cosToolsGCS, d.artifactPublicPath(extensionPath)); err != nil || len(objects) == 0 {
+		log.Errorf("Failed to list extension artifacts from public GCS: %v", err)
+		// TODO(mikewu): use flags to specify GCS directories.
+		if d.Internal {
+			if objects, err = utils.ListGCSBucket(internalGCS, d.artifactInternalPath(extensionPath)); err != nil {
+				log.Errorf("Failed to list extension artifacts from internal GCS: %v", err)
+			}
+		}
+	}
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to list extensions")
+	}
+
+	var artifacts []string
+	re := regexp.MustCompile(fmt.Sprintf(`extensions/%s/(.+)$`, extension))
+	for _, object := range objects {
+		if match := re.FindStringSubmatch(object); match != nil {
+			artifacts = append(artifacts, match[1])
+		}
+	}
+	return artifacts, nil
+}
+
+// DownloadExtensionArtifact downloads an artifact of the given extension.
+func (d *GCSDownloader) DownloadExtensionArtifact(destDir, extension, artifact string) error {
+	artifactPath := filepath.Join("extensions", extension, artifact)
+	return d.DownloadArtifact(destDir, artifactPath)
+}
+
+// GetExtensionArtifact reads the content of an artifact of the given extension.
+func (d *GCSDownloader) GetExtensionArtifact(extension, artifact string) ([]byte, error) {
+	artifactPath := filepath.Join("extensions", extension, artifact)
+	return d.GetArtifact(artifactPath)
+}