From 6117a0ff30d5f9cb9650cb8bd63b7c7948db4029 Mon Sep 17 00:00:00 2001
From: Omar Abdulaziz <omarabdul3ziz@gmail.com>
Date: Mon, 2 Dec 2024 09:44:49 +0200
Subject: [PATCH] - list partitions in devices - include the zospxe check for
 partitions as well

---
 pkg/storage/filesystem/device.go | 20 ++++---
 pkg/storage/storage.go           | 91 ++++++++++++++++++--------------
 2 files changed, 66 insertions(+), 45 deletions(-)

diff --git a/pkg/storage/filesystem/device.go b/pkg/storage/filesystem/device.go
index 3c49ae039..0808377c6 100644
--- a/pkg/storage/filesystem/device.go
+++ b/pkg/storage/filesystem/device.go
@@ -51,12 +51,13 @@ type blockDevices struct {
 type DeviceInfo struct {
 	mgr DeviceManager
 
-	Path       string `json:"path"`
-	Label      string `json:"label"`
-	Size       uint64 `json:"size"`
-	Filesystem FSType `json:"fstype"`
-	Rota       bool   `json:"rota"`
-	Subsystems string `json:"subsystems"`
+	Path       string       `json:"path"`
+	Label      string       `json:"label"`
+	Size       uint64       `json:"size"`
+	Filesystem FSType       `json:"fstype"`
+	Rota       bool         `json:"rota"`
+	Subsystems string       `json:"subsystems"`
+	Children   []DeviceInfo `json:"children,omitempty"`
 }
 
 func (i *DeviceInfo) Name() string {
@@ -77,6 +78,10 @@ func (d *DeviceInfo) Mountpoint(ctx context.Context) (string, error) {
 	return d.mgr.Mountpoint(ctx, d.Path)
 }
 
+func (d *DeviceInfo) IsPXEPartition() bool {
+	return d.Label == "ZOSPXE"
+}
+
 // lsblkDeviceManager uses the lsblk utility to scann the disk for devices, and
 // caches the result.
 //
@@ -187,6 +192,9 @@ func (l *lsblkDeviceManager) lsblk(ctx context.Context) ([]DeviceInfo, error) {
 
 	for i := range devices.BlockDevices {
 		devices.BlockDevices[i].mgr = l
+		for j := range devices.BlockDevices[i].Children {
+			devices.BlockDevices[i].Children[j].mgr = l
+		}
 	}
 
 	return devices.BlockDevices, nil
diff --git a/pkg/storage/storage.go b/pkg/storage/storage.go
index bf4daeac6..cbaa6ec42 100644
--- a/pkg/storage/storage.go
+++ b/pkg/storage/storage.go
@@ -197,6 +197,53 @@ func (s *Module) poolType(pool filesystem.Pool, vm bool) (zos.DeviceType, error)
 	return typ, nil
 }
 
+func (s *Module) mountPool(device filesystem.DeviceInfo, vm bool) {
+	log.Debug().Str("path", device.Path).Msg("mounting device")
+
+	if device.IsPXEPartition() {
+		log.Error().Str("device", device.Path).Msg("device has 'zospxe' label")
+		s.brokenDevices = append(s.brokenDevices, pkg.BrokenDevice{Path: device.Path, Err: fmt.Errorf("device is a PXE partition")})
+		return
+	}
+
+	pool, err := filesystem.NewBtrfsPool(device)
+	if err != nil {
+		log.Error().Err(err).Str("device", device.Path).Msg("failed to create pool on device")
+		s.brokenDevices = append(s.brokenDevices, pkg.BrokenDevice{Path: device.Path, Err: err})
+		return
+	}
+
+	_, err = pool.Mount()
+	if err != nil {
+		log.Error().Err(err).Str("device", device.Path).Msg("failed to mount pool")
+		s.brokenPools = append(s.brokenPools, pkg.BrokenPool{Label: pool.Name(), Err: err})
+		return
+	}
+
+	usage, err := pool.Usage()
+	if err != nil {
+		log.Error().Err(err).Str("pool", pool.Name()).Str("device", device.Path).Msg("failed to get usage of pool")
+	}
+
+	typ, err := s.poolType(pool, vm)
+	if err != nil {
+		log.Error().Str("device", device.Path).Err(err).Msg("failed to get device type")
+		s.brokenDevices = append(s.brokenDevices, pkg.BrokenDevice{Path: device.Path, Err: err})
+		return
+	}
+
+	switch typ {
+	case zos.SSDDevice:
+		s.totalSSD += usage.Size
+		s.ssds = append(s.ssds, pool)
+	case zos.HDDDevice:
+		s.totalHDD += usage.Size
+		s.hdds = append(s.hdds, pool)
+	default:
+		log.Error().Str("type", string(typ)).Str("device", device.Path).Msg("unknown device type")
+	}
+}
+
 /*
 *
 initialize, must be called at least onetime each boot.
@@ -229,47 +276,13 @@ func (s *Module) initialize(ctx context.Context) error {
 	for _, device := range devices {
 		log.Debug().Msgf("device: %+v", device)
 
-		if device.Label == PXELABEL {
-			log.Error().Err(err).Str("device", device.Path).Msg("device has 'zospxe' label")
-			s.brokenDevices = append(s.brokenDevices, pkg.BrokenDevice{Path: device.Path, Err: err})
-			continue
-		}
-
-		pool, err := filesystem.NewBtrfsPool(device)
-		if err != nil {
-			log.Error().Err(err).Str("device", device.Path).Msg("failed to create pool on device")
-			s.brokenDevices = append(s.brokenDevices, pkg.BrokenDevice{Path: device.Path, Err: err})
-			continue
-		}
-
-		_, err = pool.Mount()
-		if err != nil {
-			log.Error().Err(err).Str("device", device.Path).Msg("failed to mount pool")
-			s.brokenPools = append(s.brokenPools, pkg.BrokenPool{Label: pool.Name(), Err: err})
-			continue
-		}
-		usage, err := pool.Usage()
-		if err != nil {
-			log.Error().Err(err).Str("pool", pool.Name()).Str("device", device.Path).Msg("failed to get usage of pool")
-		}
-
-		typ, err := s.poolType(pool, vm)
-		if err != nil {
-			log.Error().Str("device", device.Path).Err(err).Msg("failed to get device type")
-			s.brokenDevices = append(s.brokenDevices, pkg.BrokenDevice{Path: device.Path, Err: err})
+		if len(device.Children) != 0 {
+			for _, part := range device.Children {
+				s.mountPool(part, vm)
+			}
 			continue
 		}
-
-		switch typ {
-		case zos.SSDDevice:
-			s.totalSSD += usage.Size
-			s.ssds = append(s.ssds, pool)
-		case zos.HDDDevice:
-			s.totalHDD += usage.Size
-			s.hdds = append(s.hdds, pool)
-		default:
-			log.Error().Str("type", string(typ)).Str("device", device.Path).Msg("unknown device type")
-		}
+		s.mountPool(device, vm)
 	}
 
 	log.Info().