github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/commands/http/multifilereader.go (about)

     1  package http
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"mime/multipart"
     8  	"net/textproto"
     9  	"net/url"
    10  	"sync"
    11  
    12  	files "github.com/ipfs/go-ipfs/commands/files"
    13  )
    14  
    15  // MultiFileReader reads from a `commands.File` (which can be a directory of files
    16  // or a regular file) as HTTP multipart encoded data.
    17  type MultiFileReader struct {
    18  	io.Reader
    19  
    20  	files       files.File
    21  	currentFile io.Reader
    22  	buf         bytes.Buffer
    23  	mpWriter    *multipart.Writer
    24  	closed      bool
    25  	mutex       *sync.Mutex
    26  
    27  	// if true, the data will be type 'multipart/form-data'
    28  	// if false, the data will be type 'multipart/mixed'
    29  	form bool
    30  }
    31  
    32  // NewMultiFileReader constructs a MultiFileReader. `file` can be any `commands.File`.
    33  // If `form` is set to true, the multipart data will have a Content-Type of 'multipart/form-data',
    34  // if `form` is false, the Content-Type will be 'multipart/mixed'.
    35  func NewMultiFileReader(file files.File, form bool) *MultiFileReader {
    36  	mfr := &MultiFileReader{
    37  		files: file,
    38  		form:  form,
    39  		mutex: &sync.Mutex{},
    40  	}
    41  	mfr.mpWriter = multipart.NewWriter(&mfr.buf)
    42  
    43  	return mfr
    44  }
    45  
    46  func (mfr *MultiFileReader) Read(buf []byte) (written int, err error) {
    47  	mfr.mutex.Lock()
    48  	defer mfr.mutex.Unlock()
    49  
    50  	// if we are closed and the buffer is flushed, end reading
    51  	if mfr.closed && mfr.buf.Len() == 0 {
    52  		return 0, io.EOF
    53  	}
    54  
    55  	// if the current file isn't set, advance to the next file
    56  	if mfr.currentFile == nil {
    57  		file, err := mfr.files.NextFile()
    58  		if err == io.EOF {
    59  			mfr.mpWriter.Close()
    60  			mfr.closed = true
    61  		} else if err != nil {
    62  			return 0, err
    63  		}
    64  
    65  		// handle starting a new file part
    66  		if !mfr.closed {
    67  
    68  			var contentType string
    69  			if s, ok := file.(*files.Symlink); ok {
    70  				mfr.currentFile = s
    71  
    72  				contentType = "application/symlink"
    73  			} else if file.IsDirectory() {
    74  				// if file is a directory, create a multifilereader from it
    75  				// (using 'multipart/mixed')
    76  				nmfr := NewMultiFileReader(file, false)
    77  				mfr.currentFile = nmfr
    78  				contentType = fmt.Sprintf("multipart/mixed; boundary=%s", nmfr.Boundary())
    79  			} else {
    80  				// otherwise, use the file as a reader to read its contents
    81  				mfr.currentFile = file
    82  				contentType = "application/octet-stream"
    83  			}
    84  
    85  			// write the boundary and headers
    86  			header := make(textproto.MIMEHeader)
    87  			filename := url.QueryEscape(file.FileName())
    88  			if mfr.form {
    89  				contentDisposition := fmt.Sprintf("form-data; name=\"file\"; filename=\"%s\"", filename)
    90  				header.Set("Content-Disposition", contentDisposition)
    91  			} else {
    92  				header.Set("Content-Disposition", fmt.Sprintf("file; filename=\"%s\"", filename))
    93  			}
    94  
    95  			header.Set("Content-Type", contentType)
    96  
    97  			_, err := mfr.mpWriter.CreatePart(header)
    98  			if err != nil {
    99  				return 0, err
   100  			}
   101  		}
   102  	}
   103  
   104  	// if the buffer has something in it, read from it
   105  	if mfr.buf.Len() > 0 {
   106  		return mfr.buf.Read(buf)
   107  	}
   108  
   109  	// otherwise, read from file data
   110  	written, err = mfr.currentFile.Read(buf)
   111  	if err == io.EOF {
   112  		mfr.currentFile = nil
   113  		return written, nil
   114  	}
   115  	return written, err
   116  }
   117  
   118  // Boundary returns the boundary string to be used to separate files in the multipart data
   119  func (mfr *MultiFileReader) Boundary() string {
   120  	return mfr.mpWriter.Boundary()
   121  }