Merge pull request #15 from phith0n/reader-support

Support parse io.Reader

README.md

@@ -33,8 +33,9 @@ import (
 )
 
 func main() {
-  data, _ := os.ReadFile("./testcases/ysoserial/CommonsCollections6.ser")
-  serialization, err := serz.FromBytes(data)
+  fs, _ := os.Open("./testcases/ysoserial/CommonsCollections6.ser")
+  defer fs.Close()
+  serialization, err := serz.FromReader(fs)
   if err != nil {
     log.Fatal("parse error")
   }

commons/stream.go

@@ -6,38 +6,55 @@ import (
 )
 
 type Stream struct {
-	io.ReadSeeker
+	reader io.Reader
+
+	buf   []byte
+	index int
 }
 
-func (s *Stream) ReadN(n int) (bs []byte, err error) {
-	bs = make([]byte, n)
-	_, err = io.ReadFull(s, bs)
-	return
+func (s *Stream) ReadN(n int) ([]byte, error) {
+	if n > len(s.buf)-s.index {
+		buf := make([]byte, n-(len(s.buf)-s.index))
+		read, err := io.ReadFull(s.reader, buf)
+		s.buf = append(s.buf, buf[0:read]...)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	start := s.index
+	s.index += n
+	return s.buf[start:s.index], nil
 }
 
-func (s *Stream) CurrentIndex() int64 {
-	n, _ := s.Seek(0, io.SeekCurrent)
-	return n
+func (s *Stream) CurrentIndex() int {
+	return s.index
 }
 
-func (s *Stream) PeekN(n int) (bs []byte, err error) {
-	bs, err = s.ReadN(n)
-	if err != nil {
-		return
-	}
-	_, err = s.Seek(int64(-n), io.SeekCurrent)
-	if err != nil {
-		return
-	}
-	return
+func (s *Stream) PeekN(n int) ([]byte, error) {
+	current := s.index
+	defer func() {
+		s.index = current
+	}()
+	return s.ReadN(n)
 }
 
-func NewStreamFromReadSeeker(rs io.ReadSeeker) *Stream {
+func (s *Stream) Seek(index int) {
+	s.index = index
+}
+
+func NewStreamFromReader(reader io.Reader) *Stream {
 	return &Stream{
-		ReadSeeker: rs,
+		reader: reader,
+		buf:    make([]byte, 0),
+		index:  0,
 	}
 }
 
+func NewStreamFromReadSeeker(rs io.ReadSeeker) *Stream {
+	return NewStreamFromReader(rs)
+}
+
 func NewStream(bs []byte) *Stream {
-	return NewStreamFromReadSeeker(bytes.NewReader(bs))
+	return NewStreamFromReader(bytes.NewReader(bs))
 }

commons/stream_test.go

@@ -1,34 +1,13 @@
 package commons
 
 import (
+	"bytes"
 	"io"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
-func TestStream_Read(t *testing.T) {
-	var bs []byte
-	var err error
-	var n int
-	s := NewStream([]byte("1111122222333334444455555"))
-	bs = make([]byte, 5)
-	n, err = s.Read(bs)
-	require.Equal(t, 5, n)
-	require.Nil(t, err)
-	require.Equal(t, []byte("11111"), bs)
-	bs = make([]byte, 22)
-	n, err = s.Read(bs)
-	require.Equal(t, 20, n)
-	require.Nil(t, nil, err)
-	require.Equal(t, []byte("22222333334444455555\x00\x00"), bs)
-	bs = make([]byte, 5)
-	n, err = s.Read(bs)
-	require.Equal(t, 0, n)
-	require.Equal(t, io.EOF, err)
-	require.Equal(t, []byte("\x00\x00\x00\x00\x00"), bs)
-}
-
 func TestStream_ReadN(t *testing.T) {
 	var bs []byte
 	var err error
@@ -58,10 +37,94 @@ func TestStream_CurrentIndex(t *testing.T) {
 	bs, err = s.PeekN(5)
 	require.Nil(t, err)
 	require.Equal(t, []byte("11111"), bs)
-	require.Equal(t, int64(0), s.CurrentIndex())
+	require.Equal(t, 0, s.CurrentIndex())
 
 	bs, err = s.ReadN(5)
 	require.Nil(t, err)
 	require.Equal(t, []byte("11111"), bs)
-	require.Equal(t, int64(5), s.CurrentIndex())
+	require.Equal(t, 5, s.CurrentIndex())
+}
+
+func TestStreamReader(t *testing.T) {
+	var s *Stream
+	var data []byte
+	var err error
+
+	s = NewStreamFromReader(bytes.NewReader([]byte("abbcccddddeeeee")))
+	data, err = s.ReadN(1)
+	require.NoError(t, err)
+	require.Equal(t, []byte("a"), data)
+	data, err = s.ReadN(5)
+	require.NoError(t, err)
+	require.Equal(t, []byte("bbccc"), data)
+	data, err = s.PeekN(3)
+	require.NoError(t, err)
+	require.Equal(t, []byte("ddd"), data)
+	data, err = s.PeekN(2)
+	require.NoError(t, err)
+	require.Equal(t, []byte("dd"), data)
+	data, err = s.PeekN(5)
+	require.NoError(t, err)
+	require.Equal(t, []byte("dddde"), data)
+	data, err = s.ReadN(4)
+	require.NoError(t, err)
+	require.Equal(t, []byte("dddd"), data)
+	data, err = s.ReadN(1)
+	require.NoError(t, err)
+	require.Equal(t, []byte("e"), data)
+	data, err = s.ReadN(3)
+	require.NoError(t, err)
+	require.Equal(t, []byte("eee"), data)
+	_, err = s.PeekN(4)
+	require.Error(t, err)
+	data, err = s.PeekN(1)
+	require.NoError(t, err)
+	require.Equal(t, []byte("e"), data)
+	_, err = s.ReadN(2)
+	require.Error(t, err)
+	data, err = s.ReadN(1)
+	require.NoError(t, err)
+	require.Equal(t, []byte("e"), data)
+	_, err = s.ReadN(1)
+	require.Error(t, err)
+}
+
+func TestEmptyStream(t *testing.T) {
+	var s *Stream
+	var data []byte
+	var err error
+
+	s = NewStreamFromReader(bytes.NewReader([]byte{}))
+	data, err = s.ReadN(2)
+	require.Error(t, err)
+	require.Nil(t, data)
+
+	_, err = s.PeekN(1)
+	require.Error(t, err)
+
+	s = NewStreamFromReader(bytes.NewReader(nil))
+	_, err = s.ReadN(1)
+	require.Error(t, err)
+}
+
+func TestStreamBuf(t *testing.T) {
+	var s *Stream
+	var data []byte
+	var err error
+
+	s = NewStreamFromReader(bytes.NewReader([]byte("abbcccddddeeeee")))
+	_, err = s.PeekN(20)
+	require.Error(t, err)
+
+	data, err = s.PeekN(3)
+	require.NoError(t, err)
+	require.Equal(t, []byte("abb"), data)
+
+	data, err = s.ReadN(5)
+	require.NoError(t, err)
+	require.Equal(t, []byte("abbcc"), data)
+
+	data, err = s.PeekN(10)
+	require.NoError(t, err)
+	require.Equal(t, []byte("cddddeeeee"), data)
 }

example/main.go

@@ -3,13 +3,18 @@ package main
 import (
 	"fmt"
 	"github.com/phith0n/zkar/serz"
-	"io/ioutil"
 	"log"
+	"os"
 )
 
 func main() {
-	data, _ := ioutil.ReadFile("./testcases/ysoserial/Jdk7u21.ser")
-	serialization, err := serz.FromBytes(data)
+	fs, err := os.Open("./testcases/ysoserial/Jdk7u21.ser")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer fs.Close()
+
+	serialization, err := serz.FromReader(fs)
 	if err != nil {
 		log.Fatal("parse error")
 	}

main.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"github.com/phith0n/zkar/serz"
 	"github.com/urfave/cli/v2"
-	"io/ioutil"
 	"log"
 	"os"
 )
@@ -79,7 +78,7 @@ func main() {
 					}
 
 					if filename != "" {
-						data, err = ioutil.ReadFile(filename)
+						data, err = os.ReadFile(filename)
 					} else {
 						data, err = base64.StdEncoding.DecodeString(b64data)
 					}

serz/buffer.go

@@ -20,9 +20,9 @@ func NewObjectStream(bs []byte) *ObjectStream {
 	}
 }
 
-func NewObjectStreamFromReadSeeker(r io.ReadSeeker) *ObjectStream {
+func NewObjectStreamFromReader(r io.Reader) *ObjectStream {
 	return &ObjectStream{
-		Stream:     commons.NewStreamFromReadSeeker(r),
+		Stream:     commons.NewStreamFromReader(r),
 		handler:    JAVA_BASE_WRITE_HANDLE,
 		references: make(map[uint32]Object),
 	}

serz/parser.go

@@ -21,8 +21,8 @@ type Serialization struct {
 	Contents      []*TCContent
 }
 
-func FromReadSeeker(r io.ReadSeeker) (*Serialization, error) {
-	var stream = NewObjectStreamFromReadSeeker(r)
+func FromReader(r io.Reader) (*Serialization, error) {
+	var stream = NewObjectStreamFromReader(r)
 	var ser = new(Serialization)
 
 	// read magic number 0xACED
@@ -53,8 +53,12 @@ func FromReadSeeker(r io.ReadSeeker) (*Serialization, error) {
 	}
 }
 
+func FromReadSeeker(r io.ReadSeeker) (*Serialization, error) {
+	return FromReader(r)
+}
+
 func FromBytes(data []byte) (*Serialization, error) {
-	return FromReadSeeker(bytes.NewReader(data))
+	return FromReader(bytes.NewReader(data))
 }
 
 func FromJDK8u20Bytes(data []byte) (*Serialization, error) {

serz/tc_classdata.go

@@ -2,7 +2,6 @@ package serz
 
 import (
 	"github.com/phith0n/zkar/commons"
-	"io"
 )
 
 type ReferenceClassInformation struct {
@@ -107,7 +106,7 @@ func readTCClassData(stream *ObjectStream, desc *TCClassDesc) (*TCClassData, err
 				// So we should clear the classData.FieldDatas and reset the position of stream
 				// Then everything will be read from objectAnnotation
 				// Example: ysoserial C3O0
-				_, _ = stream.Seek(current, io.SeekStart)
+				stream.Seek(current)
 				classData.FieldDatas = []*TCValue{}
 				break
 			} else if err != nil {