forked from StackExchange/wmi
-
Notifications
You must be signed in to change notification settings - Fork 26
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Split SWbemServices initialization out of Query method. Also add benc…
…hmark tests.
- Loading branch information
Showing
3 changed files
with
420 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,260 @@ | ||
// +build windows | ||
|
||
package wmi | ||
|
||
import ( | ||
"fmt" | ||
"reflect" | ||
"runtime" | ||
"sync" | ||
|
||
"github.com/go-ole/go-ole" | ||
"github.com/go-ole/go-ole/oleutil" | ||
) | ||
|
||
// SWbemServices is used to access wmi. See https://msdn.microsoft.com/en-us/library/aa393719(v=vs.85).aspx | ||
type SWbemServices struct { | ||
//TODO: track namespace. Not sure if we can re connect to a different namespace using the same instance | ||
cWMIClient *Client //This could also be an embedded struct, but then we would need to branch on Client vs SWbemServices in the Query method | ||
sWbemLocatorIUnknown *ole.IUnknown | ||
sWbemLocatorIDispatch *ole.IDispatch | ||
queries chan *queryRequest | ||
closeError chan error | ||
lQueryorClose sync.Mutex | ||
} | ||
|
||
type queryRequest struct { | ||
query string | ||
dst interface{} | ||
args []interface{} | ||
finished chan error | ||
} | ||
|
||
// InitializeSWbemServices will return a new SWbemServices object that can be used to query WMI | ||
func InitializeSWbemServices(c *Client, connectServerArgs ...interface{}) (*SWbemServices, error) { | ||
//fmt.Println("InitializeSWbemServices: Starting") | ||
//TODO: implement connectServerArgs as optional argument for init with connectServer call | ||
s := new(SWbemServices) | ||
s.cWMIClient = c | ||
s.queries = make(chan *queryRequest) | ||
initError := make(chan error) | ||
go s.process(initError) | ||
|
||
err, ok := <-initError | ||
if ok { | ||
return nil, err //Send error to caller | ||
} | ||
//fmt.Println("InitializeSWbemServices: Finished") | ||
return s, nil | ||
} | ||
|
||
// Close will clear and release all of the SWbemServices resources | ||
func (s *SWbemServices) Close() error { | ||
s.lQueryorClose.Lock() | ||
if s == nil || s.sWbemLocatorIDispatch == nil { | ||
s.lQueryorClose.Unlock() | ||
return fmt.Errorf("SWbemServices is not Initialized") | ||
} | ||
if s.queries == nil { | ||
s.lQueryorClose.Unlock() | ||
return fmt.Errorf("SWbemServices has been closed") | ||
} | ||
//fmt.Println("Close: sending close request") | ||
var result error | ||
ce := make(chan error) | ||
s.closeError = ce //Race condition if multiple callers to close. May need to lock here | ||
close(s.queries) //Tell background to shut things down | ||
s.lQueryorClose.Unlock() | ||
err, ok := <-ce | ||
if ok { | ||
result = err | ||
} | ||
//fmt.Println("Close: finished") | ||
return result | ||
} | ||
|
||
func (s *SWbemServices) process(initError chan error) { | ||
//fmt.Println("process: starting background thread initialization") | ||
//All OLE/WMI calls must happen on the same initialized thead, so lock this goroutine | ||
runtime.LockOSThread() | ||
defer runtime.LockOSThread() | ||
|
||
err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED) | ||
if err != nil { | ||
oleCode := err.(*ole.OleError).Code() | ||
if oleCode != ole.S_OK && oleCode != S_FALSE { | ||
initError <- fmt.Errorf("ole.CoInitializeEx error: %v", err) | ||
return | ||
} | ||
} | ||
defer ole.CoUninitialize() | ||
|
||
unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator") | ||
if err != nil { | ||
initError <- fmt.Errorf("CreateObject SWbemLocator error: %v", err) | ||
return | ||
} else if unknown == nil { | ||
initError <- ErrNilCreateObject | ||
return | ||
} | ||
defer unknown.Release() | ||
s.sWbemLocatorIUnknown = unknown | ||
|
||
dispatch, err := s.sWbemLocatorIUnknown.QueryInterface(ole.IID_IDispatch) | ||
if err != nil { | ||
initError <- fmt.Errorf("SWbemLocator QueryInterface error: %v", err) | ||
return | ||
} | ||
defer dispatch.Release() | ||
s.sWbemLocatorIDispatch = dispatch | ||
|
||
// we can't do the ConnectServer call outside the loop unless we find a way to track and re-init the connectServerArgs | ||
//fmt.Println("process: initialized. closing initError") | ||
close(initError) | ||
//fmt.Println("process: waiting for queries") | ||
for q := range s.queries { | ||
//fmt.Printf("process: new query: len(query)=%d\n", len(q.query)) | ||
errQuery := s.queryBackground(q) | ||
//fmt.Println("process: s.queryBackground finished") | ||
if errQuery != nil { | ||
q.finished <- errQuery | ||
} | ||
close(q.finished) | ||
} | ||
//fmt.Println("process: queries channel closed") | ||
s.queries = nil //set channel to nil so we know it is closed | ||
//TODO: I think the Release/Clear calls can panic if things are in a bad state. | ||
//TODO: May need to recover from panics and send error to method caller instead. | ||
close(s.closeError) | ||
} | ||
|
||
// Query runs the WQL query using a SWbemServices instance and appends the values to dst. | ||
// | ||
// dst must have type *[]S or *[]*S, for some struct type S. Fields selected in | ||
// the query must have the same name in dst. Supported types are all signed and | ||
// unsigned integers, time.Time, string, bool, or a pointer to one of those. | ||
// Array types are not supported. | ||
// | ||
// By default, the local machine and default namespace are used. These can be | ||
// changed using connectServerArgs. See | ||
// http://msdn.microsoft.com/en-us/library/aa393720.aspx for details. | ||
func (s *SWbemServices) Query(query string, dst interface{}, connectServerArgs ...interface{}) error { | ||
s.lQueryorClose.Lock() | ||
if s == nil || s.sWbemLocatorIDispatch == nil { | ||
s.lQueryorClose.Unlock() | ||
return fmt.Errorf("SWbemServices is not Initialized") | ||
} | ||
if s.queries == nil { | ||
s.lQueryorClose.Unlock() | ||
return fmt.Errorf("SWbemServices has been closed") | ||
} | ||
|
||
//fmt.Println("Query: Sending query request") | ||
qr := queryRequest{ | ||
query: query, | ||
dst: dst, | ||
args: connectServerArgs, | ||
finished: make(chan error), | ||
} | ||
s.queries <- &qr | ||
s.lQueryorClose.Unlock() | ||
err, ok := <-qr.finished | ||
if ok { | ||
//fmt.Println("Query: Finished with error") | ||
return err //Send error to caller | ||
} | ||
//fmt.Println("Query: Finished") | ||
return nil | ||
} | ||
|
||
func (s *SWbemServices) queryBackground(q *queryRequest) error { | ||
if s == nil || s.sWbemLocatorIDispatch == nil { | ||
return fmt.Errorf("SWbemServices is not Initialized") | ||
} | ||
wmi := s.sWbemLocatorIDispatch //Should just rename in the code, but this will help as we break things apart | ||
//fmt.Println("queryBackground: Starting") | ||
|
||
dv := reflect.ValueOf(q.dst) | ||
if dv.Kind() != reflect.Ptr || dv.IsNil() { | ||
return ErrInvalidEntityType | ||
} | ||
dv = dv.Elem() | ||
mat, elemType := checkMultiArg(dv) | ||
if mat == multiArgTypeInvalid { | ||
return ErrInvalidEntityType | ||
} | ||
|
||
// service is a SWbemServices | ||
serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer", q.args...) | ||
if err != nil { | ||
return err | ||
} | ||
service := serviceRaw.ToIDispatch() | ||
defer serviceRaw.Clear() | ||
|
||
// result is a SWBemObjectSet | ||
resultRaw, err := oleutil.CallMethod(service, "ExecQuery", q.query) | ||
if err != nil { | ||
return err | ||
} | ||
result := resultRaw.ToIDispatch() | ||
defer resultRaw.Clear() | ||
|
||
count, err := oleInt64(result, "Count") | ||
if err != nil { | ||
return err | ||
} | ||
|
||
enumProperty, err := result.GetProperty("_NewEnum") | ||
if err != nil { | ||
return err | ||
} | ||
defer enumProperty.Clear() | ||
|
||
enum, err := enumProperty.ToIUnknown().IEnumVARIANT(ole.IID_IEnumVariant) | ||
if err != nil { | ||
return err | ||
} | ||
if enum == nil { | ||
return fmt.Errorf("can't get IEnumVARIANT, enum is nil") | ||
} | ||
defer enum.Release() | ||
|
||
// Initialize a slice with Count capacity | ||
dv.Set(reflect.MakeSlice(dv.Type(), 0, int(count))) | ||
|
||
var errFieldMismatch error | ||
for itemRaw, length, err := enum.Next(1); length > 0; itemRaw, length, err = enum.Next(1) { | ||
if err != nil { | ||
return err | ||
} | ||
|
||
err := func() error { | ||
// item is a SWbemObject, but really a Win32_Process | ||
item := itemRaw.ToIDispatch() | ||
defer item.Release() | ||
|
||
ev := reflect.New(elemType) | ||
if err = s.cWMIClient.loadEntity(ev.Interface(), item); err != nil { | ||
if _, ok := err.(*ErrFieldMismatch); ok { | ||
// We continue loading entities even in the face of field mismatch errors. | ||
// If we encounter any other error, that other error is returned. Otherwise, | ||
// an ErrFieldMismatch is returned. | ||
errFieldMismatch = err | ||
} else { | ||
return err | ||
} | ||
} | ||
if mat != multiArgTypeStructPtr { | ||
ev = ev.Elem() | ||
} | ||
dv.Set(reflect.Append(dv, ev)) | ||
return nil | ||
}() | ||
if err != nil { | ||
return err | ||
} | ||
} | ||
//fmt.Println("queryBackground: Finished") | ||
return errFieldMismatch | ||
} |
Oops, something went wrong.