From 90ff4c18d58025caa784b74e12b9f7f1f1047871 Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Fri, 21 Feb 2025 10:31:59 +0100
Subject: [PATCH] go/runtime/host/tdx: Resize overlay image if needed

---
 .changelog/6077.bugfix.md   |  1 +
 go/runtime/host/tdx/qemu.go | 36 +++++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 .changelog/6077.bugfix.md

diff --git a/.changelog/6077.bugfix.md b/.changelog/6077.bugfix.md
new file mode 100644
index 00000000000..36f2431624e
--- /dev/null
+++ b/.changelog/6077.bugfix.md
@@ -0,0 +1 @@
+go/runtime/host/tdx: Resize overlay image if needed
diff --git a/go/runtime/host/tdx/qemu.go b/go/runtime/host/tdx/qemu.go
index 9a8148f4095..f6e8e6ea46c 100644
--- a/go/runtime/host/tdx/qemu.go
+++ b/go/runtime/host/tdx/qemu.go
@@ -2,6 +2,7 @@ package tdx
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net"
@@ -283,6 +284,25 @@ func (p *qemuProvisioner) createPersistentOverlayImage(
 		// Image already exists, perform a rebase operation to account for the backing file location
 		// changing (e.g. due to an upgrade).
 		cmd := exec.Command(
+			defaultQemuImgPath,
+			"info",
+			"--output", "json",
+			image,
+		)
+		var out strings.Builder
+		cmd.Stderr = &out
+		cmd.Stdout = &out
+		if err := cmd.Run(); err != nil {
+			return "", fmt.Errorf("failed to query base image: %s\n%w", out.String(), err)
+		}
+		var info struct {
+			VirtualSize int `json:"virtual-size"`
+		}
+		if err := json.Unmarshal([]byte(out.String()), &info); err != nil {
+			return "", fmt.Errorf("malformed base image metadata: %w", err)
+		}
+
+		cmd = exec.Command(
 			defaultQemuImgPath,
 			"rebase",
 			"-u",
@@ -291,12 +311,26 @@ func (p *qemuProvisioner) createPersistentOverlayImage(
 			"-F", format,
 			imageFn,
 		)
-		var out strings.Builder
+		out.Reset()
 		cmd.Stderr = &out
 		cmd.Stdout = &out
 		if err := cmd.Run(); err != nil {
 			return "", fmt.Errorf("failed to rebase persistent overlay image: %s\n%w", out.String(), err)
 		}
+
+		// Perform a resize if needed.
+		cmd = exec.Command(
+			defaultQemuImgPath,
+			"resize",
+			imageFn,
+			fmt.Sprintf("%d", info.VirtualSize),
+		)
+		out.Reset()
+		cmd.Stderr = &out
+		cmd.Stdout = &out
+		if err := cmd.Run(); err != nil {
+			return "", fmt.Errorf("failed to resize persistent overlay image: %s\n%w", out.String(), err)
+		}
 	case errors.Is(err, os.ErrNotExist):
 		// Create image directory if it doesn't yet exist.
 		if err := common.Mkdir(imageDir); err != nil {