feat(iaas): remote image download

stackit/internal/services/iaas/image/resource.go

@@ -3,10 +3,13 @@ package image
 import (
 	"bufio"
 	"context"
+	"crypto/md5"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strings"
 	"time"
 
@@ -59,6 +62,22 @@ type Model struct {
 	Checksum      types.Object `tfsdk:"checksum"`
 	Labels        types.Map    `tfsdk:"labels"`
 	LocalFilePath types.String `tfsdk:"local_file_path"`
+	ImageFile     types.Object `tfsdk:"image_file"`
+}
+
+type localModel struct {
+	Path                  types.String `tfsdk:"file_path"`
+	DisablePlanValidation types.Bool   `tfsdk:"disable_plan_validation"`
+}
+
+type downloadModel struct {
+	URL       types.String `tfsdk:"url"`
+	CachePath types.String `tfsdk:"cache_path"`
+}
+
+type imageFileModel struct {
+	Local    types.Object `tfsdk:"local"`
+	Download types.Object `tfsdk:"download"`
 }
 
 // Struct corresponding to Model.Config
@@ -225,7 +244,7 @@ func (r *imageResource) Schema(_ context.Context, _ resource.SchemaRequest, resp
 			},
 			"local_file_path": schema.StringAttribute{
 				Description: "The filepath of the raw image file to be uploaded.",
-				Required:    true,
+				Optional:    true,
 				PlanModifiers: []planmodifier.String{
 					stringplanmodifier.RequiresReplace(),
 				},
@@ -407,6 +426,49 @@ func (r *imageResource) Schema(_ context.Context, _ resource.SchemaRequest, resp
 				ElementType: types.StringType,
 				Optional:    true,
 			},
+			"image_file": schema.SingleNestedAttribute{
+				Description: "Representation of an image file.",
+				Computed:    false,
+				Optional:    true,
+				PlanModifiers: []planmodifier.Object{
+					objectplanmodifier.UseStateForUnknown(),
+				},
+				Attributes: map[string]schema.Attribute{
+					"local": schema.SingleNestedAttribute{
+						Description: "Representation of a local image file.",
+						Optional:    true,
+						Attributes: map[string]schema.Attribute{
+							"file_path": schema.StringAttribute{
+								Description: "Path to the local file.",
+								Required:    true,
+								Validators: []validator.String{
+									// Validating that the file exists in the plan is useful to avoid
+									// creating an image resource where the local image upload will fail
+									validate.FileExists(),
+								},
+							},
+							"disable_plan_validation": schema.BoolAttribute{
+								Description: "Wheter to disable plan-time validation.",
+								Optional:    true,
+							},
+						},
+					},
+					"download": schema.SingleNestedAttribute{
+						Description: "Remote file download settings.",
+						Optional:    true,
+						Attributes: map[string]schema.Attribute{
+							"url": schema.StringAttribute{
+								Description: "URL to downlioad the image from.",
+								Required:    true,
+							},
+							"cache_path": schema.StringAttribute{
+								Description: "Local path to cache the downloaded image.",
+								Required:    true,
+							},
+						},
+					},
+				},
+			},
 		},
 	}
 }
@@ -428,6 +490,53 @@ func (r *imageResource) Create(ctx context.Context, req resource.CreateRequest,
 
 	ctx = core.InitProviderContext(ctx)
 
+	if model.ImageFile.IsNull() || model.ImageFile.IsUnknown() {
+		core.LogAndAddError(ctx, &resp.Diagnostics, "Error creating image", "Error in config")
+	}
+
+	var imageFile imageFileModel
+	diags = model.ImageFile.As(ctx, &imageFile, basetypes.ObjectAsOptions{})
+	resp.Diagnostics.Append(diags...)
+	if resp.Diagnostics.HasError() {
+		core.LogAndAddError(ctx, &resp.Diagnostics, "Error creating image", "Error in config")
+	}
+
+	isLocal := imageFile.Download.IsNull() || imageFile.Download.IsUnknown()
+	isDownload := imageFile.Local.IsNull() || imageFile.Local.IsUnknown()
+	if isDownload && isLocal || !isDownload && !isLocal {
+		core.LogAndAddError(ctx, &resp.Diagnostics, "Error creating image", "Error in config")
+		return
+	}
+
+	var filename string
+	var err error
+	var downloadModel downloadModel
+	var localModel localModel
+	if isDownload {
+		diags = imageFile.Download.As(ctx, &downloadModel, basetypes.ObjectAsOptions{})
+		resp.Diagnostics.Append(diags...)
+		if resp.Diagnostics.HasError() {
+			core.LogAndAddError(ctx, &resp.Diagnostics, "Error creating image", "Error in config")
+			return
+		}
+		file, err := downloadImage(ctx, downloadModel.URL.ValueString())
+		if err != nil {
+			core.LogAndAddError(ctx, &resp.Diagnostics, "Error downloading image", fmt.Sprintf("Downloading Image: %v", err))
+			return
+		}
+		defer file.Close()
+		defer os.RemoveAll(filepath.Dir(file.Name()))
+		filename = file.Name()
+	} else {
+		diags = imageFile.Download.As(ctx, &localModel, basetypes.ObjectAsOptions{})
+		resp.Diagnostics.Append(diags...)
+		if resp.Diagnostics.HasError() {
+			core.LogAndAddError(ctx, &resp.Diagnostics, "Error creating image", "Error in config")
+			return
+		}
+		filename = localModel.Path.ValueString()
+	}
+
 	// Generate API request body from model
 	payload, err := toCreatePayload(ctx, &model)
 	if err != nil {
@@ -468,7 +577,7 @@ func (r *imageResource) Create(ctx context.Context, req resource.CreateRequest,
 	}
 
 	// Upload image
-	err = uploadImage(ctx, &resp.Diagnostics, model.LocalFilePath.ValueString(), imageCreateResp.UploadUrl)
+	err = uploadImage(ctx, &resp.Diagnostics, filename, imageCreateResp.UploadUrl)
 	if err != nil {
 		core.LogAndAddError(ctx, &resp.Diagnostics, "Error creating image", fmt.Sprintf("Uploading image: %v", err))
 		return
@@ -875,6 +984,13 @@ func uploadImage(ctx context.Context, diags *diag.Diagnostics, filePath, uploadU
 	if err != nil {
 		return fmt.Errorf("open file: %w", err)
 	}
+
+	defer func() {
+		if closeErr := file.Close(); closeErr != nil {
+			core.LogAndAddError(ctx, diags, "Error closing file", closeErr.Error())
+		}
+	}()
+
 	stat, err := file.Stat()
 	if err != nil {
 		return fmt.Errorf("stat file: %w", err)
@@ -902,6 +1018,76 @@ func uploadImage(ctx context.Context, diags *diag.Diagnostics, filePath, uploadU
 	if resp.StatusCode != http.StatusOK {
 		return fmt.Errorf("upload image: %s", resp.Status)
 	}
-
 	return nil
 }
+
+func downloadImage(ctx context.Context, downloadURL string) (*os.File, error) {
+	if downloadURL == "" {
+		return nil, fmt.Errorf("download URL is empty")
+	}
+
+	md5sum := fmt.Sprintf("%x", md5.Sum([]byte(downloadURL)))
+
+	tmpDir, err := os.MkdirTemp("", "tf-provider-download-*")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create temp dir: %w", err)
+	}
+
+	filename := filepath.Join(tmpDir, md5sum+".img")
+
+	cleanupOnErr := func() {
+		if err := os.RemoveAll(tmpDir); err != nil {
+			tflog.Warn(ctx, "failed to cleanup temp directory", map[string]interface{}{
+				"dir":   tmpDir,
+				"error": err.Error(),
+			})
+		}
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
+	if err != nil {
+		cleanupOnErr()
+		return nil, fmt.Errorf("create download request: %w", err)
+	}
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		cleanupOnErr()
+		return nil, fmt.Errorf("download image: %w", err)
+	}
+
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			tflog.Debug(ctx, "failed to close HTTP response body", map[string]interface{}{
+				"error": err.Error(),
+			})
+		}
+	}()
+
+	if resp.StatusCode != http.StatusOK {
+		cleanupOnErr()
+		return nil, fmt.Errorf("download image unexpected status: %s", resp.Status)
+	}
+
+	file, err := os.Create(filename)
+	if err != nil {
+		cleanupOnErr()
+		return nil, fmt.Errorf("creating file: %w", err)
+	}
+
+	_, err = io.Copy(file, resp.Body)
+	if err != nil {
+		file.Close()
+		cleanupOnErr()
+		return nil, fmt.Errorf("writing to file: %w", err)
+	}
+
+	if _, err := file.Seek(0, 0); err != nil {
+		file.Close()
+		cleanupOnErr()
+		return nil, fmt.Errorf("seeking file: %w", err)
+	}
+
+	return file, nil
+}

stackit/internal/services/iaas/image/resource_test.go

@@ -1,11 +1,13 @@
 package image
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"os"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -405,3 +407,95 @@ func Test_UploadImage(t *testing.T) {
 		})
 	}
 }
+
+func Test_DownloadImage(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/404":
+			w.WriteHeader(http.StatusNotFound)
+		case "/empty":
+			w.WriteHeader(http.StatusOK)
+		case "/large":
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write(bytes.Repeat([]byte("A"), 1024*1024))
+		case "/drop-conn":
+			hj, ok := w.(http.Hijacker)
+			if !ok {
+				http.Error(w, "webserver doesn't support hijacking", http.StatusInternalServerError)
+				return
+			}
+			conn, _, _ := hj.Hijack()
+			_ = conn.Close()
+		default:
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte("dummy content"))
+		}
+	}))
+	t.Cleanup(server.Close)
+
+	tests := []struct {
+		name        string
+		ctx         context.Context
+		downloadURL string
+		wantBytes   []byte
+		wantErr     bool
+	}{{
+		name:        "ok",
+		downloadURL: server.URL,
+		wantBytes:   []byte("dummy content"),
+		wantErr:     false,
+	},
+		{
+			name:        "invalid_url_format",
+			downloadURL: "http://127.0.0.1:0/invalid",
+			wantErr:     true,
+		},
+		{
+			name:        "status_404_not_found",
+			downloadURL: server.URL + "/404",
+			wantErr:     true,
+		},
+		{
+			name:        "empty_body_200_ok",
+			downloadURL: server.URL + "/empty",
+			wantBytes:   []byte(""),
+			wantErr:     false,
+		},
+		{
+			name:        "large_file_stream",
+			downloadURL: server.URL + "/large",
+			wantBytes:   bytes.Repeat([]byte("A"), 1024*1024),
+			wantErr:     false,
+		},
+		{
+			name:        "connection_dropped_mid_stream",
+			downloadURL: server.URL + "/drop-conn",
+			wantErr:     true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			file, err := downloadImage(tt.ctx, tt.downloadURL)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("downloadImage() error = %v, wantErr %v", err, tt.wantErr)
+			}
+
+			if file != nil {
+				t.Cleanup(func() {
+					_ = file.Close()
+					_ = os.Remove(file.Name())
+				})
+
+				gotBytes, err := os.ReadFile(file.Name())
+				if err != nil {
+					t.Fatalf("failed to read downloaded file: %v", err)
+				}
+
+				if !bytes.Equal(gotBytes, tt.wantBytes) {
+					t.Errorf("byte mismatch: got length %d, want length %d", len(gotBytes), len(tt.wantBytes))
+				}
+			}
+		})
+	}
+}