package files import ( "bytes" "fmt" "io" "mime/multipart" "net/textproto" "net/url" "path" "sync" ) // MultiFileReader reads from a `commands.Node` (which can be a directory of files // or a regular file) as HTTP multipart encoded data. type MultiFileReader struct { io.Reader // directory stack for NextFile files []DirIterator path []string currentFile Node 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.Directory`. // 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 Directory, form bool) *MultiFileReader { it := file.Entries() mfr := &MultiFileReader{ files: []DirIterator{it}, path: []string{""}, 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 entry DirEntry for entry == nil { if len(mfr.files) == 0 { mfr.mpWriter.Close() mfr.closed = true return mfr.buf.Read(buf) } if !mfr.files[len(mfr.files)-1].Next() { if mfr.files[len(mfr.files)-1].Err() != nil { return 0, mfr.files[len(mfr.files)-1].Err() } mfr.files = mfr.files[:len(mfr.files)-1] mfr.path = mfr.path[:len(mfr.path)-1] continue } entry = mfr.files[len(mfr.files)-1] } // handle starting a new file part if !mfr.closed { mfr.currentFile = entry.Node() // write the boundary and headers header := make(textproto.MIMEHeader) filename := url.QueryEscape(path.Join(path.Join(mfr.path...), entry.Name())) dispositionPrefix := "attachment" if mfr.form { dispositionPrefix = "form-data; name=\"file\"" } header.Set("Content-Disposition", fmt.Sprintf("%s; filename=\"%s\"", dispositionPrefix, filename)) var contentType string switch f := entry.Node().(type) { case *Symlink: contentType = "application/symlink" case Directory: newIt := f.Entries() mfr.files = append(mfr.files, newIt) mfr.path = append(mfr.path, entry.Name()) contentType = "application/x-directory" case File: // otherwise, use the file as a reader to read its contents contentType = "application/octet-stream" default: return 0, ErrNotSupported } header.Set("Content-Type", contentType) if rf, ok := entry.Node().(FileInfo); 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 switch f := mfr.currentFile.(type) { case File: written, err = f.Read(buf) if err != io.EOF { return written, err } } if err := mfr.currentFile.Close(); err != nil { return written, err } mfr.currentFile = nil return written, nil } // Boundary returns the boundary string to be used to separate files in the multipart data func (mfr *MultiFileReader) Boundary() string { return mfr.mpWriter.Boundary() }