Import MultiFileReader from go-ipfs-cmds

multifilereader.go

@@ -0,0 +1,124 @@
+package files
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/textproto"
+	"net/url"
+	"sync"
+)
+
+// MultiFileReader reads from a `commands.File` (which can be a directory of files
+// or a regular file) as HTTP multipart encoded data.
+type MultiFileReader struct {
+	io.Reader
+
+	files       []File
+	currentFile io.Reader
+	buf         bytes.Buffer
+	mpWriter    *multipart.Writer
+	closed      bool
+	mutex       *sync.Mutex
+
+	// if true, the data will be type 'multipart/form-data'
+	// if false, the data will be type 'multipart/mixed'
+	form bool
+}
+
+// NewMultiFileReader constructs a MultiFileReader. `file` can be any `commands.File`.
+// If `form` is set to true, the multipart data will have a Content-Type of 'multipart/form-data',
+// if `form` is false, the Content-Type will be 'multipart/mixed'.
+func NewMultiFileReader(file File, form bool) *MultiFileReader {
+	mfr := &MultiFileReader{
+		files: []File{file},
+		form:  form,
+		mutex: &sync.Mutex{},
+	}
+	mfr.mpWriter = multipart.NewWriter(&mfr.buf)
+
+	return mfr
+}
+
+func (mfr *MultiFileReader) Read(buf []byte) (written int, err error) {
+	mfr.mutex.Lock()
+	defer mfr.mutex.Unlock()
+
+	// if we are closed and the buffer is flushed, end reading
+	if mfr.closed && mfr.buf.Len() == 0 {
+		return 0, io.EOF
+	}
+
+	// if the current file isn't set, advance to the next file
+	if mfr.currentFile == nil {
+		var file File
+		for file == nil {
+			if len(mfr.files) == 0 {
+				mfr.mpWriter.Close()
+				mfr.closed = true
+				return mfr.buf.Read(buf)
+			}
+
+			nextfile, err := mfr.files[len(mfr.files)-1].NextFile()
+			if err == io.EOF {
+				mfr.files = mfr.files[:len(mfr.files)-1]
+				continue
+			} else if err != nil {
+				return 0, err
+			}
+
+			file = nextfile
+		}
+
+		// handle starting a new file part
+		if !mfr.closed {
+
+			var contentType string
+			if _, ok := file.(*Symlink); ok {
+				contentType = "application/symlink"
+			} else if file.IsDirectory() {
+				mfr.files = append(mfr.files, file)
+				contentType = "application/x-directory"
+			} else {
+				// otherwise, use the file as a reader to read its contents
+				contentType = "application/octet-stream"
+			}
+
+			mfr.currentFile = file
+
+			// write the boundary and headers
+			header := make(textproto.MIMEHeader)
+			filename := url.QueryEscape(file.FileName())
+			header.Set("Content-Disposition", fmt.Sprintf("file; filename=\"%s\"", filename))
+
+			header.Set("Content-Type", contentType)
+			if rf, ok := file.(*ReaderFile); ok {
+				header.Set("abspath", rf.AbsPath())
+			}
+
+			_, err := mfr.mpWriter.CreatePart(header)
+			if err != nil {
+				return 0, err
+			}
+		}
+	}
+
+	// if the buffer has something in it, read from it
+	if mfr.buf.Len() > 0 {
+		return mfr.buf.Read(buf)
+	}
+
+	// otherwise, read from file data
+	written, err = mfr.currentFile.Read(buf)
+	if err == io.EOF {
+		mfr.currentFile = nil
+		return written, nil
+	}
+	return written, err
+}
+
+// Boundary returns the boundary string to be used to separate files in the multipart data
+func (mfr *MultiFileReader) Boundary() string {
+	return mfr.mpWriter.Boundary()
+}

multifilereader_test.go

@@ -0,0 +1,112 @@
+package files
+
+import (
+	"io"
+	"io/ioutil"
+	"mime/multipart"
+	"strings"
+	"testing"
+)
+
+func TestOutput(t *testing.T) {
+	text := "Some text! :)"
+	fileset := []File{
+		NewReaderFile("file.txt", "file.txt", ioutil.NopCloser(strings.NewReader(text)), nil),
+		NewSliceFile("boop", "boop", []File{
+			NewReaderFile("boop/a.txt", "boop/a.txt", ioutil.NopCloser(strings.NewReader("bleep")), nil),
+			NewReaderFile("boop/b.txt", "boop/b.txt", ioutil.NopCloser(strings.NewReader("bloop")), nil),
+		}),
+		NewReaderFile("beep.txt", "beep.txt", ioutil.NopCloser(strings.NewReader("beep")), nil),
+	}
+	sf := NewSliceFile("", "", fileset)
+	buf := make([]byte, 20)
+
+	// testing output by reading it with the go stdlib "mime/multipart" Reader
+	mfr := NewMultiFileReader(sf, true)
+	mpReader := multipart.NewReader(mfr, mfr.Boundary())
+
+	part, err := mpReader.NextPart()
+	if part == nil || err != nil {
+		t.Fatal("Expected non-nil part, nil error")
+	}
+	mpf, err := NewFileFromPart(part)
+	if mpf == nil || err != nil {
+		t.Fatal("Expected non-nil MultipartFile, nil error")
+	}
+	if mpf.IsDirectory() {
+		t.Fatal("Expected file to not be a directory")
+	}
+	if mpf.FileName() != "file.txt" {
+		t.Fatal("Expected filename to be \"file.txt\"")
+	}
+	if n, err := mpf.Read(buf); n != len(text) || err != nil {
+		t.Fatal("Expected to read from file", n, err)
+	}
+	if string(buf[:len(text)]) != text {
+		t.Fatal("Data read was different than expected")
+	}
+
+	part, err = mpReader.NextPart()
+	if part == nil || err != nil {
+		t.Fatal("Expected non-nil part, nil error")
+	}
+	mpf, err = NewFileFromPart(part)
+	if mpf == nil || err != nil {
+		t.Fatal("Expected non-nil MultipartFile, nil error")
+	}
+	if !mpf.IsDirectory() {
+		t.Fatal("Expected file to be a directory")
+	}
+	if mpf.FileName() != "boop" {
+		t.Fatal("Expected filename to be \"boop\"")
+	}
+
+	part, err = mpReader.NextPart()
+	if part == nil || err != nil {
+		t.Fatal("Expected non-nil part, nil error")
+	}
+	child, err := NewFileFromPart(part)
+	if child == nil || err != nil {
+		t.Fatal("Expected to be able to read a child file")
+	}
+	if child.IsDirectory() {
+		t.Fatal("Expected file to not be a directory")
+	}
+	if child.FileName() != "boop/a.txt" {
+		t.Fatal("Expected filename to be \"some/file/path\"")
+	}
+
+	part, err = mpReader.NextPart()
+	if part == nil || err != nil {
+		t.Fatal("Expected non-nil part, nil error")
+	}
+	child, err = NewFileFromPart(part)
+	if child == nil || err != nil {
+		t.Fatal("Expected to be able to read a child file")
+	}
+	if child.IsDirectory() {
+		t.Fatal("Expected file to not be a directory")
+	}
+	if child.FileName() != "boop/b.txt" {
+		t.Fatal("Expected filename to be \"some/file/path\"")
+	}
+
+	child, err = mpf.NextFile()
+	if child != nil || err != io.EOF {
+		t.Fatal("Expected to get (nil, io.EOF)")
+	}
+
+	part, err = mpReader.NextPart()
+	if part == nil || err != nil {
+		t.Fatal("Expected non-nil part, nil error")
+	}
+	mpf, err = NewFileFromPart(part)
+	if mpf == nil || err != nil {
+		t.Fatal("Expected non-nil MultipartFile, nil error")
+	}
+
+	part, err = mpReader.NextPart()
+	if part != nil || err != io.EOF {
+		t.Fatal("Expected to get (nil, io.EOF)")
+	}
+}