From 47c58355e2e74873bc0033cae3b6e952e4c1a991 Mon Sep 17 00:00:00 2001 From: Glazier Bot Date: Sun, 9 Aug 2026 22:33:15 -0700 Subject: [PATCH] internal PiperOrigin-RevId: 961943080 --- go/device/device_test.go | 392 +++++++++++++++++++++++++++++++++++- go/device/device_windows.go | 295 ++++++++++++++++++++------- 2 files changed, 618 insertions(+), 69 deletions(-) diff --git a/go/device/device_test.go b/go/device/device_test.go index f4528e21..923b65bc 100644 --- a/go/device/device_test.go +++ b/go/device/device_test.go @@ -1,3 +1,5 @@ +//go:build windows + // Copyright 2021 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,22 +17,400 @@ package device import ( + "errors" + "fmt" "io/ioutil" "os" "path/filepath" + "strings" "testing" "github.com/google/go-cmp/cmp" ) +func TestChassisType(t *testing.T) { + origGetChassisType := getChassisType + defer func() { getChassisType = origGetChassisType }() + + tests := []struct { + name string + mock Type + mockErr error + want Type + wantErr bool + }{ + { + name: "Desktop", + mock: Desktop, + want: Desktop, + }, + { + name: "Laptop", + mock: Laptop, + want: Laptop, + }, + { + name: "Other", + mock: Other, + want: Other, + }, + { + name: "Unknown", + mock: Unknown, + want: Unknown, + }, + { + name: "Error", + mock: Unknown, + mockErr: errors.New("failed to read SMBIOS"), + want: Unknown, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + getChassisType = func() (Type, error) { + return tt.mock, tt.mockErr + } + got, err := ChassisType() + if (err != nil) != tt.wantErr { + t.Fatalf("ChassisType() error = %v, wantErr = %v", err, tt.wantErr) + } + if got != tt.want { + t.Errorf("ChassisType() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestGetDomainRole(t *testing.T) { + origRegistryGetString := registryGetString + defer func() { registryGetString = origRegistryGetString }() + + tests := []struct { + name string + mockValue string + mockErr error + want DomainRole + wantErr bool + }{ + { + name: "Workstation", + mockValue: "WinNT", + want: Workstation, + }, + { + name: "Server", + mockValue: "ServerNT", + want: Server, + }, + { + name: "DomainController", + mockValue: "LanmanNT", + want: DomainController, + }, + { + name: "UnknownRole", + mockValue: "Other", + want: RoleUnknown, + }, + { + name: "RegistryError", + mockErr: errors.New("registry read failed"), + want: RoleUnknown, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + registryGetString = func(path, name string) (string, error) { + if path != `SYSTEM\CurrentControlSet\Control\ProductOptions` || name != "ProductType" { + return "", fmt.Errorf("unexpected registry query: path=%q, name=%q", path, name) + } + return tt.mockValue, tt.mockErr + } + + got, err := GetDomainRole() + if (err != nil) != tt.wantErr { + t.Fatalf("GetDomainRole() error = %v, wantErr = %v", err, tt.wantErr) + } + if got != tt.want { + t.Errorf("GetDomainRole() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestGetDomainJoined(t *testing.T) { + origNetGetJoinInformation := netGetJoinInformation + origNetAPIBufferFree := netAPIBufferFree + defer func() { + netGetJoinInformation = origNetGetJoinInformation + netAPIBufferFree = origNetAPIBufferFree + }() + + tests := []struct { + name string + mockErr error + joinStatus uint32 + mockDomainName bool + want bool + wantErr bool + }{ + { + name: "DomainJoined", + joinStatus: 3, // NetSetupDomainName. + mockDomainName: true, + want: true, + }, + { + name: "NotDomainJoined", + joinStatus: 1, // NetSetupUnjoined + mockDomainName: true, + want: false, + }, + { + name: "APIFailure", + mockErr: errors.New("ERROR_NO_SUCH_DOMAIN"), + want: false, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + bufferFreed := false + netGetJoinInformation = func(server *uint16, name **uint16, bufType *uint32) error { + if tt.mockErr != nil { + return tt.mockErr + } + if name != nil && tt.mockDomainName { + var dummy uint16 = 'D' + *name = &dummy + } + if bufType != nil { + *bufType = tt.joinStatus + } + return nil + } + netAPIBufferFree = func(buf *byte) error { + bufferFreed = true + return nil + } + + got, err := GetDomainJoined() + if (err != nil) != tt.wantErr { + t.Fatalf("GetDomainJoined() error = %v, wantErr = %v", err, tt.wantErr) + } + if got != tt.want { + t.Errorf("GetDomainJoined() = %v, want %v", got, tt.want) + } + if tt.mockDomainName && !bufferFreed { + t.Errorf("GetDomainJoined() did not free domain name buffer") + } + }) + } +} + +func TestModel(t *testing.T) { + origRegistryGetString := registryGetString + defer func() { registryGetString = origRegistryGetString }() + + tests := []struct { + name string + manufacturer string + manufacturerErr error + family string + familyErr error + productName string + productNameErr error + want string + wantErr bool + }{ + { + name: "NonLenovo", + manufacturer: "Dell Inc.", + productName: "Precision 5530", + want: "Precision 5530", + }, + { + name: "Lenovo", + manufacturer: "LENOVO", + family: "ThinkPad P1 Gen 4", + want: "ThinkPad P1 Gen 4", + }, + { + name: "ManufacturerError", + manufacturerErr: errors.New("registry read error"), + want: "unknown", + wantErr: true, + }, + { + name: "NonLenovoProductNameError", + manufacturer: "Dell Inc.", + productNameErr: errors.New("product name missing"), + want: "unknown", + wantErr: true, + }, + { + name: "LenovoFamilyError", + manufacturer: "Lenovo", + familyErr: errors.New("family missing"), + want: "unknown", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + registryGetString = func(path, name string) (string, error) { + if path != `HARDWARE\DESCRIPTION\System\BIOS` { + return "", fmt.Errorf("unexpected registry path %q", path) + } + switch name { + case "SystemManufacturer": + return tt.manufacturer, tt.manufacturerErr + case "SystemFamily": + return tt.family, tt.familyErr + case "SystemProductName": + return tt.productName, tt.productNameErr + default: + return "", fmt.Errorf("unexpected registry value name %q", name) + } + } + + got, err := Model() + if (err != nil) != tt.wantErr { + t.Fatalf("Model() error = %v, wantErr = %v", err, tt.wantErr) + } + if got != tt.want { + t.Errorf("Model() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestSite(t *testing.T) { + origWmiQuery := wmiQuery + defer func() { wmiQuery = origWmiQuery }() + + tests := []struct { + name string + domain string + mockResult []Win32_NTDomain + queryErr error + want string + wantErr bool + }{ + { + name: "Success", + domain: "some.place.com", + mockResult: []Win32_NTDomain{ + { + ClientSiteName: "NYC", + DomainControllerName: `\\dc1.some.place.com`, + }, + }, + want: "NYC", + }, + { + name: "EmptyResult", + domain: "some.place.com", + mockResult: []Win32_NTDomain{}, + want: "", + }, + { + name: "WMIError", + domain: "some.place.com", + queryErr: errors.New("WMI connection failed"), + want: "", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var queryRan string + wmiQuery = func(query string, dst any, connectServerArgs ...any) error { + queryRan = query + if tt.queryErr != nil { + return tt.queryErr + } + res, ok := dst.(*[]Win32_NTDomain) + if !ok { + return fmt.Errorf("unexpected destination type: %T", dst) + } + *res = tt.mockResult + return nil + } + + got, err := Site(tt.domain) + if (err != nil) != tt.wantErr { + t.Fatalf("Site(%q) error = %v, wantErr = %v", tt.domain, err, tt.wantErr) + } + if got != tt.want { + t.Errorf("Site(%q) = %q, want %q", tt.domain, got, tt.want) + } + if !strings.Contains(queryRan, fmt.Sprintf("WHERE DomainName='%s'", tt.domain)) { + t.Errorf("Site(%q) query = %q, did not contain expected WHERE clause", tt.domain, queryRan) + } + }) + } +} + +func TestTPMVersion(t *testing.T) { + origGetTPMSpecVersion := getTPMSpecVersion + defer func() { getTPMSpecVersion = origGetTPMSpecVersion }() + + tests := []struct { + name string + mockVer string + mockErr error + want string + wantErr bool + }{ + { + name: "Success", + mockVer: "2.0, 0, 1.38", + want: "2.0, 0, 1.38", + }, + { + name: "Error", + mockErr: errors.New("TBS unavailable"), + want: "", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + getTPMSpecVersion = func() (string, error) { + return tt.mockVer, tt.mockErr + } + + got, err := TPMVersion() + if (err != nil) != tt.wantErr { + t.Fatalf("TPMVersion() error = %v, wantErr = %v", err, tt.wantErr) + } + if got != tt.want { + t.Errorf("TPMVersion() = %q, want %q", got, tt.want) + } + }) + } +} + func TestUserProfiles(t *testing.T) { root, err := ioutil.TempDir(os.TempDir(), "") if err != nil { t.Fatalf("ioutil.TempDir: %v", err) } + defer os.RemoveAll(root) + users := []string{"Administrator", "George", "Public"} for _, u := range users { - if err := os.MkdirAll(filepath.Join(root, "/Users/", u), 644); err != nil { + if err := os.MkdirAll(filepath.Join(root, "/Users/", u), 0755); err != nil { t.Fatalf("os.MkdirAll: %v", err) } } @@ -45,3 +425,13 @@ func TestUserProfiles(t *testing.T) { t.Errorf("UserProfiles() returned unexpected diff (-want +got):\n%s", diff) } } + +func TestUserProfiles_Error(t *testing.T) { + if err := os.Setenv("SystemDrive", `Z:\NonExistentDirectory_12345`); err != nil { + t.Fatalf("os.Setenv: %v", err) + } + _, err := UserProfiles() + if err == nil { + t.Errorf("UserProfiles() returned nil error on non-existent directory, want error") + } +} diff --git a/go/device/device_windows.go b/go/device/device_windows.go index 6e4132b6..3ab2b0a8 100644 --- a/go/device/device_windows.go +++ b/go/device/device_windows.go @@ -18,23 +18,44 @@ package device import ( + "bytes" + "encoding/binary" "errors" "fmt" "os" "strings" + "unsafe" + reg "golang.org/x/sys/windows/registry" + "golang.org/x/sys/windows" "github.com/StackExchange/wmi" ) var ( // ErrWMIEmptyResult indicates a condition where WMI failed to return the expected values. ErrWMIEmptyResult = errors.New("WMI returned without error, but zero results") + + kernelDLL = windows.NewLazySystemDLL("kernel32.dll") + procGetSystemFirmwareTable = kernelDLL.NewProc("GetSystemFirmwareTable") + + tbsDLL = windows.NewLazySystemDLL("tbs.dll") + tbsCreateContext = tbsDLL.NewProc("Tbsi_Context_Create") + tbsContextClose = tbsDLL.NewProc("Tbsip_Context_Close") + tbsSubmitCommand = tbsDLL.NewProc("Tbsip_Submit_Command") + + // Test helpers. + getChassisType = chassisTypeSMBIOS + getTPMSpecVersion = getTPMSpecVersionFromTBS + netGetJoinInformation = windows.NetGetJoinInformation + netAPIBufferFree = windows.NetApiBufferFree + registryGetString = regGetString + wmiQuery = wmi.Query ) -// Win32_SystemEnclosure models the WMI object of the same name. -type Win32_SystemEnclosure struct { - ChassisTypes []int -} +const ( + // https://github.com/digitalocean/go-smbios/blob/390a4f403a8e94ca0acdcf609eb18eeae9d6d1ac/smbios/stream_windows.go#L34 + providerRSMB = 0x52534D42 +) // Type is a device type as reported by the system enclosure. type Type string @@ -50,47 +71,62 @@ var ( Unknown Type = "Unknown" ) -// ChassisType attempts to distinguish the chassis type for the device. -func ChassisType() (Type, error) { - var result []Win32_SystemEnclosure - if err := wmi.Query(wmi.CreateQuery(&result, ""), &result); err != nil { - return Unknown, err - } - if len(result) < 1 || len(result[0].ChassisTypes) < 1 { - return Unknown, ErrWMIEmptyResult - } - switch result[0].ChassisTypes[0] { - case -1: - return Unknown, nil - case 3, 35: - return Desktop, nil - case 8, 9, 10, 11, 12, 14, 30, 31, 32: - return Laptop, nil - default: - return Other, nil +func chassisTypeSMBIOS() (Type, error) { + bufSize, _, err := procGetSystemFirmwareTable.Call(uintptr(providerRSMB), 0, 0, 0) + if bufSize == 0 { + return Unknown, fmt.Errorf("GetSystemFirmwareTable failed: %w", err) } + buf := make([]byte, bufSize) + returnCode, _, err := procGetSystemFirmwareTable.Call(uintptr(providerRSMB), 0, uintptr(unsafe.Pointer(&buf[0])), uintptr(bufSize)) + if returnCode == 0 { + return Unknown, fmt.Errorf("GetSystemFirmwareTable failed: %w", err) + } + if len(buf) < 8 { + return Unknown, errors.New("SMBIOS buffer too short") + } + tableData := buf[8:] + offset := 0 + for offset+4 <= len(tableData) { + structType := tableData[offset] + structLen := int(tableData[offset+1]) + if structLen < 4 || offset+structLen > len(tableData) { + break + } + if structType == 3 { // Type 3: System Enclosure or Chassis. + if structLen > 5 { + rawType := int(tableData[offset+5] & 0x7F) + switch rawType { + case 0: + return Unknown, nil + case 3, 35: + return Desktop, nil + case 8, 9, 10, 11, 12, 14, 30, 31, 32: + return Laptop, nil + default: + return Other, nil + } + } + } + // 127 is the end-of-table indicator. + if structType == 127 { + break + } + offset += structLen + for offset+1 < len(tableData) && (tableData[offset] != 0 || tableData[offset+1] != 0) { + offset++ + } + offset += 2 + } + return Unknown, nil } -// Win32_ComputerSystem models the WMI object of the same name. -type Win32_ComputerSystem struct { - DNSHostName string - Domain string - DomainRole int - Model string - SystemFamily string - Manufacturer string - PartOfDomain bool -} - -func sysInfo() (*Win32_ComputerSystem, error) { - var result []Win32_ComputerSystem - if err := wmi.Query(wmi.CreateQuery(&result, ""), &result); err != nil { - return nil, err - } - if len(result) < 1 { - return nil, ErrWMIEmptyResult +// ChassisType attempts to distinguish the chassis type for the device. +func ChassisType() (Type, error) { + t, err := getChassisType() + if err != nil { + return t, fmt.Errorf("failed to get chassis type from SMBIOS: %w", err) } - return &result[0], nil + return t, nil } // DomainRole indicates the role of a host on an Active Directory domain. @@ -107,18 +143,28 @@ var ( RoleUnknown DomainRole = "Unknown" ) +func regGetString(path, name string) (string, error) { + k, err := reg.OpenKey(reg.LOCAL_MACHINE, path, reg.READ) + if err != nil { + return "", err + } + defer k.Close() + val, _, err := k.GetStringValue(name) + return val, err +} + // GetDomainRole attempts to determine the host's Active Directory role. func GetDomainRole() (DomainRole, error) { - si, err := sysInfo() + productType, err := registryGetString(`SYSTEM\CurrentControlSet\Control\ProductOptions`, "ProductType") if err != nil { return RoleUnknown, err } - switch si.DomainRole { - case 0, 1: + switch productType { + case "WinNT": return Workstation, nil - case 2, 3: + case "ServerNT": return Server, nil - case 4, 5: + case "LanmanNT": return DomainController, nil default: return RoleUnknown, nil @@ -127,23 +173,35 @@ func GetDomainRole() (DomainRole, error) { // GetDomainJoined returns if a machine is joined to a domain. func GetDomainJoined() (bool, error) { - si, err := sysInfo() - if err != nil { - return false, err + var domainName *uint16 + var joinStatus uint32 + if err := netGetJoinInformation(nil, &domainName, &joinStatus); err != nil { + return false, fmt.Errorf("NetGetJoinInformation failed: %w", err) + } + if domainName != nil { + netAPIBufferFree((*byte)(unsafe.Pointer(domainName))) } - return si.PartOfDomain, nil + return joinStatus == windows.NetSetupDomainName, nil } // Model returns the system model. func Model() (string, error) { - si, err := sysInfo() + manufacturer, err := registryGetString(`HARDWARE\DESCRIPTION\System\BIOS`, "SystemManufacturer") if err != nil { - return "unknown", err + return "unknown", fmt.Errorf("failed to get SystemManufacturer from registry: %w", err) } - if strings.EqualFold(si.Manufacturer, "lenovo") { - return si.SystemFamily, nil + if strings.EqualFold(manufacturer, "lenovo") { + family, err := registryGetString(`HARDWARE\DESCRIPTION\System\BIOS`, "SystemFamily") + if err != nil { + return "unknown", fmt.Errorf("failed to get SystemFamily from registry: %w", err) + } + return family, nil } - return si.Model, nil + model, err := registryGetString(`HARDWARE\DESCRIPTION\System\BIOS`, "SystemProductName") + if err != nil { + return "unknown", fmt.Errorf("failed to get SystemProductName from registry: %w", err) + } + return model, nil } // Win32_NTDomain models the WMI object of the same name. @@ -155,7 +213,7 @@ type Win32_NTDomain struct { // Site returns the client's Active Directory site code. func Site(domain string) (string, error) { var result []Win32_NTDomain - err := wmi.Query(wmi.CreateQuery(&result, fmt.Sprintf("WHERE DomainName='%s'", domain)), &result) + err := wmiQuery(wmi.CreateQuery(&result, fmt.Sprintf("WHERE DomainName='%s'", domain)), &result) if err == nil && len(result) > 0 { return result[0].ClientSiteName, nil } @@ -164,7 +222,7 @@ func Site(domain string) (string, error) { // UserProfiles returns a list of user profiles on the local device. func UserProfiles() ([]string, error) { - users := []string{} + var users []string files, err := os.ReadDir(os.Getenv("SystemDrive") + `\Users`) if err != nil { return users, err @@ -178,20 +236,121 @@ func UserProfiles() ([]string, error) { return users, nil } -// Win32_Tpm models the WMI object of the same name. -type Win32_Tpm struct { - SpecVersion string +type tbsContextParams2 struct { + version uint32 + flags uint32 +} + +func getTPMSpecVersionFromTBS() (string, error) { + var tbsContext uintptr + params := tbsContextParams2{ + version: 2, // TPMVersion20 + flags: 6, // IncludeTPM12 (2) | IncludeTPM20 (4) + } + + returnCode, _, _ := tbsCreateContext.Call(uintptr(unsafe.Pointer(¶ms)), uintptr(unsafe.Pointer(&tbsContext))) + if returnCode != 0 { + return "", fmt.Errorf("Tbsi_Context_Create failed: %w", windows.Errno(returnCode)) + } + defer tbsContextClose.Call(tbsContext) + + // Send TPM2_GetCapability for TPM_PT_FAMILY_INDICATOR (0x100). + cmd := []byte{ + 0x80, 0x01, // TPM_ST_NO_SESSIONS + 0x00, 0x00, 0x00, 0x16, // Command Size = 22 bytes + 0x00, 0x00, 0x01, 0x7A, // TPM_CC_GetCapability + 0x00, 0x00, 0x00, 0x06, // TPM_CAP_TPM_PROPERTIES + 0x00, 0x00, 0x01, 0x00, // Property = TPM_PT_FAMILY_INDICATOR (0x100) + 0x00, 0x00, 0x00, 0x04, // PropertyCount = 4 + } + + resp := make([]byte, 1024) + respLen := uint32(len(resp)) + + const normalPriority = 200 + const commandLocalityZero = 0 + returnCode, _, _ = tbsSubmitCommand.Call( + tbsContext, + commandLocalityZero, + normalPriority, + uintptr(unsafe.Pointer(&cmd[0])), + uintptr(len(cmd)), + uintptr(unsafe.Pointer(&resp[0])), + uintptr(unsafe.Pointer(&respLen)), + ) + if returnCode == 0 && respLen >= 10 { + rc := binary.BigEndian.Uint32(resp[6:10]) + if rc == 0 && respLen >= 19 { + count := binary.BigEndian.Uint32(resp[15:19]) + var family string = "2.0" + var level uint32 = 0 + var rev uint32 = 0 + var foundRev bool + + offset := 19 + for i := uint32(0); i < count && offset+8 <= int(respLen); i++ { + prop := binary.BigEndian.Uint32(resp[offset : offset+4]) + val := binary.BigEndian.Uint32(resp[offset+4 : offset+8]) + switch prop { + case 0x100: // TPM_PT_FAMILY_INDICATOR (ASCII "2.0\0"). + b := resp[offset+4 : offset+8] + n := bytes.IndexByte(b, 0) + if n == -1 { + n = len(b) + } + family = string(b[:n]) + case 0x101: // TPM_PT_LEVEL + level = val + case 0x102: // TPM_PT_REVISION + rev = val + foundRev = true + } + offset += 8 + } + + if foundRev { + return fmt.Sprintf("%s, %d, %d.%02d", family, level, rev/100, rev%100), nil + } + } + } + + // Try TPM 1.2 GetCapability command if TPM 2.0 failed or on TPM 1.2. + cmd12 := []byte{ + 0x00, 0xC1, // TPM_TAG_RQU_COMMAND + 0x00, 0x00, 0x00, 0x12, // Command Size = 18 bytes + 0x00, 0x00, 0x00, 0x65, // TPM_ORD_GetCapability + 0x00, 0x00, 0x00, 0x1A, // TPM_CAP_VERSION_VAL + 0x00, 0x00, 0x00, 0x00, // SubCapSize = 0 + } + respLen = uint32(len(resp)) + returnCode, _, _ = tbsSubmitCommand.Call( + tbsContext, + commandLocalityZero, + normalPriority, + uintptr(unsafe.Pointer(&cmd12[0])), + uintptr(len(cmd12)), + uintptr(unsafe.Pointer(&resp[0])), + uintptr(unsafe.Pointer(&respLen)), + ) + if returnCode == 0 && respLen >= 20 { + rc := binary.BigEndian.Uint32(resp[6:10]) + if rc == 0 { + major := resp[16] + minor := resp[17] + revMajor := resp[18] + revMinor := resp[19] + return fmt.Sprintf("%d.%d, %d, %d.%d", major, minor, revMajor, revMinor, 0), nil + } + } + + return "", errors.New("unable to retrieve TPM spec version from TBS") } // TPMVersion returns the version of the TPM on the host. func TPMVersion() (string, error) { - var result []Win32_Tpm - query := "SELECT * FROM Win32_Tpm" - if err := wmi.QueryNamespace(query, &result, `root\CIMV2\Security\MicrosoftTpm`); err != nil { - return "", fmt.Errorf("WMI query for Win32_Tpm failed: %w", err) - } - if len(result) < 1 { - return "", nil + ver, err := getTPMSpecVersion() + if err != nil { + return "", err } - return result[0].SpecVersion, nil + return ver, nil }