multipart.go raw

   1  package message
   2  
   3  import (
   4  	"io"
   5  
   6  	"github.com/emersion/go-message/textproto"
   7  )
   8  
   9  // MultipartReader is an iterator over parts in a MIME multipart body.
  10  type MultipartReader interface {
  11  	io.Closer
  12  
  13  	// NextPart returns the next part in the multipart or an error. When there are
  14  	// no more parts, the error io.EOF is returned.
  15  	//
  16  	// Entity.Body must be read completely before the next call to NextPart,
  17  	// otherwise it will be discarded.
  18  	NextPart() (*Entity, error)
  19  }
  20  
  21  type multipartReader struct {
  22  	r *textproto.MultipartReader
  23  }
  24  
  25  // NextPart implements MultipartReader.
  26  func (r *multipartReader) NextPart() (*Entity, error) {
  27  	p, err := r.r.NextPart()
  28  	if err != nil {
  29  		return nil, err
  30  	}
  31  	return New(Header{p.Header}, p)
  32  }
  33  
  34  // Close implements io.Closer.
  35  func (r *multipartReader) Close() error {
  36  	return nil
  37  }
  38  
  39  type multipartBody struct {
  40  	header Header
  41  	parts  []*Entity
  42  
  43  	r *io.PipeReader
  44  	w *Writer
  45  
  46  	i int
  47  }
  48  
  49  // Read implements io.Reader.
  50  func (m *multipartBody) Read(p []byte) (n int, err error) {
  51  	if m.r == nil {
  52  		r, w := io.Pipe()
  53  		m.r = r
  54  
  55  		var err error
  56  		m.w, err = createWriter(w, &m.header)
  57  		if err != nil {
  58  			return 0, err
  59  		}
  60  
  61  		// Prevent calls to NextPart to succeed
  62  		m.i = len(m.parts)
  63  
  64  		go func() {
  65  			if err := m.writeBodyTo(m.w); err != nil {
  66  				w.CloseWithError(err)
  67  				return
  68  			}
  69  
  70  			if err := m.w.Close(); err != nil {
  71  				w.CloseWithError(err)
  72  				return
  73  			}
  74  
  75  			w.Close()
  76  		}()
  77  	}
  78  
  79  	return m.r.Read(p)
  80  }
  81  
  82  // Close implements io.Closer.
  83  func (m *multipartBody) Close() error {
  84  	if m.r != nil {
  85  		m.r.Close()
  86  	}
  87  	return nil
  88  }
  89  
  90  // NextPart implements MultipartReader.
  91  func (m *multipartBody) NextPart() (*Entity, error) {
  92  	if m.i >= len(m.parts) {
  93  		return nil, io.EOF
  94  	}
  95  
  96  	part := m.parts[m.i]
  97  	m.i++
  98  	return part, nil
  99  }
 100  
 101  func (m *multipartBody) writeBodyTo(w *Writer) error {
 102  	for _, p := range m.parts {
 103  		pw, err := w.CreatePart(p.Header)
 104  		if err != nil {
 105  			return err
 106  		}
 107  
 108  		if err := p.writeBodyTo(pw); err != nil {
 109  			return err
 110  		}
 111  		if err := pw.Close(); err != nil {
 112  			return err
 113  		}
 114  	}
 115  	return nil
 116  }
 117