diff --git a/hcloudimages/client.go b/hcloudimages/client.go index c30cef9..fa51b55 100644 --- a/hcloudimages/client.go +++ b/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 +} diff --git a/hcloudimages/client_test.go b/hcloudimages/client_test.go index 4905bdf..86f8961 100644 --- a/hcloudimages/client_test.go +++ b/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("")), - ) - - 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) + } + }) + } } diff --git a/hcloudimages/doc_test.go b/hcloudimages/doc_test.go new file mode 100644 index 0000000..4905bdf --- /dev/null +++ b/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("")), + ) + + 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) +}