common_windows.go

Documentation: github.com/shirou/gopsutil/internal/common

     1  // +build windows
     2  
     3  package common
     4  
     5  import (
     6  	"context"
     7  	"fmt"
     8  	"path/filepath"
     9  	"reflect"
    10  	"strings"
    11  	"syscall"
    12  	"unsafe"
    13  
    14  	"github.com/yusufpapurcu/wmi"
    15  	"golang.org/x/sys/windows"
    16  )
    17  
    18  // for double values
    19  type PDH_FMT_COUNTERVALUE_DOUBLE struct {
    20  	CStatus     uint32
    21  	DoubleValue float64
    22  }
    23  
    24  // for 64 bit integer values
    25  type PDH_FMT_COUNTERVALUE_LARGE struct {
    26  	CStatus    uint32
    27  	LargeValue int64
    28  }
    29  
    30  // for long values
    31  type PDH_FMT_COUNTERVALUE_LONG struct {
    32  	CStatus   uint32
    33  	LongValue int32
    34  	padding   [4]byte
    35  }
    36  
    37  // windows system const
    38  const (
    39  	ERROR_SUCCESS        = 0
    40  	ERROR_FILE_NOT_FOUND = 2
    41  	DRIVE_REMOVABLE      = 2
    42  	DRIVE_FIXED          = 3
    43  	HKEY_LOCAL_MACHINE   = 0x80000002
    44  	RRF_RT_REG_SZ        = 0x00000002
    45  	RRF_RT_REG_DWORD     = 0x00000010
    46  	PDH_FMT_LONG         = 0x00000100
    47  	PDH_FMT_DOUBLE       = 0x00000200
    48  	PDH_FMT_LARGE        = 0x00000400
    49  	PDH_INVALID_DATA     = 0xc0000bc6
    50  	PDH_INVALID_HANDLE   = 0xC0000bbc
    51  	PDH_NO_DATA          = 0x800007d5
    52  
    53  	STATUS_BUFFER_OVERFLOW      = 0x80000005
    54  	STATUS_BUFFER_TOO_SMALL     = 0xC0000023
    55  	STATUS_INFO_LENGTH_MISMATCH = 0xC0000004
    56  )
    57  
    58  const (
    59  	ProcessBasicInformation = 0
    60  	ProcessWow64Information = 26
    61  	ProcessQueryInformation = windows.PROCESS_DUP_HANDLE | windows.PROCESS_QUERY_INFORMATION
    62  
    63  	SystemExtendedHandleInformationClass = 64
    64  )
    65  
    66  var (
    67  	Modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
    68  	ModNt       = windows.NewLazySystemDLL("ntdll.dll")
    69  	ModPdh      = windows.NewLazySystemDLL("pdh.dll")
    70  	ModPsapi    = windows.NewLazySystemDLL("psapi.dll")
    71  
    72  	ProcGetSystemTimes                   = Modkernel32.NewProc("GetSystemTimes")
    73  	ProcNtQuerySystemInformation         = ModNt.NewProc("NtQuerySystemInformation")
    74  	ProcRtlGetNativeSystemInformation    = ModNt.NewProc("RtlGetNativeSystemInformation")
    75  	ProcRtlNtStatusToDosError            = ModNt.NewProc("RtlNtStatusToDosError")
    76  	ProcNtQueryInformationProcess        = ModNt.NewProc("NtQueryInformationProcess")
    77  	ProcNtReadVirtualMemory              = ModNt.NewProc("NtReadVirtualMemory")
    78  	ProcNtWow64QueryInformationProcess64 = ModNt.NewProc("NtWow64QueryInformationProcess64")
    79  	ProcNtWow64ReadVirtualMemory64       = ModNt.NewProc("NtWow64ReadVirtualMemory64")
    80  
    81  	PdhOpenQuery                = ModPdh.NewProc("PdhOpenQuery")
    82  	PdhAddEnglishCounterW       = ModPdh.NewProc("PdhAddEnglishCounterW")
    83  	PdhCollectQueryData         = ModPdh.NewProc("PdhCollectQueryData")
    84  	PdhGetFormattedCounterValue = ModPdh.NewProc("PdhGetFormattedCounterValue")
    85  	PdhCloseQuery               = ModPdh.NewProc("PdhCloseQuery")
    86  
    87  	procQueryDosDeviceW = Modkernel32.NewProc("QueryDosDeviceW")
    88  )
    89  
    90  type FILETIME struct {
    91  	DwLowDateTime  uint32
    92  	DwHighDateTime uint32
    93  }
    94  
    95  // borrowed from net/interface_windows.go
    96  func BytePtrToString(p *uint8) string {
    97  	a := (*[10000]uint8)(unsafe.Pointer(p))
    98  	i := 0
    99  	for a[i] != 0 {
   100  		i++
   101  	}
   102  	return string(a[:i])
   103  }
   104  
   105  // CounterInfo struct is used to track a windows performance counter
   106  // copied from https://github.com/mackerelio/mackerel-agent/
   107  type CounterInfo struct {
   108  	PostName    string
   109  	CounterName string
   110  	Counter     windows.Handle
   111  }
   112  
   113  // CreateQuery with a PdhOpenQuery call
   114  // copied from https://github.com/mackerelio/mackerel-agent/
   115  func CreateQuery() (windows.Handle, error) {
   116  	var query windows.Handle
   117  	r, _, err := PdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&query)))
   118  	if r != 0 {
   119  		return 0, err
   120  	}
   121  	return query, nil
   122  }
   123  
   124  // CreateCounter with a PdhAddEnglishCounterW call
   125  func CreateCounter(query windows.Handle, pname, cname string) (*CounterInfo, error) {
   126  	var counter windows.Handle
   127  	r, _, err := PdhAddEnglishCounterW.Call(
   128  		uintptr(query),
   129  		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(cname))),
   130  		0,
   131  		uintptr(unsafe.Pointer(&counter)))
   132  	if r != 0 {
   133  		return nil, err
   134  	}
   135  	return &CounterInfo{
   136  		PostName:    pname,
   137  		CounterName: cname,
   138  		Counter:     counter,
   139  	}, nil
   140  }
   141  
   142  // GetCounterValue get counter value from handle
   143  // adapted from https://github.com/mackerelio/mackerel-agent/
   144  func GetCounterValue(counter windows.Handle) (float64, error) {
   145  	var value PDH_FMT_COUNTERVALUE_DOUBLE
   146  	r, _, err := PdhGetFormattedCounterValue.Call(uintptr(counter), PDH_FMT_DOUBLE, uintptr(0), uintptr(unsafe.Pointer(&value)))
   147  	if r != 0 && r != PDH_INVALID_DATA {
   148  		return 0.0, err
   149  	}
   150  	return value.DoubleValue, nil
   151  }
   152  
   153  type Win32PerformanceCounter struct {
   154  	PostName    string
   155  	CounterName string
   156  	Query       windows.Handle
   157  	Counter     windows.Handle
   158  }
   159  
   160  func NewWin32PerformanceCounter(postName, counterName string) (*Win32PerformanceCounter, error) {
   161  	query, err := CreateQuery()
   162  	if err != nil {
   163  		return nil, err
   164  	}
   165  	var counter = Win32PerformanceCounter{
   166  		Query:       query,
   167  		PostName:    postName,
   168  		CounterName: counterName,
   169  	}
   170  	r, _, err := PdhAddEnglishCounterW.Call(
   171  		uintptr(counter.Query),
   172  		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(counter.CounterName))),
   173  		0,
   174  		uintptr(unsafe.Pointer(&counter.Counter)),
   175  	)
   176  	if r != 0 {
   177  		return nil, err
   178  	}
   179  	return &counter, nil
   180  }
   181  
   182  func (w *Win32PerformanceCounter) GetValue() (float64, error) {
   183  	r, _, err := PdhCollectQueryData.Call(uintptr(w.Query))
   184  	if r != 0 && err != nil {
   185  		if r == PDH_NO_DATA {
   186  			return 0.0, fmt.Errorf("%w: this counter has not data", err)
   187  		}
   188  		return 0.0, err
   189  	}
   190  
   191  	return GetCounterValue(w.Counter)
   192  }
   193  
   194  func ProcessorQueueLengthCounter() (*Win32PerformanceCounter, error) {
   195  	return NewWin32PerformanceCounter("processor_queue_length", `\System\Processor Queue Length`)
   196  }
   197  
   198  // WMIQueryWithContext - wraps wmi.Query with a timed-out context to avoid hanging
   199  func WMIQueryWithContext(ctx context.Context, query string, dst interface{}, connectServerArgs ...interface{}) error {
   200  	if _, ok := ctx.Deadline(); !ok {
   201  		ctxTimeout, cancel := context.WithTimeout(ctx, Timeout)
   202  		defer cancel()
   203  		ctx = ctxTimeout
   204  	}
   205  
   206  	errChan := make(chan error, 1)
   207  	go func() {
   208  		errChan <- wmi.Query(query, dst, connectServerArgs...)
   209  	}()
   210  
   211  	select {
   212  	case <-ctx.Done():
   213  		return ctx.Err()
   214  	case err := <-errChan:
   215  		return err
   216  	}
   217  }
   218  
   219  // Convert paths using native DOS format like:
   220  //   "\Device\HarddiskVolume1\Windows\systemew\file.txt"
   221  // into:
   222  //   "C:\Windows\systemew\file.txt"
   223  func ConvertDOSPath(p string) string {
   224  	rawDrive := strings.Join(strings.Split(p, `\`)[:3], `\`)
   225  
   226  	for d := 'A'; d <= 'Z'; d++ {
   227  		szDeviceName := string(d) + ":"
   228  		szTarget := make([]uint16, 512)
   229  		ret, _, _ := procQueryDosDeviceW.Call(uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(szDeviceName))),
   230  			uintptr(unsafe.Pointer(&szTarget[0])),
   231  			uintptr(len(szTarget)))
   232  		if ret != 0 && windows.UTF16ToString(szTarget[:]) == rawDrive {
   233  			return filepath.Join(szDeviceName, p[len(rawDrive):])
   234  		}
   235  	}
   236  	return p
   237  }
   238  
   239  type NtStatus uint32
   240  
   241  func (s NtStatus) Error() error {
   242  	if s == 0 {
   243  		return nil
   244  	}
   245  	return fmt.Errorf("NtStatus 0x%08x", uint32(s))
   246  }
   247  
   248  func (s NtStatus) IsError() bool {
   249  	return s>>30 == 3
   250  }
   251  
   252  type SystemExtendedHandleTableEntryInformation struct {
   253  	Object                uintptr
   254  	UniqueProcessId       uintptr
   255  	HandleValue           uintptr
   256  	GrantedAccess         uint32
   257  	CreatorBackTraceIndex uint16
   258  	ObjectTypeIndex       uint16
   259  	HandleAttributes      uint32
   260  	Reserved              uint32
   261  }
   262  
   263  type SystemExtendedHandleInformation struct {
   264  	NumberOfHandles uintptr
   265  	Reserved        uintptr
   266  	Handles         [1]SystemExtendedHandleTableEntryInformation
   267  }
   268  
   269  // CallWithExpandingBuffer https://github.com/hillu/go-ntdll
   270  func CallWithExpandingBuffer(fn func() NtStatus, buf *[]byte, resultLength *uint32) NtStatus {
   271  	for {
   272  		if st := fn(); st == STATUS_BUFFER_OVERFLOW || st == STATUS_BUFFER_TOO_SMALL || st == STATUS_INFO_LENGTH_MISMATCH {
   273  			if int(*resultLength) <= cap(*buf) {
   274  				(*reflect.SliceHeader)(unsafe.Pointer(buf)).Len = int(*resultLength)
   275  			} else {
   276  				*buf = make([]byte, int(*resultLength))
   277  			}
   278  			continue
   279  		} else {
   280  			if !st.IsError() {
   281  				*buf = (*buf)[:int(*resultLength)]
   282  			}
   283  			return st
   284  		}
   285  	}
   286  }
   287  
   288  func NtQuerySystemInformation(
   289  	SystemInformationClass uint32,
   290  	SystemInformation *byte,
   291  	SystemInformationLength uint32,
   292  	ReturnLength *uint32,
   293  ) NtStatus {
   294  	r0, _, _ := ProcNtQuerySystemInformation.Call(
   295  		uintptr(SystemInformationClass),
   296  		uintptr(unsafe.Pointer(SystemInformation)),
   297  		uintptr(SystemInformationLength),
   298  		uintptr(unsafe.Pointer(ReturnLength)))
   299  	return NtStatus(r0)
   300  }
   301
View as plain text