cos-customizer: Support non-prod GCE endpoints
Users can now control the GCE endpoint through a flag that is available
on all subcommands. Every step in the workflow must set the flag for it
to work properly.
BUG=b/176990931
TEST=Run one workflow manually in non-prod GCE, run ./run_tests to
regression check
Change-Id: I0336418a3ad8a680e854085afcb6bbb5f5efce04
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/25821
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Varsha Teratipally <teratipally@google.com>
Reviewed-by: Roy Yang <royyang@google.com>
Tested-by: Robert Kolchmeyer <rkolchmeyer@google.com>
diff --git a/src/cmd/cos_customizer/finish_image_build.go b/src/cmd/cos_customizer/finish_image_build.go
index 6d052e7..1707bac 100644
--- a/src/cmd/cos_customizer/finish_image_build.go
+++ b/src/cmd/cos_customizer/finish_image_build.go
@@ -21,6 +21,7 @@
"log"
"os/exec"
"strconv"
+ "strings"
"time"
"cos.googlesource.com/cos/tools.git/src/pkg/config"
@@ -302,6 +303,10 @@
log.Println(err)
return subcommands.ExitFailure
}
+ // Set non-default GCE endpoints in the buildConfig.
+ if !strings.Contains(svc.BasePath, "compute.googleapis.com/compute/v1") {
+ buildConfig.GCEEndpoint = svc.BasePath
+ }
if err := validateOEM(buildConfig, provConfig); err != nil {
log.Println(err)
return subcommands.ExitFailure
diff --git a/src/cmd/cos_customizer/main.go b/src/cmd/cos_customizer/main.go
index 62e7580..320ca0e 100644
--- a/src/cmd/cos_customizer/main.go
+++ b/src/cmd/cos_customizer/main.go
@@ -35,6 +35,8 @@
var persistentDir = flag.String("local-state-workdir", ".cos-customizer-workdir",
"Name of the directory in $HOME to use for storing local state.")
+var computeEndpoint = flag.String("compute-endpoint", "", "If set, used as the endpoint for the GCE API.")
+
func clients(ctx context.Context, anonymousCreds bool) (*compute.Service, *storage.Client, error) {
var httpClient *http.Client
var err error
@@ -46,7 +48,11 @@
return nil, nil, err
}
}
- svc, err := compute.New(httpClient)
+ computeOpts := []option.ClientOption{option.WithHTTPClient(httpClient)}
+ if *computeEndpoint != "" {
+ computeOpts = append(computeOpts, option.WithEndpoint(*computeEndpoint))
+ }
+ svc, err := compute.NewService(ctx, computeOpts...)
if err != nil {
return nil, nil, err
}
diff --git a/src/pkg/config/config.go b/src/pkg/config/config.go
index 9f84d2e..d680a28 100644
--- a/src/pkg/config/config.go
+++ b/src/pkg/config/config.go
@@ -62,14 +62,15 @@
// Build stores configuration data associated with the image build session.
type Build struct {
- GCSBucket string
- GCSDir string
- Project string
- Zone string
- DiskSize int
- GPUType string
- Timeout string
- GCSFiles []string
+ GCSBucket string
+ GCSDir string
+ Project string
+ Zone string
+ DiskSize int
+ GPUType string
+ Timeout string
+ GCSFiles []string
+ GCEEndpoint string
}
// SaveConfigToFile clears the target config file and then saves the new config
diff --git a/src/pkg/preloader/preload.go b/src/pkg/preloader/preload.go
index 860cfb6..e5aa5ab 100644
--- a/src/pkg/preloader/preload.go
+++ b/src/pkg/preloader/preload.go
@@ -23,6 +23,7 @@
"fmt"
"io/ioutil"
"log"
+ "net/url"
"os"
"os/exec"
"path"
@@ -270,6 +271,17 @@
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,
