| // Copyright 2018 Google LLC |
| // |
| // 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 preloader contains functionality for preloading a COS image from |
| // provided configuration. |
| package preloader |
| |
| import ( |
| "context" |
| _ "embed" |
| "encoding/json" |
| "fmt" |
| "io/ioutil" |
| "log" |
| "net/url" |
| "os" |
| "os/exec" |
| "path" |
| "path/filepath" |
| "strconv" |
| "strings" |
| "text/template" |
| |
| "cos.googlesource.com/cos/tools.git/src/pkg/config" |
| "cos.googlesource.com/cos/tools.git/src/pkg/fs" |
| "cos.googlesource.com/cos/tools.git/src/pkg/provisioner" |
| "cos.googlesource.com/cos/tools.git/src/pkg/utils" |
| |
| "cloud.google.com/go/storage" |
| ) |
| |
| //go:embed cidata.img |
| var ciDataImg []byte |
| |
| // storeInGCS stores the given files in GCS using the given gcsManager. |
| // Files to store are provided in a map where each key is a file on the local |
| // file system and each value is the relative path in GCS at which to store the |
| // corresponding key. The provided relative paths in GCS must be unique. |
| func storeInGCS(ctx context.Context, gcs *gcsManager, files map[string]string) error { |
| gcsRelPaths := make(map[string]bool) |
| for _, gcsRelPath := range files { |
| if gcsRelPaths[gcsRelPath] { |
| return fmt.Errorf("storeInGCS: collision in relative path %q", gcsRelPath) |
| } |
| gcsRelPaths[gcsRelPath] = true |
| } |
| for file, gcsRelPath := range files { |
| r, err := os.Open(file) |
| if err != nil { |
| return fmt.Errorf("error opening %q: %v", file, err) |
| } |
| defer r.Close() |
| if err := gcs.store(ctx, r, gcsRelPath); err != nil { |
| return err |
| } |
| } |
| return nil |
| } |
| |
| func needDiskResize(provConfig *provisioner.Config, buildSpec *config.Build) bool { |
| // We need to resize the disk during provisioning if: |
| // 1. The requested disk size is larger than default, and |
| // 2. Partitions need to be relocated, i.e. we are enlarging the OEM partition |
| // or reclaiming /dev/sda3 |
| return buildSpec.DiskSize > 10 && (provConfig.BootDisk.OEMSize != "" || provConfig.BootDisk.ReclaimSDA3) |
| } |
| |
| // writeDaisyWorkflow templates the given Daisy workflow and writes the result to a temporary file. |
| // The given workflow should be the one at //data/build_image.wf.json. |
| func writeDaisyWorkflow(inputWorkflow string, outputImage *config.Image, buildSpec *config.Build, provConfig *provisioner.Config) (string, error) { |
| tmplContents, err := ioutil.ReadFile(inputWorkflow) |
| if err != nil { |
| return "", err |
| } |
| labelsJSON, err := json.Marshal(outputImage.Labels) |
| if err != nil { |
| return "", err |
| } |
| acceleratorsJSON, err := json.Marshal([]map[string]interface{}{}) |
| if err != nil { |
| return "", err |
| } |
| if buildSpec.GPUType != "" { |
| acceleratorType := fmt.Sprintf("projects/%s/zones/%s/acceleratorTypes/%s", |
| buildSpec.Project, buildSpec.Zone, buildSpec.GPUType) |
| acceleratorsJSON, err = json.Marshal([]map[string]interface{}{ |
| {"acceleratorType": acceleratorType, "acceleratorCount": 1}}) |
| if err != nil { |
| return "", err |
| } |
| } |
| licensesJSON, err := json.Marshal(outputImage.Licenses) |
| if err != nil { |
| return "", err |
| } |
| |
| // template content for the step resize-disk. |
| // If the oem-size is set, or need to reclaim sda3 (with disk-size-gb set), |
| // create the disk with the default size, and then resize the disk. |
| // Otherwise, a place holder is used. The disk is created with provided disk-size-gb or |
| // the default size. And the disk will not be resized. |
| // The place holder is needed because ResizeDisk API requires a larger size than the original disk. |
| var resizeDiskJSON string |
| var waitResizeJSON string |
| if needDiskResize(provConfig, buildSpec) { |
| // actual disk size |
| resizeDiskJSON = fmt.Sprintf(`"ResizeDisks": [{"Name": "boot-disk","SizeGb": "%d"}]`, buildSpec.DiskSize) |
| waitResizeJSON = ` |
| "WaitForInstancesSignal": [ |
| { |
| "Name": "preload-vm", |
| "Interval": "10s", |
| "SerialOutput": { |
| "Port": 3, |
| "SuccessMatch": "waiting for the boot disk size to change", |
| "FailureMatch": "BuildFailed:" |
| } |
| } |
| ]` |
| } else { |
| // placeholder |
| resizeDiskJSON = `"WaitForInstancesSignal": [{"Name": "preload-vm","Interval": "10s","SerialOutput": {"Port": 3,"SuccessMatch": "BuildStatus:"}}]` |
| waitResizeJSON = `"WaitForInstancesSignal": [{"Name": "preload-vm","Interval": "10s","SerialOutput": {"Port": 3,"SuccessMatch": "BuildStatus:"}}]` |
| } |
| tmpl, err := template.New("workflow").Parse(string(tmplContents)) |
| if err != nil { |
| return "", err |
| } |
| w, err := ioutil.TempFile(fs.ScratchDir, "daisy-") |
| if err != nil { |
| return "", err |
| } |
| if err := tmpl.Execute(w, struct { |
| Labels string |
| Accelerators string |
| Licenses string |
| ResizeDisks string |
| WaitResize string |
| }{ |
| string(labelsJSON), |
| string(acceleratorsJSON), |
| string(licensesJSON), |
| resizeDiskJSON, |
| waitResizeJSON, |
| }); err != nil { |
| w.Close() |
| os.Remove(w.Name()) |
| return "", err |
| } |
| if err := w.Close(); err != nil { |
| os.Remove(w.Name()) |
| return "", err |
| } |
| return w.Name(), nil |
| } |
| |
| func writeCIDataImage(files *fs.Files) (path string, err error) { |
| img, err := ioutil.TempFile(fs.ScratchDir, "cidata-") |
| if err != nil { |
| return "", err |
| } |
| _, writeErr := img.Write(ciDataImg) |
| closeErr := img.Close() |
| if writeErr != nil { |
| return "", writeErr |
| } |
| if closeErr != nil { |
| return "", closeErr |
| } |
| if err := utils.RunCommand([]string{"mcopy", "-i", img.Name(), files.ProvConfig, "::/config.json"}, "", nil); err != nil { |
| return "", err |
| } |
| out, err := ioutil.TempFile(fs.ScratchDir, "cidata-tar-") |
| if err != nil { |
| return "", err |
| } |
| if err := out.Close(); err != nil { |
| return "", err |
| } |
| // tar with the "z" option requires a shell to be installed in the container. |
| // To avoid the shell dependency, gzip the tar ourselves. |
| if err := utils.RunCommand([]string{ |
| "tar", |
| "cf", out.Name(), |
| "--transform", fmt.Sprintf("s|%s|disk.raw|g", strings.TrimLeft(img.Name(), "/")), |
| img.Name(), |
| }, "", nil); err != nil { |
| return "", err |
| } |
| if err := fs.GzipFile(out.Name(), out.Name()+".gz"); err != nil { |
| return "", err |
| } |
| return out.Name() + ".gz", err |
| } |
| |
| func updateProvConfig(provConfig *provisioner.Config, buildSpec *config.Build, buildContexts map[string]string, gcs *gcsManager, files *fs.Files) error { |
| if needDiskResize(provConfig, buildSpec) { |
| provConfig.BootDisk.WaitForDiskResize = true |
| } |
| provConfig.BuildContexts = buildContexts |
| for idx := range provConfig.Steps { |
| if provConfig.Steps[idx].Type == "InstallGPU" { |
| var step provisioner.InstallGPUStep |
| if err := json.Unmarshal(provConfig.Steps[idx].Args, &step); err != nil { |
| return err |
| } |
| if step.GCSDepsPrefix != "" { |
| step.GCSDepsPrefix = gcs.managedDirURL() + "/gcs_files" |
| } |
| buf, err := json.Marshal(&step) |
| if err != nil { |
| return err |
| } |
| provConfig.Steps[idx].Args = json.RawMessage(buf) |
| } |
| } |
| buf, err := json.Marshal(provConfig) |
| if err != nil { |
| return err |
| } |
| log.Printf("Using provisioner config: %s", string(buf)) |
| return config.SaveConfigToPath(files.ProvConfig, provConfig) |
| } |
| |
| func sanitize(output *config.Image) { |
| var licenses []string |
| for _, l := range output.Licenses { |
| if l != "" { |
| licenses = append(licenses, strings.TrimPrefix(l, "https://www.googleapis.com/compute/v1/")) |
| } |
| } |
| output.Licenses = licenses |
| } |
| |
| // daisyArgs computes the parameters to the cos-customizer Daisy workflow (//data/build_image.wf.json) |
| // and uploads dependencies to GCS. |
| func daisyArgs(ctx context.Context, gcs *gcsManager, files *fs.Files, input *config.Image, output *config.Image, buildSpec *config.Build, provConfig *provisioner.Config) ([]string, error) { |
| sanitize(output) |
| buildContexts := map[string]string{ |
| "user": gcs.managedDirURL() + "/" + filepath.Base(files.UserBuildContextArchive), |
| } |
| toUpload := map[string]string{ |
| files.UserBuildContextArchive: filepath.Base(files.UserBuildContextArchive), |
| } |
| for _, gcsFile := range buildSpec.GCSFiles { |
| toUpload[gcsFile] = path.Join("gcs_files", filepath.Base(gcsFile)) |
| } |
| if err := storeInGCS(ctx, gcs, toUpload); err != nil { |
| return nil, err |
| } |
| daisyWorkflow, err := writeDaisyWorkflow(files.DaisyWorkflow, output, buildSpec, provConfig) |
| if err != nil { |
| return nil, err |
| } |
| if err := updateProvConfig(provConfig, buildSpec, buildContexts, gcs, files); err != nil { |
| return nil, err |
| } |
| ciDataFile, err := writeCIDataImage(files) |
| if err != nil { |
| return nil, err |
| } |
| var args []string |
| if buildSpec.GCEEndpoint != "" { |
| // Adding the "projects/" suffix is a Daisy-specific quirk. |
| u, err := url.Parse(buildSpec.GCEEndpoint) |
| if err != nil { |
| return nil, fmt.Errorf("error parsing %q: %v", buildSpec.GCEEndpoint, err) |
| } |
| u.Path = path.Join(u.Path, "projects") |
| // We add a trailing "/" here because Daisy needs it, and u.String() trims |
| // it. |
| args = append(args, "-compute_endpoint_override="+u.String()+"/") |
| } |
| if provConfig.BootDisk.OEMSize == "" && buildSpec.DiskSize > 10 && !provConfig.BootDisk.ReclaimSDA3 { |
| // If the oem-size is set, or need to reclaim sda3, |
| // create the disk with default size, |
| // and then resize the disk in the template step "resize-disk". |
| // Otherwise, create the disk with the provided disk-size-gb. |
| args = append(args, "-var:disk_size_gb", strconv.Itoa(buildSpec.DiskSize)) |
| } |
| if output.Family != "" { |
| args = append(args, "-var:output_image_family", output.Family) |
| } |
| hostMaintenance := "MIGRATE" |
| if buildSpec.GPUType != "" { |
| hostMaintenance = "TERMINATE" |
| } |
| args = append( |
| args, |
| "-var:source_image", |
| input.URL(), |
| "-var:output_image_name", |
| output.Name, |
| "-var:output_image_project", |
| output.Project, |
| "-var:cidata_img", |
| ciDataFile, |
| "-var:machine_type", |
| buildSpec.MachineType, |
| "-var:disk_type", |
| buildSpec.DiskType, |
| "-var:host_maintenance", |
| hostMaintenance, |
| "-var:network", |
| buildSpec.Network, |
| "-var:subnet", |
| buildSpec.Subnet, |
| "-gcs_path", |
| gcs.managedDirURL(), |
| "-project", |
| buildSpec.Project, |
| "-zone", |
| buildSpec.Zone, |
| "-default_timeout", |
| buildSpec.Timeout, |
| "-disable_gcs_logging", |
| daisyWorkflow, |
| ) |
| return args, nil |
| } |
| |
| // BuildImage builds a customized image using Daisy. |
| func BuildImage(ctx context.Context, gcsClient *storage.Client, files *fs.Files, input, output *config.Image, |
| buildSpec *config.Build, provConfig *provisioner.Config) error { |
| gcs := &gcsManager{gcsClient, buildSpec.GCSBucket, buildSpec.GCSDir} |
| defer gcs.cleanup(ctx) |
| args, err := daisyArgs(ctx, gcs, files, input, output, buildSpec, provConfig) |
| if err != nil { |
| return err |
| } |
| cmd := exec.Command(files.DaisyBin, args...) |
| cmd.Stdout = os.Stdout |
| cmd.Stderr = os.Stdout |
| return cmd.Run() |
| } |
