From 720eb845ca533ae7ebe919aa350d95a7a18f2833 Mon Sep 17 00:00:00 2001
From: CapCreeperGR <nickvavour@gmail.com>
Date: Sun, 9 Jun 2024 10:12:36 +0300
Subject: [PATCH] Improved partition detection

---
 go.mod   |  1 +
 go.sum   |  3 +++
 utils.go | 40 ++++++++--------------------------------
 3 files changed, 12 insertions(+), 32 deletions(-)

diff --git a/go.mod b/go.mod
index 979a6b8..185d2e6 100644
--- a/go.mod
+++ b/go.mod
@@ -14,6 +14,7 @@ require (
 	github.com/jackmordaunt/pcidb v1.0.1 // indirect
 	github.com/jackmordaunt/wmi v1.2.4 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
+	github.com/moby/sys/mountinfo v0.7.1 // indirect
 	golang.org/x/sys v0.3.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
 )
diff --git a/go.sum b/go.sum
index fb61053..5610274 100644
--- a/go.sum
+++ b/go.sum
@@ -15,7 +15,10 @@ github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/mitchellh/go-ps v1.0.0 h1:i6ampVEEF4wQFF+bkYfwYgY+F/uYJDktmvLPf7qIgjc=
 github.com/mitchellh/go-ps v1.0.0/go.mod h1:J4lOc8z8yJs6vUwklHw2XEIiT4z4C40KtWVN3nvg8Pg=
+github.com/moby/sys/mountinfo v0.7.1 h1:/tTvQaSJRr2FshkhXiIpux6fQ2Zvc4j7tAhMTStAG2g=
+github.com/moby/sys/mountinfo v0.7.1/go.mod h1:IJb6JQeOklcdMU9F5xQ8ZALD+CUr5VlGpwtX+VE0rpI=
 golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.3.0 h1:w8ZOecv6NaNa/zC8944JTU3vz4u6Lagfk4RPQxv92NQ=
 golang.org/x/sys v0.3.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
diff --git a/utils.go b/utils.go
index a3972df..42e6d5d 100644
--- a/utils.go
+++ b/utils.go
@@ -7,6 +7,7 @@ import (
 	"github.com/BurntSushi/xgb/xinerama"
 	"github.com/jackmordaunt/ghw"
 	"github.com/mitchellh/go-ps"
+	"github.com/moby/sys/mountinfo"
 	"log"
 	"math"
 	"os"
@@ -329,11 +330,9 @@ type partition struct {
 }
 
 func getMountedPartitions() []partition {
-	partuuids, err := os.ReadDir("/dev/disk/by-partuuid")
-	if err != nil {
-		return nil
-	}
-
+	mounts, err := mountinfo.GetMounts(func(info *mountinfo.Info) (skip, stop bool) {
+		return !strings.HasPrefix(info.Source, "/dev/"), false
+	})
 	partlabels, err := os.ReadDir("/dev/disk/by-partlabel")
 	if err != nil && !os.IsNotExist(err) {
 		return nil
@@ -347,40 +346,17 @@ func getMountedPartitions() []partition {
 		labels[link] = entry.Name()
 	}
 
-	bytes, err := os.ReadFile("/proc/mounts")
-	if err != nil {
-		return nil
-	}
-
-	mountsRaw := strings.Split(string(bytes), "\n")
-	mounts := make(map[string]string)
-
-	for i := 0; i < len(mountsRaw); i++ {
-		split := strings.Split(mountsRaw[i], " ")
-		if len(split) <= 2 {
-			continue
-		}
-		mounts[split[0]] = split[1]
-	}
-
 	var partitions []partition
-	for _, entry := range partuuids {
-		link, err := filepath.EvalSymlinks(filepath.Join("/dev/disk/by-partuuid/", entry.Name()))
-		if err != nil {
-			continue
-		}
-		if _, ok := mounts[link]; !ok {
-			continue
-		}
+	for _, entry := range mounts {
 		p := partition{
-			link,
-			mounts[link],
+			entry.Source,
+			entry.Mountpoint,
 			"",
 			0,
 			0,
 			0,
 		}
-		if value, ok := labels[link]; ok {
+		if value, ok := labels[entry.Source]; ok {
 			p.Label = value
 		}
 		buf := new(syscall.Statfs_t)