feat: 增加对流式解析的支持

commons/stream.go

@@ -1,34 +1,39 @@
 package commons
 
 import (
+	"bytes"
 	"errors"
 	"io"
 )
 
-type Stream struct {
-	bs      []byte
-	current int64
+type CommonStream interface {
+	io.ReadSeeker
+	ReadN(n int) (bs []byte, err error)
+	PeekN(n int) (bs []byte, err error)
+	EOF() bool
+	CurrentIndex() int64
 }
 
-func NewStream(bs []byte) *Stream {
-	return &Stream{
-		bs:      bs,
-		current: int64(0),
-	}
+type Stream struct {
+	io.ReadSeeker
+	current int64
+	err     error
 }
 
-// Read implement io.Reader
 func (s *Stream) Read(b []byte) (n int, err error) {
-	if s.EOF() {
-		return 0, io.EOF
+	if _, err = s.Seek(s.current, io.SeekStart); err != nil {
+		s.err = err
+		return
+	}
+	n, err = s.ReadSeeker.Read(b)
+	if err != nil {
+		s.err = err
+		return
 	}
-	n = copy(b, s.bs[s.current:])
 	s.current += int64(n)
 	return
 }
 
-// ReadN read n bytes into byte array bs, it returns the byte array and any error encountered.
-// If read data size < n, an error will be returned with nil bs
 func (s *Stream) ReadN(n int) (bs []byte, err error) {
 	oldCurrent := s.current
 	bs = make([]byte, n)
@@ -37,43 +42,48 @@ func (s *Stream) ReadN(n int) (bs []byte, err error) {
 		s.current = oldCurrent
 		bs = nil
 	}
-
 	return
 }
 
-// PeekN read n bytes into byte array bs, it returns the byte array and any error encountered.
-// This method is similar as ReadN, but the stream pointer wouldn't move
 func (s *Stream) PeekN(n int) (bs []byte, err error) {
-	oldCurrent := s.current
+	oldOffset := s.current
 	bs, err = s.ReadN(n)
-	s.current = oldCurrent
+	s.current = oldOffset
 	return
 }
 
-// Seek implement io.Seeker
+func (s *Stream) EOF() bool {
+	return s.err != nil
+}
+
+func (s *Stream) CurrentIndex() int64 {
+	return s.current
+}
+
 func (s *Stream) Seek(offset int64, whence int) (int64, error) {
 	var abs int64
 	switch whence {
 	case io.SeekStart:
 		abs = offset
 	case io.SeekCurrent:
 		abs = s.current + offset
-	case io.SeekEnd:
-		abs = int64(len(s.bs)) + offset
 	default:
 		return 0, errors.New("bytes.Reader.Seek: invalid whence")
 	}
 	if abs < 0 {
 		return 0, errors.New("bytes.Reader.Seek: negative position")
 	}
 	s.current = abs
-	return abs, nil
+	return s.ReadSeeker.Seek(abs, whence)
 }
 
-func (s *Stream) EOF() bool {
-	return s.current >= int64(len(s.bs))
+func NewStreamFromReadSeeker(rs io.ReadSeeker) *Stream {
+	return &Stream{
+		ReadSeeker: rs,
+		current:    0,
+	}
 }
 
-func (s *Stream) CurrentIndex() int64 {
-	return s.current
+func NewStream(bs []byte) *Stream {
+	return NewStreamFromReadSeeker(bytes.NewReader(bs))
 }

serz/buffer.go

@@ -2,19 +2,28 @@ package serz
 
 import (
 	"github.com/phith0n/zkar/commons"
+	"io"
 )
 
 type ObjectStream struct {
-	*commons.Stream
+	commons.CommonStream
 	handler    uint32
 	references map[uint32]Object
 }
 
 func NewObjectStream(bs []byte) *ObjectStream {
 	return &ObjectStream{
-		Stream:     commons.NewStream(bs),
-		handler:    JAVA_BASE_WRITE_HANDLE,
-		references: make(map[uint32]Object),
+		CommonStream: commons.NewStream(bs),
+		handler:      JAVA_BASE_WRITE_HANDLE,
+		references:   make(map[uint32]Object),
+	}
+}
+
+func NewObjectStreamFromReadSeeker(r io.ReadSeeker) *ObjectStream {
+	return &ObjectStream{
+		CommonStream: commons.NewStreamFromReadSeeker(r),
+		handler:      JAVA_BASE_WRITE_HANDLE,
+		references:   make(map[uint32]Object),
 	}
 }
 

serz/parser.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"github.com/phith0n/zkar/commons"
+	"io"
 	"os"
 )
 
@@ -19,14 +20,12 @@ type Serialization struct {
 	Contents      []*TCContent
 }
 
-func FromBytes(data []byte) (*Serialization, error) {
-	var bs []byte
-	var err error
-	var stream = NewObjectStream(data)
+func FromReadSeeker(r io.ReadSeeker) (*Serialization, error) {
+	var stream = NewObjectStreamFromReadSeeker(r)
 	var ser = new(Serialization)
 
 	// read magic number 0xACED
-	bs, err = stream.ReadN(2)
+	bs, err := stream.ReadN(2)
 	if err != nil || !bytes.Equal(bs, JAVA_STREAM_MAGIC) {
 		return nil, fmt.Errorf("invalid magic number")
 	}
@@ -39,17 +38,22 @@ func FromBytes(data []byte) (*Serialization, error) {
 	}
 	ser.StreamVersion = bs
 
-	for !stream.EOF() {
+	for {
 		var content *TCContent
-		content, err = readTCContent(stream)
+		content, err = ReadTCContent(stream)
 		if err != nil {
+			if err == io.EOF {
+				return ser, nil
+			}
 			return nil, err
 		}
 
 		ser.Contents = append(ser.Contents, content)
 	}
+}
 
-	return ser, nil
+func FromBytes(data []byte) (*Serialization, error) {
+	return FromReadSeeker(bytes.NewReader(data))
 }
 
 func FromJDK8u20Bytes(data []byte) (*Serialization, error) {
@@ -62,6 +66,16 @@ func FromJDK8u20Bytes(data []byte) (*Serialization, error) {
 	return FromBytes(data)
 }
 
+func FromJDK8u20ReadSeeker(data []byte) (*Serialization, error) {
+	data = bytes.Replace(
+		data,
+		[]byte{0x00, 0x7e, 0x00, 0x09},
+		[]byte{0x00, 0x7e, 0x00, 0x09, JAVA_TC_ENDBLOCKDATA},
+		1,
+	)
+	return FromReadSeeker(bytes.NewReader(data))
+}
+
 func (ois *Serialization) ToString() string {
 	var b = commons.NewPrinter()
 	b.Printf("@Magic - %s", commons.Hexify(ois.MagicNumber))

serz/parser_test.go

@@ -27,6 +27,19 @@ func extractPackage(name string) string {
 	return name
 }
 
+func TestJDK8u20FromReadSeeker(t *testing.T) {
+	var filename = "../testcases/pwntester/JDK8u20.ser"
+	data, err := ioutil.ReadFile(filename)
+	require.Nil(t, err)
+
+	ser, err := FromJDK8u20ReadSeeker(data)
+	require.Nilf(t, err, "an error is occurred in file %v", filename)
+	serFromBytes, err := FromJDK8u20Bytes(data)
+	require.Equal(t, ser, serFromBytes)
+	require.Nilf(t, err, "an error is occurred in file %v", filename)
+	require.Truef(t, bytes.Equal(data, ser.ToJDK8u20Bytes()), "original data is different from generation data in file %v", filename)
+}
+
 func TestYsoserial(t *testing.T) {
 	walkAndTest("../testcases/ysoserial/*.ser", t, func(filename string, data []byte, ser *Serialization) {
 		require.Truef(

serz/tc_content.go

@@ -96,11 +96,14 @@ func (c *TCContent) Walk(callback WalkCallback) error {
 	return obj.Walk(callback)
 }
 
-func readTCContent(stream *ObjectStream) (*TCContent, error) {
+func ReadTCContent(stream *ObjectStream) (*TCContent, error) {
 	var err error = nil
 	var content = new(TCContent)
 
-	next, _ := stream.PeekN(1)
+	next, err := stream.PeekN(1)
+	if err != nil {
+		return nil, err
+	}
 	content.Flag = next[0]
 	switch next[0] {
 	case JAVA_TC_STRING, JAVA_TC_LONGSTRING:

serz/tc_normalclassdesc.go

@@ -197,7 +197,7 @@ func readTCAnnotation(stream *ObjectStream) ([]*TCContent, error) {
 			break
 		}
 
-		content, err := readTCContent(stream)
+		content, err := ReadTCContent(stream)
 		if err != nil {
 			return nil, err
 		}