msrpc: wmi: fix wmi parse object array

fix for smart enum implementation, where classes are preserved
across different wco smartenum calls and are allowed to go in
any order.

examples/samples_with_config/wmic.go

@@ -8,6 +8,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"runtime"
 	"time"
 
 	"github.com/oiweiwei/go-msrpc/dcerpc"
@@ -180,6 +181,8 @@ func main() {
 		flags |= wmi.QueryFlagType(wmi.GenericFlagTypeForwardOnly)
 	}
 
+	now := time.Now()
+
 	enum, err := svcs.ExecQuery(ctx, &iwbemservices.ExecQueryRequest{
 		This:          &dcom.ORPCThis{Version: srv.COMVersion},
 		QueryLanguage: &oaut.String{Data: "WQL"},
@@ -244,14 +247,15 @@ func main() {
 		return
 	}
 
-	var cls wmio.Class
-
-	now := time.Now()
-
 	if limit > 0 && limit < page {
 		page = limit
 	}
 
+	// classes should store the class definitions across the calls.
+	var classes = make(map[string]*wmio.Class)
+
+	count, buffer := 0, 0
+
 	for i := 0; limit == 0 || i < limit; i += page {
 
 		ret, err := wco.Next(ctx, &iwbemwcosmartenum.NextRequest{
@@ -270,7 +274,7 @@ func main() {
 			break
 		}
 
-		oa, err := wmi.UnmarshalObjectArrayWithClass(ret.Buffer, cls)
+		oa, err := wmi.UnmarshalObjectArrayWithClasses(ret.Buffer, classes)
 		if err != nil {
 			fmt.Fprintln(os.Stderr, "unmarshal_object_array_with_class", err)
 			return
@@ -280,13 +284,33 @@ func main() {
 			if po.Object.Class != nil {
 				fmt.Println(J(po.Object.Properties()))
 			} else {
-				cls = po.Object.Instance.CurrentClass
 				fmt.Println(J(po.Object.Values()))
 			}
 		}
+
+		count += len(oa.Objects)
+		buffer += len(ret.Buffer)
 	}
 
+	fmt.Fprintln(os.Stderr, "buffer fetched:", buffer/1024, "KiB")
+	fmt.Fprintln(os.Stderr, "records fetched:", count)
 	fmt.Fprintln(os.Stderr, "query execution time:", time.Now().Sub(now))
 	fmt.Fprintln(os.Stderr, "script execution time:", time.Now().Sub(startTime))
+	PrintMemUsage()
+
+}
+
+func PrintMemUsage() {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	// For info on each, see: https://golang.org/pkg/runtime/#MemStats
+	fmt.Fprintf(os.Stderr, "alloc / total / sys / numgc: %v MiB / %v MiB / %v MiB / %v\n",
+		bToMb(m.Alloc),
+		bToMb(m.TotalAlloc),
+		bToMb(m.Sys),
+		m.NumGC)
+}
 
+func bToMb(b uint64) uint64 {
+	return b / 1024 / 1024
 }

examples/wmic.go

@@ -192,6 +192,8 @@ func main() {
 		flags |= wmi.QueryFlagType(wmi.GenericFlagTypeForwardOnly)
 	}
 
+	now := time.Now()
+
 	enum, err := svcs.ExecQuery(ctx, &iwbemservices.ExecQueryRequest{
 		This:          &dcom.ORPCThis{Version: srv.COMVersion},
 		QueryLanguage: &oaut.String{Data: "WQL"},
@@ -256,9 +258,7 @@ func main() {
 		return
 	}
 
-	var cls wmio.Class
-
-	now := time.Now()
+	var classes = make(map[string]*wmio.Class)
 
 	if limit > 0 && limit < page {
 		page = limit
@@ -282,7 +282,7 @@ func main() {
 			break
 		}
 
-		oa, err := wmi.UnmarshalObjectArrayWithClass(ret.Buffer, cls)
+		oa, err := wmi.UnmarshalObjectArrayWithClasses(ret.Buffer, classes)
 		if err != nil {
 			fmt.Fprintln(os.Stderr, "unmarshal_object_array_with_class", err)
 			return
@@ -292,7 +292,6 @@ func main() {
 			if po.Object.Class != nil {
 				fmt.Println(j(po.Object.Properties()))
 			} else {
-				cls = po.Object.Instance.CurrentClass
 				fmt.Println(j(po.Object.Values()))
 			}
 		}

msrpc/dcom/wmi/object_array.go

@@ -13,6 +13,14 @@ func UnmarshalObjectArrayWithClass(b []byte, cls wmio.Class) (*ObjectArray, erro
 	return o, o.DecodeWithClass(wmio.NewCodec(b), cls)
 }
 
+// UnmarshalObjectArrayWithClasses decodes the ObjectArray structure
+// using the classes lookup map. Note that classes lookup map is updated
+// during the call and must be reused for the next object decoding.
+func UnmarshalObjectArrayWithClasses(b []byte, classes map[string]*wmio.Class) (*ObjectArray, error) {
+	o := &ObjectArray{}
+	return o, o.DecodeWithClasses(wmio.NewCodec(b), classes)
+}
+
 // The ObjectArray structure MUST be used to encode multiple CIM
 // objects that are returned in response to the IWbemWCOSmartEnum::Next
 // method. This structure is also used to encode parameters of the
@@ -102,7 +110,25 @@ type DataPacketObject struct {
 	Object *wmio.Object
 }
 
+var (
+	// Object is type WBEMOBJECT_CLASS:
+	// Structure contains the complete CIM Class definition.
+	ObjectTypeClass uint8 = 1
+	// Object is type WBEMOBJECT_INSTANCE:
+	// Structure contains the complete CIM Instance definition.
+	ObjectTypeInstance uint8 = 2
+	// Object is type WBEMOBJECT_INSTANCE_NOCLASS.
+	// Structure contains CIM Instance without the CIM Class definition.
+	ObjectTypeInstanceNoClass uint8 = 3
+)
+
+var DefaultClassID = ""
+
 func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error {
+	return o.DecodeWithClasses(r, map[string]*wmio.Class{DefaultClassID: &cls})
+}
+
+func (o *ObjectArray) DecodeWithClasses(r *wmio.Codec, classes map[string]*wmio.Class) error {
 
 	r.ReadData(&o.ByteOrdering)
 	o.Signature = make([]byte, 8)
@@ -120,6 +146,12 @@ func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error {
 
 	for i := 0; i < int(o.NumObjects); i++ {
 
+		var cls wmio.Class
+
+		if _, ok := classes[DefaultClassID]; ok {
+			cls = *classes[DefaultClassID]
+		}
+
 		var po DataPacketObject
 
 		r.ReadData(&po.SizeOfHeader)
@@ -136,12 +168,15 @@ func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error {
 		var err error
 
 		switch po.ObjectType {
-		case 1:
+		case ObjectTypeClass:
+
 			if po.Object, err = wmio.UnmarshalWithClass(r.Bytes()[:po.SizeOfData2], wmio.Class{}); err != nil {
 				return err
 			}
+
 			cls = po.Object.Class.CurrentClass
-		case 2:
+
+		case ObjectTypeInstance:
 
 			if po.ClassID, err = ReadClassID(r); err != nil {
 				return err
@@ -152,18 +187,28 @@ func (o *ObjectArray) DecodeWithClass(r *wmio.Codec, cls wmio.Class) error {
 			}
 
 			cls = po.Object.Instance.CurrentClass
-		case 3:
-			if cls.Name == "" {
-				return fmt.Errorf("class name is empty")
-			}
+
+			// store classes into map.
+			classes[po.ClassID.String()] = &cls
+
+		case ObjectTypeInstanceNoClass:
 
 			if po.ClassID, err = ReadClassID(r); err != nil {
 				return err
 			}
 
+			if _, ok := classes[po.ClassID.String()]; ok {
+				cls = *(classes[po.ClassID.String()])
+			}
+
+			if cls.Name == "" {
+				return fmt.Errorf("class name is empty")
+			}
+
 			if po.Object, err = wmio.UnmarshalWithClass(r.Bytes()[:po.SizeOfData2], cls); err != nil {
 				return err
 			}
+
 		default:
 			return fmt.Errorf("unknown class type %d", po.ObjectType)
 		}

msrpc/dcom/wmio/codec.go

@@ -157,6 +157,10 @@ func (c *Codec) DecodeWithSize32(sz uint32, data any) error {
 		return nil
 	}
 
+	if sz > uint32(c.buf.Len()) {
+		return c.Errf("decode_with_size32: expected size is greater than actual: %d > %d", sz, c.buf.Len())
+	}
+
 	b := make([]byte, sz)
 	c.ReadData(b)
 

msrpc/dcom/wmio/wmio.go

@@ -123,7 +123,18 @@ func (o *Object) Values() Values {
 	var values = make(Values)
 
 	for _, prop := range o.Instance.Properties {
-		values[prop.Name] = prop.Value.Value
+		switch value := prop.Value.Value.(type) {
+		case *Object:
+			values[prop.Name] = value.Values()
+		case []*Object:
+			vls := make([]any, len(value))
+			for i := range value {
+				vls[i] = value[i].Values()
+			}
+			values[prop.Name] = vls
+		default:
+			values[prop.Name] = value
+		}
 	}
 
 	return values