Merge pull request #5 from phith0n/feature/walker

added walker for the Objects

Author: phith0n

example/main.go

@@ -1,12 +1,23 @@
 package main
 
 import (
-	"encoding/hex"
 	"fmt"
+	"github.com/phith0n/zkar/serz"
+	"io/ioutil"
+	"log"
 )
 
 func main() {
-	var s = "\x8c\xA3\x8B"
-	var bs = []byte(s)
-	fmt.Println(bs, s, hex.EncodeToString(bs))
+	data, _ := ioutil.ReadFile("./testcases/ysoserial/Jdk7u21.ser")
+	serialization, err := serz.FromBytes(data)
+	if err != nil {
+		log.Fatal("parse error")
+	}
+
+	desc := serz.FindClassDesc(serialization, "sun.reflect.annotation.AnnotationInvocationHandler")
+	if desc != nil {
+		fmt.Println(desc.ToString())
+	} else {
+		log.Fatal("class desc not found")
+	}
 }

serz/parser.go

@@ -10,6 +10,7 @@ import (
 type Object interface {
 	ToBytes() []byte
 	ToString() string
+	AllowWalked
 }
 
 type Serialization struct {
@@ -91,3 +92,17 @@ func (ois *Serialization) ToJDK8u20Bytes() []byte {
 		1,
 	)
 }
+
+func (ois *Serialization) Walk(callback WalkCallback) error {
+	for _, content := range ois.Contents {
+		if err := callback(content); err != nil {
+			return err
+		}
+
+		if err := content.Walk(callback); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

serz/parser_test.go

@@ -28,18 +28,14 @@ func extractPackage(name string) string {
 }
 
 func TestYsoserial(t *testing.T) {
-	files, err := filepath.Glob("../testcases/ysoserial/*.ser")
-	require.Nil(t, err)
-	require.NotZero(t, len(files))
-
-	for _, name := range files {
-		data, err := ioutil.ReadFile(name)
-		require.Nil(t, err)
-
-		ser, err := FromBytes(data)
-		require.Nilf(t, err, "an error is occurred in file %v", name)
-		require.Truef(t, bytes.Equal(data, ser.ToBytes()), "original serz data is different from generation data in file %v", name)
-	}
+	walkAndTest("../testcases/ysoserial/*.ser", t, func(filename string, data []byte, ser *Serialization) {
+		require.Truef(
+			t,
+			bytes.Equal(data, ser.ToBytes()),
+			"original data is different from generation data in file %v",
+			filename,
+		)
+	})
 }
 
 func TestJDK8u20(t *testing.T) {
@@ -49,7 +45,7 @@ func TestJDK8u20(t *testing.T) {
 
 	ser, err := FromJDK8u20Bytes(data)
 	require.Nilf(t, err, "an error is occurred in file %v", filename)
-	require.Truef(t, bytes.Equal(data, ser.ToJDK8u20Bytes()), "original serz data is different from generation data in file %v", filename)
+	require.Truef(t, bytes.Equal(data, ser.ToJDK8u20Bytes()), "original data is different from generation data in file %v", filename)
 }
 
 func TestMain(m *testing.M) {
@@ -113,3 +109,19 @@ func TestMain(m *testing.M) {
 cleanup:
 	os.Exit(exitCode)
 }
+
+func walkAndTest(pathGlob string, t *testing.T, callback func(filename string, data []byte, ser *Serialization)) {
+	files, err := filepath.Glob(pathGlob)
+	require.Nil(t, err)
+	require.NotZero(t, len(files))
+
+	for _, name := range files {
+		data, err := ioutil.ReadFile(name)
+		require.Nil(t, err)
+
+		ser, err := FromBytes(data)
+		require.Nilf(t, err, "an error is occurred in file %v", name)
+
+		callback(name, data, ser)
+	}
+}

serz/tc_array.go

@@ -58,6 +58,28 @@ func (t *TCArray) ToString() string {
 	return b.String()
 }
 
+func (t *TCArray) Walk(callback WalkCallback) error {
+	if err := callback(t.ClassPointer); err != nil {
+		return err
+	}
+
+	if err := t.ClassPointer.Walk(callback); err != nil {
+		return err
+	}
+
+	for _, v := range t.ArrayData {
+		if err := callback(v); err != nil {
+			return err
+		}
+
+		if err := v.Walk(callback); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func (t *TCArray) DumpByteArray() string {
 	var builder = &strings.Builder{}
 	var dumper = hex.Dumper(builder)

serz/tc_blockdata.go

@@ -35,6 +35,10 @@ func (bd *TCBlockData) ToString() string {
 	return b.String()
 }
 
+func (bd *TCBlockData) Walk(callback WalkCallback) error {
+	return nil
+}
+
 func readTCBlockData(stream *ObjectStream) (*TCBlockData, error) {
 	// read JAVA_TC_BLOCKDATA or JAVA_TC_BLOCKDATALONG Flag
 	flag, _ := stream.ReadN(1)

serz/tc_class.go

@@ -22,6 +22,14 @@ func (c *TCClass) ToString() string {
 	return b.String()
 }
 
+func (c *TCClass) Walk(callback WalkCallback) error {
+	if err := callback(c.ClassPointer); err != nil {
+		return err
+	}
+
+	return c.ClassPointer.Walk(callback)
+}
+
 func readTCClass(stream *ObjectStream) (*TCClass, error) {
 	var class = new(TCClass)
 	var err error

serz/tc_classdata.go

@@ -62,6 +62,34 @@ func (cd *TCClassData) ToString() string {
 	return b.String()
 }
 
+func (cd *TCClassData) Walk(callback WalkCallback) error {
+	for _, data := range cd.FieldDatas {
+		if err := callback(data); err != nil {
+			return err
+		}
+
+		if err := data.Walk(callback); err != nil {
+			return err
+		}
+	}
+
+	if !cd.HasAnnotation {
+		return nil
+	}
+
+	for _, anno := range cd.ObjectAnnotation {
+		if err := callback(anno); err != nil {
+			return err
+		}
+
+		if err := anno.Walk(callback); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func readTCClassData(stream *ObjectStream, desc *TCClassDesc) (*TCClassData, error) {
 	var err error
 	var classData = &TCClassData{

serz/tc_classpointer.go

@@ -48,6 +48,28 @@ func (cp *TCClassPointer) ToString() string {
 	return result
 }
 
+func (cp *TCClassPointer) Walk(callback WalkCallback) error {
+	var obj Object
+	switch cp.Flag {
+	case JAVA_TC_NULL:
+		obj = cp.Null
+	case JAVA_TC_REFERENCE:
+		obj = cp.Reference
+	case JAVA_TC_CLASSDESC:
+		obj = cp.NormalClassDesc
+	case JAVA_TC_PROXYCLASSDESC:
+		obj = cp.ProxyClassDesc
+	default:
+		panic("unexpected TCClassPointer Flag")
+	}
+
+	if err := callback(obj); err != nil {
+		return err
+	}
+
+	return obj.Walk(callback)
+}
+
 func (cp *TCClassPointer) FindClassBag(stream *ObjectStream) (*ClassBag, error) {
 	var normalClassDesc *TCClassDesc
 	var proxyClassDesc *TCProxyClassDesc

serz/tc_content.go

@@ -68,6 +68,34 @@ func (c *TCContent) ToString() string {
 	return bs
 }
 
+func (c *TCContent) Walk(callback WalkCallback) error {
+	var obj Object
+	switch c.Flag {
+	case JAVA_TC_STRING, JAVA_TC_LONGSTRING:
+		obj = c.String
+	case JAVA_TC_BLOCKDATA, JAVA_TC_BLOCKDATALONG:
+		obj = c.BlockData
+	case JAVA_TC_CLASS:
+		obj = c.Class
+	case JAVA_TC_OBJECT:
+		obj = c.Object
+	case JAVA_TC_NULL:
+		obj = c.Null
+	case JAVA_TC_REFERENCE:
+		obj = c.Reference
+	case JAVA_TC_ENUM:
+		obj = c.Enum
+	case JAVA_TC_ARRAY:
+		obj = c.Array
+	}
+
+	if err := callback(obj); err != nil {
+		return err
+	}
+
+	return obj.Walk(callback)
+}
+
 func readTCContent(stream *ObjectStream) (*TCContent, error) {
 	var err error = nil
 	var content = new(TCContent)

serz/tc_enum.go

@@ -25,6 +25,22 @@ func (e *TCEnum) ToString() string {
 	return b.String()
 }
 
+func (e *TCEnum) Walk(callback WalkCallback) error {
+	if err := callback(e.ClassPointer); err != nil {
+		return err
+	}
+
+	if err := e.ClassPointer.Walk(callback); err != nil {
+		return err
+	}
+
+	if err := callback(e.ConstantName); err != nil {
+		return err
+	}
+
+	return e.ConstantName.Walk(callback)
+}
+
 func readTCEnum(stream *ObjectStream) (*TCEnum, error) {
 	var enum = new(TCEnum)
 	var err error

serz/tc_fielddesc.go

@@ -54,6 +54,26 @@ func (f *TCFieldDesc) ToString() string {
 	return b.String()
 }
 
+func (f *TCFieldDesc) Walk(callback WalkCallback) error {
+	if err := callback(f.FieldName); err != nil {
+		return err
+	}
+
+	if err := f.FieldName.Walk(callback); err != nil {
+		return err
+	}
+
+	if f.TypeCode == "L" || f.TypeCode == "[" {
+		if err := callback(f.ClassName); err != nil {
+			return err
+		}
+
+		return f.ClassName.Walk(callback)
+	}
+
+	return nil
+}
+
 func readTCField(stream *ObjectStream) (*TCFieldDesc, error) {
 	var fieldDesc = new(TCFieldDesc)
 	typeCode, err := stream.ReadN(1)

serz/tc_normalclassdesc.go

@@ -74,6 +74,42 @@ func (desc *TCClassDesc) ToString() string {
 	return b.String()
 }
 
+func (desc *TCClassDesc) Walk(callback WalkCallback) error {
+	if err := callback(desc.ClassName); err != nil {
+		return err
+	}
+
+	if err := desc.ClassName.Walk(callback); err != nil {
+		return err
+	}
+
+	for _, field := range desc.Fields {
+		if err := callback(field); err != nil {
+			return err
+		}
+
+		if err := field.Walk(callback); err != nil {
+			return err
+		}
+	}
+
+	for _, anno := range desc.ClassAnnotation {
+		if err := callback(anno); err != nil {
+			return err
+		}
+
+		if err := anno.Walk(callback); err != nil {
+			return err
+		}
+	}
+
+	if err := callback(desc.SuperClassPointer); err != nil {
+		return err
+	}
+
+	return desc.SuperClassPointer.Walk(callback)
+}
+
 // HasFlag Check if a TCClassDesc object has a flag
 func (desc *TCClassDesc) HasFlag(flag byte) bool {
 	return (desc.ClassDescFlags & flag) == flag

serz/tc_null.go

@@ -16,6 +16,10 @@ func (n *TCNull) ToString() string {
 	return b.String()
 }
 
+func (n *TCNull) Walk(callback WalkCallback) error {
+	return nil
+}
+
 func readTCNull(stream *ObjectStream) *TCNull {
 	_, _ = stream.ReadN(1)
 	return new(TCNull)

serz/tc_object.go

@@ -33,6 +33,28 @@ func (oo *TCObject) ToString() string {
 	return b.String()
 }
 
+func (oo *TCObject) Walk(callback WalkCallback) error {
+	if err := callback(oo.ClassPointer); err != nil {
+		return err
+	}
+
+	if err := oo.ClassPointer.Walk(callback); err != nil {
+		return err
+	}
+
+	for _, data := range oo.ClassDatas {
+		if err := callback(data); err != nil {
+			return err
+		}
+
+		if err := data.Walk(callback); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func readTCObject(stream *ObjectStream) (*TCObject, error) {
 	var obj = new(TCObject)
 	var err error