Writer struct allowing io.Copy or io.TeeReader from a reader

filetype_test.go

@@ -1,7 +1,11 @@
 package filetype
 
 import (
+	"bytes"
+	"io"
+	"math/rand"
 	"testing"
+	"time"
 
 	"github.com/h2non/filetype/types"
 )
@@ -121,3 +125,65 @@ func TestGetType(t *testing.T) {
 		t.Fatalf("Type should not be supported")
 	}
 }
+
+func TestMatchWriter(t *testing.T) {
+	cases := []struct {
+		mime   string
+		reader io.Reader
+		read   int64
+		err    error
+	}{
+		{"image/jpeg", bytes.NewReader([]byte{0xFF, 0xD8, 0xFF}), 3, nil},
+		{"image/png", bytes.NewReader([]byte{0x89, 0x50, 0x4E, 0x47}), 4, nil},
+		{types.Unknown.MIME.Value, bytes.NewReader([]byte{}), 0, ErrEmptyBuffer},
+		{types.Unknown.MIME.Value, nil, 0, ErrEmptyBuffer},
+	}
+	for _, test := range cases {
+		var w int64
+		var err error
+		var mimeType types.Type
+		mw := NewMatcherWriter()
+		if test.reader != nil {
+			w, err = io.Copy(mw, test.reader)
+		}
+		if err != nil {
+			t.Fatalf("Error matching %s error: %v", test.mime, err)
+		}
+		mimeType, err = mw.Match()
+		if err != test.err {
+			t.Fatalf("Invalid error match: %v, expected %s", err, test.err)
+		}
+		if mimeType.MIME.Value != test.mime {
+			t.Fatalf("Invalid mime match: %s, expected %s", mimeType.MIME.Value, test.mime)
+		}
+		if w != test.read {
+			t.Fatalf("Invalid read match: %d, expected %d", w, test.read)
+		}
+	}
+}
+
+func generateRandomSlice(size int) []byte {
+	slice := make([]byte, size, size)
+	rand.Seed(time.Now().UnixNano())
+	for i := 0; i < size; i++ {
+		slice[i] = byte(rand.Intn(999))
+	}
+	return slice
+}
+
+func TestMatchWriterBuffer(t *testing.T) {
+	randomSlice := generateRandomSlice(maxBufSize + 500)
+	reader := bytes.NewReader(randomSlice)
+	mw := NewMatcherWriter()
+	_, err := io.Copy(mw, reader)
+	if err != nil {
+		t.Fatalf("error copying bytes to reader %v", err)
+	}
+	bufLen := len(mw.buf)
+	if bufLen != maxBufSize {
+		t.Fatalf("expected buffer len to be %d but is %d", maxBufSize, bufLen)
+	}
+	if bytes.Compare(mw.buf, randomSlice[0:maxBufSize]) != 0 {
+		t.Fatalf("expected buffer to equal re sliced buffer")
+	}
+}

match.go

@@ -17,6 +17,46 @@ var MatcherKeys = &matchers.MatcherKeys
 // NewMatcher is an alias to matchers.NewMatcher
 var NewMatcher = matchers.NewMatcher
 
+const maxBufSize = 8192
+
+// MatcherWriter is a matcher that coplies to the writer interface
+type MatcherWriter struct {
+	buf []byte
+}
+
+func (mb *MatcherWriter) Write(p []byte) (n int, err error) {
+	incomingSize := len(p)
+	currentSize := len(mb.buf)
+	if currentSize < maxBufSize {
+		// write only when the current size of the buffer is less than the max buffer size
+		newSize := currentSize + incomingSize
+		overflow := newSize - maxBufSize
+		reSlice := p
+		if overflow > 0 {
+			// if the maxBufSize is exceeded by the new size, we need to do some re slicing
+			maxLen := incomingSize - overflow
+			reSlice = p[0:maxLen]
+		}
+		mb.buf = append(mb.buf, reSlice...)
+	}
+	return incomingSize, nil
+}
+
+// Match calls the Match function with the inner buffer of the MatcherWriter
+func (mb *MatcherWriter) Match() (types.Type, error) {
+	if mb.buf == nil {
+		return types.Unknown, ErrEmptyBuffer
+	}
+	return Match(mb.buf)
+}
+
+// NewMatcherWriter creates a matcher which is a io.Writer
+func NewMatcherWriter() *MatcherWriter {
+	return &MatcherWriter{
+		buf: make([]byte, 0, maxBufSize),
+	}
+}
+
 // Match infers the file type of a given buffer inspecting its magic numbers signature
 func Match(buf []byte) (types.Type, error) {
 	length := len(buf)