Allow client to set the snapshot id.

- ID values are constrained by the following regex "^([a-z0-9\\-]+)$".

Author: ser-io

e2etests/go.mod

@@ -65,4 +65,7 @@ require (
 )
 
 replace github.com/google/android-cuttlefish/frontend/src/libhoclient => ../frontend/src/libhoclient
+
 replace github.com/google/android-cuttlefish/frontend/src/liboperator => ../frontend/src/liboperator
+
+replace github.com/google/android-cuttlefish/frontend/src/host_orchestrator => ../frontend/src/host_orchestrator

e2etests/orchestration/snapshot_test/main_test.go

@@ -19,7 +19,6 @@ import (
 	"errors"
 	"fmt"
 	"log"
-	"net/http"
 	"testing"
 
 	"github.com/google/android-cuttlefish/e2etests/orchestration/common"
@@ -31,7 +30,7 @@ import (
 
 func TestSnapshot(t *testing.T) {
 	t.Skip("b/403233780: Skipping snapshot tests due to out of disk issues on standard runners.")
-	
+
 	ctx, err := common.Setup()
 	if err != nil {
 		t.Fatal(err)
@@ -67,11 +66,15 @@ func TestSnapshot(t *testing.T) {
 	if _, err := dh.ExecADBShellCommand(cID, adbBin, cvd.ADBSerial, []string{"touch", tmpFile}); err != nil {
 		t.Fatal(err)
 	}
+	snapshotID := "00000000-0000-0000-0000-000000000000"
 	// Create a snapshot containing the temporary file.
-	createSnapshotRes, err := createSnapshot(ctx.ServiceURL, groupName, cvd.Name)
+	createSnapshotRes, err := srv.CreateSnapshot(groupName, cvd.Name, &hoapi.CreateSnapshotRequest{SnapshotID: snapshotID})
 	if err != nil {
 		t.Fatal(err)
 	}
+	if diff := cmp.Diff(snapshotID, createSnapshotRes.SnapshotID); diff != "" {
+		t.Fatalf("snapshot id mismatch (-want +got):\n%s", diff)
+	}
 	// Powerwash the device removing the temporary file.
 	if err := srv.Powerwash(groupName, cvd.Name); err != nil {
 		t.Fatal(err)
@@ -165,26 +168,6 @@ func createDevice(srv hoclient.HostOrchestratorService, group_name, artifactsDir
 	return res.CVDs[0], nil
 }
 
-// TODO(b/370550070) Remove once this method is added to the client implementation.
-func createSnapshot(srvURL, group, name string) (*hoapi.CreateSnapshotResponse, error) {
-	helper := hoclient.HTTPHelper{
-		Client:       http.DefaultClient,
-		RootEndpoint: srvURL,
-	}
-	op := &hoapi.Operation{}
-	path := fmt.Sprintf("/cvds/%s/%s/snapshots", group, name)
-	rb := helper.NewPostRequest(path, nil)
-	if err := rb.JSONResDo(op); err != nil {
-		return nil, err
-	}
-	srv := hoclient.NewHostOrchestratorService(srvURL)
-	res := &hoapi.CreateSnapshotResponse{}
-	if err := srv.WaitForOperation(op.Name, res); err != nil {
-		return nil, err
-	}
-	return res, nil
-}
-
 func getCVD(srv hoclient.HostOrchestratorService) (*hoapi.CVD, error) {
 	cvds, err := srv.ListCVDs()
 	if err != nil {

frontend/src/host_orchestrator/api/v1/messages.go

@@ -134,6 +134,12 @@ type ListUploadDirectoriesResponse struct {
 	Items []*UploadDirectory `json:"items"`
 }
 
+type CreateSnapshotRequest struct {
+	// [Optional]
+	// Value must match regex: "^([a-z0-9\\-]+)$".
+	SnapshotID string `json:"snapshot_id,omitempty"`
+}
+
 type CreateSnapshotResponse struct {
 	SnapshotID string `json:"snapshot_id"`
 }

frontend/src/host_orchestrator/orchestrator/controller.go

@@ -294,10 +294,16 @@ func newCreateSnapshotHandler(c Config, om OperationManager) *createSnapshotHand
 }
 
 func (h *createSnapshotHandler) Handle(r *http.Request) (interface{}, error) {
+	req := &apiv1.CreateSnapshotRequest{}
+	err := json.NewDecoder(r.Body).Decode(req)
+	if err != nil {
+		return nil, operator.NewBadRequestError("malformed JSON in request", err)
+	}
 	vars := mux.Vars(r)
 	group := vars["group"]
 	name := vars["name"]
 	opts := CreateSnapshotActionOpts{
+		Request:          req,
 		Selector:         cvd.Selector{Group: group, Instance: name},
 		Paths:            h.Config.Paths,
 		OperationManager: h.OM,

frontend/src/host_orchestrator/orchestrator/createsnapshotaction.go

@@ -15,9 +15,11 @@
 package orchestrator
 
 import (
+	"fmt"
 	"log"
 	"os/user"
 	"path/filepath"
+	"regexp"
 
 	apiv1 "github.com/google/android-cuttlefish/frontend/src/host_orchestrator/api/v1"
 	"github.com/google/android-cuttlefish/frontend/src/host_orchestrator/orchestrator/cvd"
@@ -26,6 +28,7 @@ import (
 )
 
 type CreateSnapshotActionOpts struct {
+	Request          *apiv1.CreateSnapshotRequest
 	Selector         cvd.Selector
 	Paths            IMPaths
 	OperationManager OperationManager
@@ -34,6 +37,7 @@ type CreateSnapshotActionOpts struct {
 }
 
 type CreateSnapshotAction struct {
+	req      *apiv1.CreateSnapshotRequest
 	selector cvd.Selector
 	paths    IMPaths
 	om       OperationManager
@@ -42,6 +46,7 @@ type CreateSnapshotAction struct {
 
 func NewCreateSnapshotAction(opts CreateSnapshotActionOpts) *CreateSnapshotAction {
 	return &CreateSnapshotAction{
+		req:      opts.Request,
 		selector: opts.Selector,
 		paths:    opts.Paths,
 		om:       opts.OperationManager,
@@ -50,6 +55,9 @@ func NewCreateSnapshotAction(opts CreateSnapshotActionOpts) *CreateSnapshotActio
 }
 
 func (a *CreateSnapshotAction) Run() (apiv1.Operation, error) {
+	if err := a.validateRequest(); err != nil {
+		return apiv1.Operation{}, err
+	}
 	if err := createDir(a.paths.SnapshotsRootDir); err != nil {
 		return apiv1.Operation{}, err
 	}
@@ -64,12 +72,29 @@ func (a *CreateSnapshotAction) Run() (apiv1.Operation, error) {
 	return op, nil
 }
 
+var snapshotIDRegex = regexp.MustCompile(`^([a-z0-9\-]+)$`)
+
+func (a *CreateSnapshotAction) validateRequest() error {
+	if a.req.SnapshotID != "" {
+		if !snapshotIDRegex.MatchString(a.req.SnapshotID) {
+			return operator.NewBadRequestError(
+				"invalid snapshot id value", fmt.Errorf("%s does not match %s", a.req.SnapshotID, snapshotIDRegex))
+		}
+	}
+	return nil
+}
+
 func (a *CreateSnapshotAction) createSnapshot(op apiv1.Operation) (*apiv1.CreateSnapshotResponse, error) {
 	if err := a.cvdCLI.Suspend(a.selector); err != nil {
 		return nil, operator.NewInternalError("cvd suspend failed", err)
 	}
-	snapshotID := op.Name
+	snapshotID := a.req.SnapshotID
+	if snapshotID == "" {
+		snapshotID = op.Name
+	}
 	dir := filepath.Join(a.paths.SnapshotsRootDir, snapshotID)
+	// snapshot_util_cvd makes sure the directory is not being used before.
+	// https://github.com/google/android-cuttlefish/blob/7fb47855a2c937e4531044437616bd82e11e3b2e/base/cvd/cuttlefish/host/commands/snapshot_util_cvd/main.cc#L97
 	if err := a.cvdCLI.TakeSnapshot(a.selector, dir); err != nil {
 		return nil, operator.NewInternalError("cvd snapshot_take failed", err)
 	}

frontend/src/libhoclient/go.mod

@@ -47,3 +47,5 @@ require (
 )
 
 replace github.com/google/android-cuttlefish/frontend/src/liboperator => ../liboperator
+
+replace github.com/google/android-cuttlefish/frontend/src/host_orchestrator => ../host_orchestrator

frontend/src/libhoclient/host_orchestrator_client.go

@@ -124,7 +124,7 @@ type HostOrchestratorService interface {
 	Start(groupName, instanceName string, req *hoapi.StartCVDRequest) error
 
 	// Create device snapshot.
-	CreateSnapshot(groupName, instanceName string) (*hoapi.CreateSnapshotResponse, error)
+	CreateSnapshot(groupName, instanceName string, req *hoapi.CreateSnapshotRequest) (*hoapi.CreateSnapshotResponse, error)
 }
 
 const (
@@ -429,9 +429,9 @@ func (c *HostOrchestratorServiceImpl) Start(groupName, instanceName string, req
 	return c.doEmptyResponseRequest(rb)
 }
 
-func (c *HostOrchestratorServiceImpl) CreateSnapshot(groupName, instanceName string) (*hoapi.CreateSnapshotResponse, error) {
+func (c *HostOrchestratorServiceImpl) CreateSnapshot(groupName, instanceName string, req *hoapi.CreateSnapshotRequest) (*hoapi.CreateSnapshotResponse, error) {
 	path := fmt.Sprintf("/cvds/%s/%s/snapshots", groupName, instanceName)
-	rb := c.HTTPHelper.NewPostRequest(path, nil)
+	rb := c.HTTPHelper.NewPostRequest(path, req)
 	op := &hoapi.Operation{}
 	if err := rb.JSONResDo(op); err != nil {
 		return nil, err