Merge pull request #40 from crytic/fix-droplet-destruction

Fix droplet destruction

CONTRIBUTING.md

@@ -21,19 +21,16 @@ On MacOS you can install trunk with `brew install trunk-io`. If you're using nix
 
 To lint all files that you've modified since the last commit, run `trunk check`. To fix any linter issues trunk can autoformat, run `trunk fmt`. To lint or format all files in the repo, specify -a like `trunk check -a`.
 
-To rescan the repo for any new linters that should be configured, like when adding a new language, run `trunk upgrade`. This will also update linters to their latest version.
-
 ### Overview
 
 When introducing changes to the project, note the following requirements:
 
 - All changes to the main branch should be introduced via pull requests.
-- All branches created for pull requests should follow the `dev/*` naming convention, e.g. `dev/coverage-reports`.
 - Every pull request **must** be reviewed by at least one other peer prior to being merged into the main branch.
 - Code **must** be supported on Linux, macOS, and Windows.
 - Code **must** be sufficiently commented:
   - Every type, function, const, and other variables should be accompanied by [doc comments](https://tip.golang.org/doc/comment).
-  - Inline comments should be provided for every block of code. These should explain what the block of code is aiming to achieve in plain english.
+  - Inline comments should be provided for every block of code. These should explain what the block of code is aiming to achieve in plain English.
   - Inline comments should specify any non-trivial caveats with a given piece of code, considerations to make when maintaining it, etc.
   - Any considerations regarding potential weaknesses or future improvements to be made within your code should be accompanied by an inline comment prefixed with `// TODO: `
   - Comments should provide some value to a new contributor and improve code readability.

Makefile

@@ -24,17 +24,3 @@ lint:
 	go fmt pkg/ssh/*.go
 	go fmt pkg/state/*.go
 	shellcheck cmd/cloudexec/user_data.sh.tmpl
-
-# smuggles git info in the VERSION file to work around nix flake hermicity
-nix-install:
-	@set -e; \
-	cp -f ./VERSION ./.VERSION.backup; \
-	trap 'mv -f ./.VERSION.backup ./VERSION' EXIT; \
-	echo "commit=$(GIT_COMMIT)" >> ./VERSION; \
-	echo "date=$(GIT_DATE)" >> ./VERSION; \
-	echo "nix build"; \
-	nix build; \
-	echo nix profile remove $(shell nix profile list | grep cloudexec | cut -d " " -f 1); \
-	nix profile remove $(shell nix profile list | grep cloudexec | cut -d " " -f 1); \
-	echo nix profile install ./result; \
-	nix profile install ./result

cmd/cloudexec/clean.go

@@ -9,18 +9,19 @@ import (
 	"github.com/crytic/cloudexec/pkg/s3"
 )
 
-func ConfirmDeleteDroplets(config config.Config, dropletName string, instanceToJobs map[int64][]int64) error {
+func ConfirmDeleteDroplets(config config.Config, dropletName string, instanceToJobs map[int64][]int64) ([]int64, error) {
+	var confirmedToDelete []int64
 	instances, err := do.GetDropletsByName(config, dropletName)
 	if err != nil {
-		return fmt.Errorf("Failed to get droplets by name: %w", err)
+		return confirmedToDelete, fmt.Errorf("Failed to get droplets by name: %w", err)
+	}
+	if instanceToJobs == nil {
+		return confirmedToDelete, fmt.Errorf("Given instanceToJobs argument must not be nil")
 	}
 	if len(instances) > 0 {
 		fmt.Printf("Existing %s instance(s) found:\n", dropletName)
 		for _, instance := range instances {
 			// get a pretty string describing the jobs associated with this instance
-			if instanceToJobs == nil {
-				return fmt.Errorf("Given instanceToJobs argument must not be nil")
-			}
 			jobs := instanceToJobs[int64(instance.ID)]
 			var prettyJobs string
 			if len(jobs) == 0 {
@@ -41,14 +42,15 @@ func ConfirmDeleteDroplets(config config.Config, dropletName string, instanceToJ
 				fmt.Printf("Destroying droplet %v...\n", instance.ID)
 				err = do.DeleteDroplet(config, instance.ID)
 				if err != nil {
-					return fmt.Errorf("Failed to destroy droplet: %w", err)
+					return confirmedToDelete, fmt.Errorf("Failed to destroy droplet: %w", err)
 				}
+				confirmedToDelete = append(confirmedToDelete, instance.ID)
 			}
 		}
 	} else {
 		fmt.Printf("Zero %s instances found\n", dropletName)
 	}
-	return nil
+	return confirmedToDelete, nil
 }
 
 func ResetBucket(config config.Config, bucketName string, spacesAccessKey string, spacesSecretKey string, spacesRegion string) error {

cmd/cloudexec/main.go

@@ -210,13 +210,13 @@ func main() {
 						return err
 					}
 
-					err = ConfirmDeleteDroplets(config, dropletName, instanceToJobs)
+					// deletes droplets per user feedback & returns a list of job IDs for state updates
+					confirmedToDelete, err := ConfirmDeleteDroplets(config, dropletName, instanceToJobs)
 					if err != nil {
 						return err
 					}
-					err = ssh.DeleteSSHConfig(user, "cloudexec")
-					if err != nil {
-						return err
+					if len(confirmedToDelete) == 0 {
+						return nil
 					}
 
 					existingState, err := state.GetState(config, bucketName)
@@ -225,7 +225,12 @@ func main() {
 					}
 
 					// mark any running jobs as cancelled
-					err = existingState.CancelRunningJobs(config, bucketName)
+					err = existingState.CancelRunningJobs(config, bucketName, confirmedToDelete)
+					if err != nil {
+						return err
+					}
+
+					err = ssh.DeleteSSHConfig(user, "cloudexec")
 					if err != nil {
 						return err
 					}
@@ -259,13 +264,15 @@ func main() {
 					if err != nil {
 						return err
 					}
-					err = ConfirmDeleteDroplets(config, dropletName, instanceToJobs)
+					confirmedToDelete, err := ConfirmDeleteDroplets(config, dropletName, instanceToJobs)
 					if err != nil {
 						return err
 					}
-					err = ssh.DeleteSSHConfig(user, "cloudexec")
-					if err != nil {
-						return err
+					if len(confirmedToDelete) > 0 {
+						err = ssh.DeleteSSHConfig(user, "cloudexec")
+						if err != nil {
+							return err
+						}
 					}
 					return nil
 				},
@@ -480,6 +487,7 @@ func main() {
 					if err != nil {
 						return err
 					}
+					// If we got a job Id, get that job's state, else continue
 					latestJob := existingState.GetLatestJob()
 					jobStatus := latestJob.Status
 

flake.nix

@@ -32,7 +32,7 @@
             pname = "cloudexec";
             version = "${version}";
             src = ./.;
-            vendorSha256 = "sha256-xiiMcjo+hRllttjYXB3F2Ms2gX43r7/qgwxr4THNhsk=";
+            vendorHash = "sha256-xiiMcjo+hRllttjYXB3F2Ms2gX43r7/qgwxr4THNhsk=";
             nativeBuildInputs = [
               pkgs.git
               pkgs.go_1_20

justfile

@@ -0,0 +1,23 @@
+
+fmt:
+	go fmt cmd/cloudexec/*.go
+	go fmt pkg/digitalocean/*.go
+	go fmt pkg/ssh/*.go
+	go fmt pkg/state/*.go
+
+trunk:
+	trunk fmt
+	trunk check
+
+build:
+  nix build
+
+install:
+	nix build
+	echo nix profile remove $(nix profile list | grep cloudexec | cut -d " " -f 1)
+	nix profile remove $(nix profile list | grep cloudexec | cut -d " " -f 1)
+	echo nix profile install ./result
+	nix profile install ./result
+
+launch-example:
+  cd example && cloudexec launch

pkg/ssh/ssh.go

@@ -40,7 +40,8 @@ type HostConfig struct {
 }
 
 func EnsureSSHIncludeConfig(usr *user.User) error {
-	includeString := "# Added by cloudexec\nInclude config.d/*\n\n"
+	commentString := "# Added by cloudexec\n"
+	includeString := "Include config.d/*\n"
 	sshDir := filepath.Join(usr.HomeDir, ".ssh")
 	configPath := filepath.Join(sshDir, "config")
 
@@ -54,8 +55,8 @@ func EnsureSSHIncludeConfig(usr *user.User) error {
 
 	// Check if the config file exists
 	if _, err = os.Stat(configPath); os.IsNotExist(err) {
-		// If the config file does not exist, create it with the "Include config.d/*" line
-		configFileContent = includeString
+		// If the config file does not exist, create it with the includeString line
+		configFileContent = commentString + includeString
 	} else {
 		// If the config file exists, read its content
 		content, err := os.ReadFile(configPath)
@@ -65,10 +66,10 @@ func EnsureSSHIncludeConfig(usr *user.User) error {
 
 		configFileContent = string(content)
 
-		// Check if the "Include config.d/*" line is present
-		if !strings.Contains(configFileContent, "Include config.d/*") {
+		// Check if the includeString line is present
+		if !strings.Contains(configFileContent, includeString) {
 			// If not present, add the line to the top of the file
-			configFileContent = includeString + configFileContent
+			configFileContent = commentString + includeString + configFileContent
 		} else {
 			// If the line is already present, no further action is required
 			return nil

pkg/state/state.go

@@ -178,11 +178,20 @@ func (s *State) DeleteJob(jobID int64) {
 	}
 }
 
-func (s *State) CancelRunningJobs(config config.Config, bucketName string) error {
+func (s *State) CancelRunningJobs(config config.Config, bucketName string, toCancelSlice []int64) error {
+	// Gather list of ids to mark as cancelled
+	toCancel := make(map[int64]bool)
+	for _, id := range toCancelSlice {
+		toCancel[id] = true
+	}
+
 	// Mark any running jobs as cancelled
 	for i, job := range s.Jobs {
 		if job.Status == Running || job.Status == Provisioning {
-			s.Jobs[i].Status = Cancelled
+			if _, exists := toCancel[job.ID]; exists {
+				fmt.Printf("Setting status of job %d to 'Cancelled'\n", job.ID)
+				s.Jobs[i].Status = Cancelled
+			}
 		}
 	}