logical_disk: safely enumerate all volume mount paths (#2584)

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/collector/logical_disk/logical_disk.go

@@ -720,7 +720,7 @@ func getVolumeInfo(volumes map[string]string, rootDrive string) (volumeInfo, err
 
 func getAllMountedVolumes() (map[string]string, error) {
 	guidBuf := make([]uint16, windows.MAX_PATH+1)
-	guidBufLen := uint32(len(guidBuf) * 2)
+	guidBufLen := uint32(len(guidBuf))
 
 	hFindVolume, err := windows.FindFirstVolume(&guidBuf[0], guidBufLen)
 	if err != nil {
@@ -743,38 +743,57 @@ func getAllMountedVolumes() (map[string]string, error) {
 			}
 		}
 
-		var rootPathLen uint32
-
-		rootPathBuf := make([]uint16, windows.MAX_PATH+1)
-		rootPathBufLen := uint32(len(rootPathBuf) * 2)
-
-		for {
-			err = windows.GetVolumePathNamesForVolumeName(&guidBuf[0], &rootPathBuf[0], rootPathBufLen, &rootPathLen)
-			if err == nil {
-				break
-			}
-
-			if errors.Is(err, windows.ERROR_FILE_NOT_FOUND) {
-				// the volume is not mounted
-				break
-			}
-
-			if errors.Is(err, windows.ERROR_NO_MORE_FILES) {
-				rootPathBuf = make([]uint16, (rootPathLen+1)/2)
-
-				continue
-			}
-
+		mountPoints, err := getVolumeMountPaths(&guidBuf[0], windows.GetVolumePathNamesForVolumeName)
+		if err != nil {
 			return nil, fmt.Errorf("GetVolumePathNamesForVolumeName: %w", err)
 		}
 
-		mountPoint := windows.UTF16ToString(rootPathBuf)
+		volumeName := strings.TrimSuffix(windows.UTF16ToString(guidBuf), `\`)
+		for _, mountPoint := range mountPoints {
+			volumes[strings.TrimSuffix(mountPoint, `\`)] = volumeName
+		}
+	}
+}
+
+// getVolumeMountPaths reads the MULTI_SZ mount paths. The API measures both
+// buffer capacity and required capacity in UTF-16 characters, not bytes.
+func getVolumeMountPaths(volumeName *uint16, query func(*uint16, *uint16, uint32, *uint32) error) ([]string, error) {
+	buf := make([]uint16, windows.MAX_PATH+1)
+
+	for {
+		var required uint32
+
+		err := query(volumeName, &buf[0], uint32(len(buf)), &required)
+		if errors.Is(err, windows.ERROR_FILE_NOT_FOUND) {
+			return nil, nil
+		}
+
+		if errors.Is(err, windows.ERROR_MORE_DATA) {
+			if required <= uint32(len(buf)) {
+				return nil, fmt.Errorf("invalid required mount path buffer size: %d", required)
+			}
+
+			buf = make([]uint16, required)
 
-		// Skip unmounted volumes
-		if len(mountPoint) == 0 {
 			continue
 		}
 
-		volumes[strings.TrimSuffix(mountPoint, `\`)] = strings.TrimSuffix(windows.UTF16ToString(guidBuf), `\`)
+		if err != nil {
+			return nil, err
+		}
+
+		var paths []string
+
+		for len(buf) > 0 && buf[0] != 0 {
+			end := slices.Index(buf, uint16(0))
+			if end < 0 {
+				return nil, errors.New("unterminated mount path")
+			}
+
+			paths = append(paths, windows.UTF16ToString(buf[:end]))
+			buf = buf[end+1:]
+		}
+
+		return paths, nil
 	}
 }

internal/collector/logical_disk/volume_paths_test.go

@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Copyright The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build windows
+
+package logical_disk
+
+import (
+	"fmt"
+	"testing"
+	"unsafe"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/windows"
+)
+
+func TestVolumeMountPathsResize(t *testing.T) {
+	t.Parallel()
+
+	want := make([]string, 0, 30)
+
+	var data []uint16
+
+	for i := range 30 {
+		path := fmt.Sprintf(`C:\mounts\volume-%02d\`, i)
+		want = append(want, path)
+		encoded, err := windows.UTF16FromString(path)
+		require.NoError(t, err)
+
+		data = append(data, encoded...)
+	}
+
+	data = append(data, 0)
+	calls := 0
+	query := func(_ *uint16, dst *uint16, capacity uint32, required *uint32) error {
+		calls++
+		*required = uint32(len(data))
+
+		if calls == 1 {
+			require.Equal(t, uint32(windows.MAX_PATH+1), capacity)
+
+			return windows.ERROR_MORE_DATA
+		}
+
+		require.Equal(t, uint32(len(data)), capacity)
+		copy(unsafe.Slice(dst, capacity), data)
+
+		return nil
+	}
+
+	paths, err := getVolumeMountPaths(nil, query)
+	require.NoError(t, err)
+	require.Equal(t, want, paths)
+	require.Equal(t, 2, calls)
+}
+
+func TestMountedVolumeEnumeration(t *testing.T) {
+	t.Parallel()
+
+	volumes, err := getAllMountedVolumes()
+	require.NoError(t, err)
+	require.NotEmpty(t, volumes)
+
+	for _, volume := range volumes {
+		name, err := windows.UTF16PtrFromString(volume + `\`)
+		require.NoError(t, err)
+		paths, err := getVolumeMountPaths(name, windows.GetVolumePathNamesForVolumeName)
+		require.NoError(t, err)
+		require.NotEmpty(t, paths)
+	}
+}