From d981cc7f3e155c9c4c604c04986883a03df16af2 Mon Sep 17 00:00:00 2001
From: Dan Petersen <dp@danitso.com>
Date: Thu, 12 Dec 2019 23:15:13 +0100
Subject: [PATCH] Changed file upload code to use storage devices instead of
 memory

---
 proxmox/virtual_environment.go            |  7 ++++
 proxmox/virtual_environment_datastores.go | 48 +++++++++++++++++++----
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/proxmox/virtual_environment.go b/proxmox/virtual_environment.go
index 6701e605..3764c9dd 100644
--- a/proxmox/virtual_environment.go
+++ b/proxmox/virtual_environment.go
@@ -43,6 +43,7 @@ type VirtualEnvironmentClient struct {
 type VirtualEnvironmentMultiPartData struct {
 	Boundary string
 	Reader   io.Reader
+	Size     *int64
 }
 
 // NewVirtualEnvironmentClient creates and initializes a VirtualEnvironmentClient instance.
@@ -85,6 +86,7 @@ func NewVirtualEnvironmentClient(endpoint, username, password string, insecure b
 // DoRequest performs a HTTP request against a JSON API endpoint.
 func (c *VirtualEnvironmentClient) DoRequest(method, path string, requestBody interface{}, responseBody interface{}) error {
 	var reqBodyReader io.Reader
+	var reqContentLength *int64
 
 	log.Printf("[DEBUG] Performing HTTP %s request (path: %s)", method, path)
 
@@ -98,6 +100,7 @@ func (c *VirtualEnvironmentClient) DoRequest(method, path string, requestBody in
 		if multipart {
 			reqBodyReader = multipartData.Reader
 			reqBodyType = fmt.Sprintf("multipart/form-data; boundary=%s", multipartData.Boundary)
+			reqContentLength = multipartData.Size
 
 			log.Printf("[DEBUG] Added multipart request body to HTTP %s request (path: %s)", method, modifiedPath)
 		} else if pipedBody {
@@ -140,6 +143,10 @@ func (c *VirtualEnvironmentClient) DoRequest(method, path string, requestBody in
 
 	req.Header.Add("Accept", "application/json")
 
+	if reqContentLength != nil {
+		req.ContentLength = *reqContentLength
+	}
+
 	if reqBodyType != "" {
 		req.Header.Add("Content-Type", reqBodyType)
 	}
diff --git a/proxmox/virtual_environment_datastores.go b/proxmox/virtual_environment_datastores.go
index 0d3f1b39..6bb9e599 100644
--- a/proxmox/virtual_environment_datastores.go
+++ b/proxmox/virtual_environment_datastores.go
@@ -5,11 +5,12 @@
 package proxmox
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"mime/multipart"
+	"os"
 	"sort"
 )
 
@@ -99,18 +100,51 @@ func (c *VirtualEnvironmentClient) UploadFileToDatastore(d *VirtualEnvironmentDa
 		}
 	}()
 
-	// Due to Proxmox VE not supporting chunked transfers, we sadly need to load the file into memory.
-	// This is not optimal for large files but there's no alternative right now.
-	workaroundReader := new(bytes.Buffer)
-	workaroundReader.ReadFrom(r)
+	// We need to store the multipart content in a temporary file to avoid using high amounts of memory.
+	// This is necessary due to Proxmox VE not supporting chunked transfers in v6.1 and earlier versions.
+	tempMultipartFile, err := ioutil.TempFile("", "multipart")
+
+	if err != nil {
+		return nil, err
+	}
+
+	tempMultipartFileName := tempMultipartFile.Name()
+
+	io.Copy(tempMultipartFile, r)
+
+	err = tempMultipartFile.Close()
+
+	if err != nil {
+		return nil, err
+	}
+
+	defer os.Remove(tempMultipartFileName)
+
+	// Now that the multipart data is stored in a file, we can go ahead and do a HTTP POST request.
+	fileReader, err := os.Open(tempMultipartFileName)
+
+	if err != nil {
+		return nil, err
+	}
+
+	defer fileReader.Close()
+
+	fileInfo, err := fileReader.Stat()
+
+	if err != nil {
+		return nil, err
+	}
+
+	fileSize := fileInfo.Size()
 
 	reqBody := &VirtualEnvironmentMultiPartData{
 		Boundary: m.Boundary(),
-		Reader:   workaroundReader,
+		Reader:   fileReader,
+		Size:     &fileSize,
 	}
 
 	resBody := &VirtualEnvironmentDatastoreUploadResponseBody{}
-	err := c.DoRequest(hmPOST, fmt.Sprintf("nodes/%s/storage/%s/upload", d.NodeName, d.DatastoreID), reqBody, resBody)
+	err = c.DoRequest(hmPOST, fmt.Sprintf("nodes/%s/storage/%s/upload", d.NodeName, d.DatastoreID), reqBody, resBody)
 
 	if err != nil {
 		return nil, err