From f49c405509e05dfaec065bcfeb339aa0d328bdad Mon Sep 17 00:00:00 2001
From: authrequest <admin@hypedcarts.com>
Date: Wed, 12 Feb 2025 14:39:43 -0800
Subject: [PATCH] Support of img.xz images on JetKVM storage mount

---
 go.mod              |  1 +
 go.sum              |  2 ++
 usb_mass_storage.go | 62 ++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 5ddcfb6..72f2480 100644
--- a/go.mod
+++ b/go.mod
@@ -20,6 +20,7 @@ require (
 	github.com/pion/webrtc/v4 v4.0.0
 	github.com/pojntfx/go-nbd v0.3.2
 	github.com/psanford/httpreadat v0.1.0
+	github.com/ulikunitz/xz v0.5.12
 	github.com/vishvananda/netlink v1.3.0
 	golang.org/x/crypto v0.28.0
 	golang.org/x/net v0.30.0
diff --git a/go.sum b/go.sum
index be21917..2084906 100644
--- a/go.sum
+++ b/go.sum
@@ -138,6 +138,8 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS
 github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
 github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE=
 github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
+github.com/ulikunitz/xz v0.5.12 h1:37Nm15o69RwBkXM0J6A5OlE67RZTfzUxTj8fB3dfcsc=
+github.com/ulikunitz/xz v0.5.12/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
 github.com/vishvananda/netlink v1.3.0 h1:X7l42GfcV4S6E4vHTsw48qbrV+9PVojNfIhZcwQdrZk=
 github.com/vishvananda/netlink v1.3.0/go.mod h1:i6NetklAujEcC6fK0JPjT8qSwWyO0HLn4UKG+hGqeJs=
 github.com/vishvananda/netns v0.0.4 h1:Oeaw1EM2JMxD51g9uhtC0D7erkIjgmj8+JZc26m1YX8=
diff --git a/usb_mass_storage.go b/usb_mass_storage.go
index b897c20..154c84c 100644
--- a/usb_mass_storage.go
+++ b/usb_mass_storage.go
@@ -22,6 +22,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/pion/webrtc/v4"
+	"github.com/ulikunitz/xz"
 )
 
 const massStorageName = "mass_storage.usb0"
@@ -269,16 +270,48 @@ func rpcMountWithStorage(filename string, mode VirtualMediaMode) error {
 		return fmt.Errorf("failed to get file info: %w", err)
 	}
 
-	err = setMassStorageImage(fullPath)
+	// Handle XZ compressed images
+	if strings.HasSuffix(filename, ".img.xz") {
+		logger.Info("Mounting compressed XZ image")
+
+		// Create temporary file for decompressed image
+		decompressedPath := filepath.Join(imagesFolder, strings.TrimSuffix(filename, ".xz")+".temp")
+		defer os.Remove(decompressedPath) // Clean up temp file after mounting
+
+		err = decompressXZImage(fullPath, decompressedPath)
+		if err != nil {
+			return fmt.Errorf("failed to decompress XZ image: %w", err)
+		}
+
+		// Mount the decompressed image
+		err = mountImage(decompressedPath)
+		if err != nil {
+			return fmt.Errorf("failed to mount decompressed image: %w", err)
+		}
+
+		currentVirtualMediaState = &VirtualMediaState{
+			Source:   Storage,
+			Mode:     mode,
+			Filename: filename,
+			Size:     fileInfo.Size(),
+		}
+
+		return nil
+	}
+
+	// Handle regular images
+	err = mountImage(fullPath)
 	if err != nil {
 		return fmt.Errorf("failed to set mass storage image: %w", err)
 	}
+
 	currentVirtualMediaState = &VirtualMediaState{
 		Source:   Storage,
 		Mode:     mode,
 		Filename: filename,
 		Size:     fileInfo.Size(),
 	}
+
 	return nil
 }
 
@@ -556,3 +589,30 @@ func handleUploadHttp(c *gin.Context) {
 
 	c.JSON(http.StatusOK, gin.H{"message": "Upload completed"})
 }
+
+// Add this helper function to handle XZ decompression
+func decompressXZImage(sourcePath, destPath string) error {
+	sourceFile, err := os.Open(sourcePath)
+	if err != nil {
+		return fmt.Errorf("failed to open source file: %w", err)
+	}
+	defer sourceFile.Close()
+
+	reader, err := xz.NewReader(sourceFile)
+	if err != nil {
+		return fmt.Errorf("failed to create XZ reader: %w", err)
+	}
+
+	destFile, err := os.Create(destPath)
+	if err != nil {
+		return fmt.Errorf("failed to create destination file: %w", err)
+	}
+	defer destFile.Close()
+
+	_, err = io.Copy(destFile, reader)
+	if err != nil {
+		return fmt.Errorf("failed to decompress file: %w", err)
+	}
+
+	return nil
+}