fix: upload from local image generates broken command (#98)

While adding support for qcow2 images in #69 I broke support for local
images. Building a shell pipeline through string concatenation is not a
good idea...

The specific issue was fixed and I also moved building the shell
pipeline to a separate function and added unit tests for all cases, so
it should be easier to spot these issues in the future.

Closes #97

hcloudimages/client.go

@@ -331,34 +331,12 @@ func (s *Client) Upload(ctx context.Context, options UploadOptions) (*hcloud.Ima
 
 	// 6. SSH On Server: Download Image, Decompress, Write to Root Disk
 	logger.InfoContext(ctx, "# Step 6: Downloading image and writing to disk")
-	cmd := ""
-	if options.ImageURL != nil {
-		cmd += fmt.Sprintf("wget --no-verbose -O - %q", options.ImageURL.String())
+
+	cmd, err := assembleCommand(options)
+	if err != nil {
+		return nil, err
 	}
 
-	if options.ImageCompression != CompressionNone {
-		switch options.ImageCompression {
-		case CompressionBZ2:
-			cmd += " | bzip2 -cd"
-		case CompressionXZ:
-			cmd += " | xz -cd"
-		default:
-			return nil, fmt.Errorf("unknown compression: %q", options.ImageCompression)
-		}
-	}
-
-	switch options.ImageFormat {
-	case FormatRaw:
-		cmd += " | dd of=/dev/sda bs=4M"
-	case FormatQCOW2:
-		cmd += " > image.qcow2 && qemu-img dd -f qcow2 -O raw if=image.qcow2 of=/dev/sda bs=4M"
-	}
-
-	cmd += " && sync"
-
-	// Make sure that we fail early, ie. if the image url does not work.
-	// the pipefail does not work correctly without wrapping in bash.
-	cmd = fmt.Sprintf("bash -c 'set -euo pipefail && %s'", cmd)
 	logger.DebugContext(ctx, "running download, decompress and write to disk command", "cmd", cmd)
 
 	output, err := sshsession.Run(sshClient, cmd, options.ImageReader)
@@ -522,3 +500,39 @@ func (s *Client) cleanupTempSSHKeys(ctx context.Context, logger *slog.Logger, se
 
 	return nil
 }
+
+func assembleCommand(options UploadOptions) (string, error) {
+	// Make sure that we fail early, ie. if the image url does not work
+	cmd := "set -euo pipefail && "
+
+	if options.ImageURL != nil {
+		cmd += fmt.Sprintf("wget --no-verbose -O - %q | ", options.ImageURL.String())
+	}
+
+	if options.ImageCompression != CompressionNone {
+		switch options.ImageCompression {
+		case CompressionBZ2:
+			cmd += "bzip2 -cd | "
+		case CompressionXZ:
+			cmd += "xz -cd | "
+		default:
+			return "", fmt.Errorf("unknown compression: %q", options.ImageCompression)
+		}
+	}
+
+	switch options.ImageFormat {
+	case FormatRaw:
+		cmd += "dd of=/dev/sda bs=4M"
+	case FormatQCOW2:
+		cmd += "tee image.qcow2 > /dev/null && qemu-img dd -f qcow2 -O raw if=image.qcow2 of=/dev/sda bs=4M"
+	default:
+		return "", fmt.Errorf("unknown format: %q", options.ImageFormat)
+	}
+
+	cmd += " && sync"
+
+	// the pipefail does not work correctly without wrapping in bash.
+	cmd = fmt.Sprintf("bash -c '%s'", cmd)
+
+	return cmd, nil
+}

hcloudimages/client_test.go

@@ -1,33 +1,110 @@
-package hcloudimages_test
+package hcloudimages
 
 import (
-	"context"
-	"fmt"
 	"net/url"
-
-	"github.com/hetznercloud/hcloud-go/v2/hcloud"
-
-	"github.com/apricote/hcloud-upload-image/hcloudimages"
+	"testing"
 )
 
-func ExampleClient_Upload() {
-	client := hcloudimages.NewClient(
-		hcloud.NewClient(hcloud.WithToken("<your token>")),
-	)
-
-	imageURL, err := url.Parse("https://example.com/disk-image.raw.bz2")
+func mustParseURL(s string) *url.URL {
+	u, err := url.Parse(s)
 	if err != nil {
 		panic(err)
 	}
 
-	image, err := client.Upload(context.TODO(), hcloudimages.UploadOptions{
-		ImageURL:         imageURL,
-		ImageCompression: hcloudimages.CompressionBZ2,
-		Architecture:     hcloud.ArchitectureX86,
-	})
-	if err != nil {
-		panic(err)
-	}
-
-	fmt.Printf("Uploaded Image: %d", image.ID)
+	return u
+}
+
+func TestAssembleCommand(t *testing.T) {
+	tests := []struct {
+		name    string
+		options UploadOptions
+		want    string
+		wantErr bool
+	}{
+		{
+			name:    "local raw",
+			options: UploadOptions{},
+			want:    "bash -c 'set -euo pipefail && dd of=/dev/sda bs=4M && sync'",
+		},
+		{
+			name: "remote raw",
+			options: UploadOptions{
+				ImageURL: mustParseURL("https://example.com/image.xz"),
+			},
+			want: "bash -c 'set -euo pipefail && wget --no-verbose -O - \"https://example.com/image.xz\" | dd of=/dev/sda bs=4M && sync'",
+		},
+		{
+			name: "local xz",
+			options: UploadOptions{
+				ImageCompression: CompressionXZ,
+			},
+			want: "bash -c 'set -euo pipefail && xz -cd | dd of=/dev/sda bs=4M && sync'",
+		},
+		{
+			name: "remote xz",
+			options: UploadOptions{
+				ImageURL:         mustParseURL("https://example.com/image.xz"),
+				ImageCompression: CompressionXZ,
+			},
+			want: "bash -c 'set -euo pipefail && wget --no-verbose -O - \"https://example.com/image.xz\" | xz -cd | dd of=/dev/sda bs=4M && sync'",
+		},
+		{
+			name: "local bz2",
+			options: UploadOptions{
+				ImageCompression: CompressionBZ2,
+			},
+			want: "bash -c 'set -euo pipefail && bzip2 -cd | dd of=/dev/sda bs=4M && sync'",
+		},
+		{
+			name: "remote bz2",
+			options: UploadOptions{
+				ImageURL:         mustParseURL("https://example.com/image.bz2"),
+				ImageCompression: CompressionXZ,
+			},
+			want: "bash -c 'set -euo pipefail && wget --no-verbose -O - \"https://example.com/image.bz2\" | xz -cd | dd of=/dev/sda bs=4M && sync'",
+		},
+		{
+			name: "local qcow2",
+			options: UploadOptions{
+				ImageFormat: FormatQCOW2,
+			},
+			want: "bash -c 'set -euo pipefail && tee image.qcow2 > /dev/null && qemu-img dd -f qcow2 -O raw if=image.qcow2 of=/dev/sda bs=4M && sync'",
+		},
+		{
+			name: "remote qcow2",
+			options: UploadOptions{
+				ImageURL:    mustParseURL("https://example.com/image.qcow2"),
+				ImageFormat: FormatQCOW2,
+			},
+			want: "bash -c 'set -euo pipefail && wget --no-verbose -O - \"https://example.com/image.qcow2\" | tee image.qcow2 > /dev/null && qemu-img dd -f qcow2 -O raw if=image.qcow2 of=/dev/sda bs=4M && sync'",
+		},
+
+		{
+			name: "unknown compression",
+			options: UploadOptions{
+				ImageCompression: "noodle",
+			},
+			wantErr: true,
+		},
+
+		{
+			name: "unknown format",
+			options: UploadOptions{
+				ImageFormat: "poodle",
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := assembleCommand(tt.options)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("assembleCommand() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("assembleCommand() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
 }

hcloudimages/doc_test.go

@@ -0,0 +1,33 @@
+package hcloudimages_test
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+
+	"github.com/hetznercloud/hcloud-go/v2/hcloud"
+
+	"github.com/apricote/hcloud-upload-image/hcloudimages"
+)
+
+func ExampleClient_Upload() {
+	client := hcloudimages.NewClient(
+		hcloud.NewClient(hcloud.WithToken("<your token>")),
+	)
+
+	imageURL, err := url.Parse("https://example.com/disk-image.raw.bz2")
+	if err != nil {
+		panic(err)
+	}
+
+	image, err := client.Upload(context.TODO(), hcloudimages.UploadOptions{
+		ImageURL:         imageURL,
+		ImageCompression: hcloudimages.CompressionBZ2,
+		Architecture:     hcloud.ArchitectureX86,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Printf("Uploaded Image: %d", image.ID)
+}