fix(fs): Expand SystemRoot in file paths

Author: rabbitstack

pkg/fs/dev.go

@@ -23,6 +23,7 @@ package fs
 
 import (
 	"github.com/rabbitstack/fibratus/pkg/sys"
+	"os"
 	"strings"
 )
 
@@ -36,14 +37,16 @@ type DevMapper interface {
 }
 
 type mapper struct {
-	cache map[string]string
+	cache   map[string]string
+	sysroot string
 }
 
 // NewDevMapper creates a new instance of the DOS device replacer.
 func NewDevMapper() DevMapper {
 	m := &mapper{
 		cache: make(map[string]string),
 	}
+
 	// loop through logical drives and query the DOS device name
 	for _, drive := range sys.GetLogicalDrives() {
 		device, err := sys.QueryDosDevice(drive)
@@ -52,30 +55,49 @@ func NewDevMapper() DevMapper {
 		}
 		m.cache[device] = drive
 	}
+
+	// resolve the SystemRoot environment variable
+	m.sysroot = os.Getenv("SystemRoot")
+	if m.sysroot == "" {
+		m.sysroot = os.Getenv("SYSTEMROOT")
+	}
+
 	return m
 }
 
 func (m *mapper) Convert(filename string) string {
 	if filename == "" || len(filename) < deviceOffset {
 		return filename
 	}
-	i := strings.Index(filename[deviceOffset:], "\\")
-	if i < 0 {
+
+	// find the backslash index
+	n := strings.Index(filename[deviceOffset:], "\\")
+	if n < 0 {
 		if f, ok := m.cache[filename]; ok {
 			return f
 		}
 		return filename
 	}
-	dev := filename[:i+deviceOffset]
-	// convert Windows Sandbox path to native path
-	if dev == vmsmbDevice {
-		n := strings.Index(filename, "os")
-		if n > 0 {
-			return "C:" + filename[n+2:]
-		}
-	}
+
+	dev := filename[:n+deviceOffset]
 	if drive, ok := m.cache[dev]; ok {
+		// the mapping for the DOS device exists
 		return strings.Replace(filename, dev, drive, 1)
 	}
+
+	switch {
+	case dev == vmsmbDevice:
+		// convert Windows Sandbox path to native path
+		if n := strings.Index(filename, "os"); n > 0 {
+			return "C:" + filename[n+2:]
+		}
+	case strings.HasPrefix(filename, "\\SystemRoot"):
+		// normalize paths starting with SystemRoot
+		return strings.Replace(filename, "\\SystemRoot", m.sysroot, 1)
+	case strings.HasPrefix(filename, "\\SYSTEMROOT"):
+		// normalize paths starting with SYSTEMROOT
+		return strings.Replace(filename, "\\SYSTEMROOT", m.sysroot, 1)
+	}
+
 	return filename
 }

pkg/fs/dev_test.go

@@ -57,23 +57,40 @@ var drives = []string{
 	"Z"}
 
 func TestConvertDosDevice(t *testing.T) {
-	mapper := NewDevMapper()
+	m := NewDevMapper()
 	files := make([]string, 0, len(drives))
+
 	for _, drive := range drives {
 		files = append(files, fmt.Sprintf("%s:\\Windows\\system32\\kernel32.dll", drive))
 	}
+
 	var filename string
 	for i := 0; i < len(drives); i++ {
-		filename = mapper.Convert(fmt.Sprintf("\\Device\\HarddiskVolume%d\\Windows\\system32\\kernel32.dll", i))
+		filename = m.Convert(fmt.Sprintf("\\Device\\HarddiskVolume%d\\Windows\\system32\\kernel32.dll", i))
 		if !strings.HasPrefix(filename, "\\Device") {
 			break
 		}
 	}
 	assert.Contains(t, files, filename)
-}
 
-func TestConvertVmsmbDevice(t *testing.T) {
-	mapper := NewDevMapper()
-	path := "\\Device\\vmsmb\\VSMB-{dcc079ae-60ba-4d07-847c-3493609c0870}\\os\\Windows\\System32\\ntdll.dll"
-	assert.Equal(t, "C:\\Windows\\System32\\ntdll.dll", mapper.Convert(path))
+	m.(*mapper).cache["\\Device\\HarddiskVolume1"] = "C:"
+	m.(*mapper).sysroot = "C:\\Windows"
+
+	var tests = []struct {
+		inputFilename    string
+		expectedFilename string
+	}{
+		{"\\Device\\HarddiskVolume1\\Windows\\system32\\kernel32.dll", "C:\\Windows\\system32\\kernel32.dll"},
+		{"\\Device\\HarddiskVolume5\\Windows\\system32\\kernel32.dll", "\\Device\\HarddiskVolume5\\Windows\\system32\\kernel32.dll"},
+		{"\\Device\\vmsmb\\VSMB-{dcc079ae-60ba-4d07-847c-3493609c0870}\\os\\Windows\\System32\\ntdll.dll", "C:\\Windows\\System32\\ntdll.dll"},
+		{"\\SystemRoot\\system32\\drivers\\wd\\WdNisDrv.sys", "C:\\Windows\\system32\\drivers\\wd\\WdNisDrv.sys"},
+		{"\\SYSTEMROOT\\system32\\drivers\\wd\\WdNisDrv.sys", "C:\\Windows\\system32\\drivers\\wd\\WdNisDrv.sys"},
+		{"\\Device\\Mup", "\\Device\\Mup"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.inputFilename, func(t *testing.T) {
+			assert.Equal(t, tt.expectedFilename, m.Convert(tt.inputFilename))
+		})
+	}
 }