path: root/vendor/github.com/StackExchange/wmi/swbemservices.go

// +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
}