Merge pull request #27 from projectdiscovery/issue-26-reusable-reader

Adding reusable reader

.github/workflows/build-test.yml

@@ -1,7 +1,6 @@
 name: 🔨 Build Test
 
 on:
-  push:
   pull_request:
   workflow_dispatch:
 

.github/workflows/lint-test.yml

@@ -2,7 +2,6 @@
 
 name: 🙏🏻 Lint Test
 on:
-  push:
   pull_request:
   workflow_dispatch:
 

file/frozen_reader.go → reader/frozen_reader.go

@@ -1,4 +1,4 @@
-package fileutil
+package reader
 
 import (
 	"io"

file/frozen_reader_test.go → reader/frozen_reader_test.go

@@ -1,4 +1,4 @@
-package fileutil
+package reader
 
 import (
 	"io"
@@ -12,7 +12,7 @@ func TestFrozenReader(t *testing.T) {
 		wrappedStdin := FrozenReader{}
 		_, err := io.Copy(os.Stdout, wrappedStdin)
 		if err != nil {
-			return 
+			return
 		}
 	}
 	go forever()

reader/length.go

@@ -0,0 +1,15 @@
+package reader
+
+import "io"
+
+// LenReader is an interface implemented by many in-memory io.Reader's. Used
+// for automatically sending the right Content-Length header when possible.
+type LenReader interface {
+	Len() int
+}
+
+// GetLength returns Length of reader using LenReader Interface
+func GetLength(reader io.Reader) (int64, bool) {
+	len, ok := reader.(LenReader)
+	return int64(len.Len()), ok
+}

reader/reusable_read_closer.go

@@ -0,0 +1,95 @@
+package reader
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"strings"
+)
+
+// ReusableReadCloser is a reusable reader with no-op close
+type ReusableReadCloser struct {
+	io.Reader
+	readBuf *bytes.Buffer
+	backBuf *bytes.Buffer
+}
+
+// NewReusableReadCloser is returned for any type of input
+func NewReusableReadCloser(raw interface{}) (*ReusableReadCloser, error) {
+	readBuf := bytes.Buffer{}
+	backBuf := bytes.Buffer{}
+	if raw != nil {
+		switch body := raw.(type) {
+
+		case []byte:
+			// if a byte array , create buffer from bytes and use it
+			readBuf = *bytes.NewBuffer(body)
+
+		case *[]byte:
+			// if *[]byte, create buffer from bytes and use it
+			readBuf = *bytes.NewBuffer(*body)
+
+		case string:
+			// if a string , create buffer from string and use it
+			readBuf = *bytes.NewBufferString(body)
+
+		case *bytes.Buffer:
+			// if *bytes.Buffer is given , use it
+			readBuf = *body
+
+		case *bytes.Reader:
+			// if *bytes.Reader , make buffer read from reader
+			if _, er := readBuf.ReadFrom(body); er != nil {
+				return nil, er
+			}
+
+		case *strings.Reader:
+			// if *strings.Reader , make buffer read from reader
+			if _, er := readBuf.ReadFrom(body); er != nil {
+				return nil, er
+			}
+
+		case io.ReadSeeker:
+			// if io.ReadSeeker , make buffer read from reader
+			if _, er := readBuf.ReadFrom(body); er != nil {
+				return nil, er
+			}
+
+		case io.Reader:
+			// if io.Reader , make buffer read from reader
+			if _, er := readBuf.ReadFrom(body); er != nil {
+				return nil, er
+			}
+		default:
+			// type not implemented or cannot handle
+			return nil, fmt.Errorf("cannot handle type %T", body)
+		}
+
+	}
+	reusableReadCloser := &ReusableReadCloser{
+		io.TeeReader(&readBuf, &backBuf),
+		&readBuf,
+		&backBuf,
+	}
+
+	return reusableReadCloser, nil
+
+}
+
+// Read []byte from Reader
+func (r ReusableReadCloser) Read(p []byte) (int, error) {
+	n, err := r.Reader.Read(p)
+	if err == io.EOF {
+		r.reset()
+	}
+	return n, err
+}
+
+func (r ReusableReadCloser) reset() {
+	_, _ = io.Copy(r.readBuf, r.backBuf)
+}
+
+// Close is a no-op close of ReusableReadCloser
+func (r ReusableReadCloser) Close() error {
+	return nil
+}

reader/reusable_read_closer_test.go

@@ -0,0 +1,31 @@
+package reader
+
+import (
+	"bytes"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestReusableReader(t *testing.T) {
+	testcases := []interface{}{
+		strings.NewReader("test"),
+		bytes.NewBuffer([]byte("test")),
+		bytes.NewBufferString("test"),
+		bytes.NewReader([]byte("test")),
+		[]byte("test"),
+		"test",
+	}
+	for _, v := range testcases {
+		reusableReader, err := NewReusableReadCloser(v)
+		require.Nil(t, err)
+
+		for i := 0; i < 100; i++ {
+			n, err := io.Copy(io.Discard, reusableReader)
+			require.Nil(t, err)
+			require.Positive(t, n)
+		}
+	}
+}

file/timeout_reader.go → reader/timeout_reader.go

@@ -1,4 +1,4 @@
-package fileutil
+package reader
 
 import (
 	"errors"

file/timeout_reader_test.go → reader/timeout_reader_test.go

@@ -1,4 +1,4 @@
-package fileutil
+package reader
 
 import (
 	"io"