-
-
Notifications
You must be signed in to change notification settings - Fork 1.6k
/
Copy pathdisk_windows.go
241 lines (219 loc) · 7.44 KB
1
//go:build windows
2
3
// +build windows
4
package disk
5
6
import (
7
"bytes"
8
"context"
9
10
"fmt"
"syscall"
11
"unsafe"
12
13
"github.com/shirou/gopsutil/v3/internal/common"
14
"golang.org/x/sys/windows"
15
"golang.org/x/sys/windows/registry"
16
17
18
)
var (
19
20
21
procGetDiskFreeSpaceExW = common.Modkernel32.NewProc("GetDiskFreeSpaceExW")
procGetLogicalDriveStringsW = common.Modkernel32.NewProc("GetLogicalDriveStringsW")
procGetDriveType = common.Modkernel32.NewProc("GetDriveTypeW")
22
procGetVolumeInformation = common.Modkernel32.NewProc("GetVolumeInformationW")
23
24
25
)
var (
26
27
fileFileCompression = int64(16) // 0x00000010
fileReadOnlyVolume = int64(524288) // 0x00080000
28
)
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
// diskPerformance is an equivalent representation of DISK_PERFORMANCE in the Windows API.
// https://docs.microsoft.com/fr-fr/windows/win32/api/winioctl/ns-winioctl-disk_performance
type diskPerformance struct {
BytesRead int64
BytesWritten int64
ReadTime int64
WriteTime int64
IdleTime int64
ReadCount uint32
WriteCount uint32
QueueDepth uint32
SplitCount uint32
QueryTime int64
StorageDeviceNumber uint32
StorageManagerName [8]uint16
alignmentPadding uint32 // necessary for 32bit support, see https://github.com/elastic/beats/pull/16553
46
47
}
48
49
50
51
52
func init() {
// enable disk performance counters on Windows Server editions (needs to run as admin)
key, err := registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM\CurrentControlSet\Services\PartMgr`, registry.SET_VALUE)
if err == nil {
key.SetDWordValue("EnableCounterForIoctl", 1)
53
key.Close()
54
55
56
}
}
57
func UsageWithContext(ctx context.Context, path string) (*UsageStat, error) {
58
59
60
lpFreeBytesAvailable := int64(0)
lpTotalNumberOfBytes := int64(0)
lpTotalNumberOfFreeBytes := int64(0)
61
diskret, _, err := procGetDiskFreeSpaceExW.Call(
62
uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(path))),
63
64
65
66
uintptr(unsafe.Pointer(&lpFreeBytesAvailable)),
uintptr(unsafe.Pointer(&lpTotalNumberOfBytes)),
uintptr(unsafe.Pointer(&lpTotalNumberOfFreeBytes)))
if diskret == 0 {
67
68
return nil, err
}
69
ret := &UsageStat{
70
71
72
73
74
Path: path,
Total: uint64(lpTotalNumberOfBytes),
Free: uint64(lpTotalNumberOfFreeBytes),
Used: uint64(lpTotalNumberOfBytes) - uint64(lpTotalNumberOfFreeBytes),
UsedPercent: (float64(lpTotalNumberOfBytes) - float64(lpTotalNumberOfFreeBytes)) / float64(lpTotalNumberOfBytes) * 100,
75
76
77
78
// InodesTotal: 0,
// InodesFree: 0,
// InodesUsed: 0,
// InodesUsedPercent: 0,
79
}
80
81
return ret, nil
}
82
83
84
// PartitionsWithContext returns disk partitions.
// Since GetVolumeInformation doesn't have a timeout, this method uses context to set deadline by users.
85
func PartitionsWithContext(ctx context.Context, all bool) ([]PartitionStat, error) {
86
warnings := Warnings{
87
88
Verbose: true,
}
89
90
91
92
93
94
95
96
97
98
var errLogicalDrives error
retChan := make(chan PartitionStat)
quitChan := make(chan struct{})
defer close(quitChan)
getPartitions := func() {
defer close(retChan)
lpBuffer := make([]byte, 254)
99
100
101
102
103
diskret, _, err := procGetLogicalDriveStringsW.Call(
uintptr(len(lpBuffer)),
uintptr(unsafe.Pointer(&lpBuffer[0])))
if diskret == 0 {
104
errLogicalDrives = err
105
106
107
108
109
110
111
112
113
return
}
for _, v := range lpBuffer {
if v >= 65 && v <= 90 {
path := string(v) + ":"
typepath, _ := windows.UTF16PtrFromString(path)
typeret, _, _ := procGetDriveType.Call(uintptr(unsafe.Pointer(typepath)))
if typeret == 0 {
err := windows.GetLastError()
114
115
116
warnings.Add(err)
continue
}
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
// 2: DRIVE_REMOVABLE 3: DRIVE_FIXED 4: DRIVE_REMOTE 5: DRIVE_CDROM
if typeret == 2 || typeret == 3 || typeret == 4 || typeret == 5 {
lpVolumeNameBuffer := make([]byte, 256)
lpVolumeSerialNumber := int64(0)
lpMaximumComponentLength := int64(0)
lpFileSystemFlags := int64(0)
lpFileSystemNameBuffer := make([]byte, 256)
volpath, _ := windows.UTF16PtrFromString(string(v) + ":/")
driveret, _, err := procGetVolumeInformation.Call(
uintptr(unsafe.Pointer(volpath)),
uintptr(unsafe.Pointer(&lpVolumeNameBuffer[0])),
uintptr(len(lpVolumeNameBuffer)),
uintptr(unsafe.Pointer(&lpVolumeSerialNumber)),
uintptr(unsafe.Pointer(&lpMaximumComponentLength)),
uintptr(unsafe.Pointer(&lpFileSystemFlags)),
uintptr(unsafe.Pointer(&lpFileSystemNameBuffer[0])),
uintptr(len(lpFileSystemNameBuffer)))
if driveret == 0 {
if typeret == 5 || typeret == 2 {
continue // device is not ready will happen if there is no disk in the drive
}
warnings.Add(err)
continue
}
opts := []string{"rw"}
if lpFileSystemFlags&fileReadOnlyVolume != 0 {
opts = []string{"ro"}
}
if lpFileSystemFlags&fileFileCompression != 0 {
opts = append(opts, "compress")
}
149
150
151
select {
case retChan <- PartitionStat{
152
153
Mountpoint: path,
Device: path,
154
Fstype: string(bytes.ReplaceAll(lpFileSystemNameBuffer, []byte("\x00"), []byte(""))),
155
Opts: opts,
156
157
158
}:
case <-quitChan:
return
159
}
160
}
161
162
}
}
163
164
}
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
go getPartitions()
var ret []PartitionStat
for {
select {
case p, ok := <-retChan:
if !ok {
if errLogicalDrives != nil {
return ret, errLogicalDrives
}
return ret, warnings.Reference()
}
ret = append(ret, p)
case <-ctx.Done():
return ret, ctx.Err()
}
181
182
}
}
183
184
func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOCountersStat, error) {
185
186
187
// https://github.com/giampaolo/psutil/blob/544e9daa4f66a9f80d7bf6c7886d693ee42f0a13/psutil/arch/windows/disk.c#L83
drivemap := make(map[string]IOCountersStat, 0)
var diskPerformance diskPerformance
188
189
190
lpBuffer := make([]uint16, 254)
lpBufferLen, err := windows.GetLogicalDriveStrings(uint32(len(lpBuffer)), &lpBuffer[0])
191
if err != nil {
192
return drivemap, err
193
}
194
195
for _, v := range lpBuffer[:lpBufferLen] {
if 'A' <= v && v <= 'Z' {
196
path := string(rune(v)) + ":"
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
typepath, _ := windows.UTF16PtrFromString(path)
typeret := windows.GetDriveType(typepath)
if typeret == 0 {
return drivemap, windows.GetLastError()
}
if typeret != windows.DRIVE_FIXED {
continue
}
szDevice := fmt.Sprintf(`\\.\%s`, path)
const IOCTL_DISK_PERFORMANCE = 0x70020
h, err := windows.CreateFile(syscall.StringToUTF16Ptr(szDevice), 0, windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE, nil, windows.OPEN_EXISTING, 0, 0)
if err != nil {
if err == windows.ERROR_FILE_NOT_FOUND {
continue
}
return drivemap, err
}
defer windows.CloseHandle(h)
215
216
217
218
219
220
221
var diskPerformanceSize uint32
err = windows.DeviceIoControl(h, IOCTL_DISK_PERFORMANCE, nil, 0, (*byte)(unsafe.Pointer(&diskPerformance)), uint32(unsafe.Sizeof(diskPerformance)), &diskPerformanceSize, nil)
if err != nil {
return drivemap, err
}
drivemap[path] = IOCountersStat{
222
223
224
225
226
227
228
ReadBytes: uint64(diskPerformance.BytesRead),
WriteBytes: uint64(diskPerformance.BytesWritten),
ReadCount: uint64(diskPerformance.ReadCount),
WriteCount: uint64(diskPerformance.WriteCount),
ReadTime: uint64(diskPerformance.ReadTime / 10000 / 1000), // convert to ms: https://github.com/giampaolo/psutil/issues/1012
WriteTime: uint64(diskPerformance.WriteTime / 10000 / 1000),
Name: path,
229
}
230
231
}
}
232
return drivemap, nil
233
}
234
235
func SerialNumberWithContext(ctx context.Context, name string) (string, error) {
236
return "", common.ErrNotImplementedError
237
238
239
}
func LabelWithContext(ctx context.Context, name string) (string, error) {
240
return "", common.ErrNotImplementedError
241
}