From b22151f7dc94734ebc67431240282cd569027423 Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Fri, 20 May 2011 09:03:43 -0700 Subject: [PATCH] mime/multipart: add a multipart Writer Fixes #1823 R=golang-dev, adg, robert.hencke CC=golang-dev https://golang.org/cl/4530054 --- src/pkg/mime/multipart/Makefile | 1 + src/pkg/mime/multipart/writer.go | 160 ++++++++++++++++++++++++++ src/pkg/mime/multipart/writer_test.go | 71 ++++++++++++ 3 files changed, 232 insertions(+) create mode 100644 src/pkg/mime/multipart/writer.go create mode 100644 src/pkg/mime/multipart/writer_test.go diff --git a/src/pkg/mime/multipart/Makefile b/src/pkg/mime/multipart/Makefile index 5051f0df1c5..de1a439f2a4 100644 --- a/src/pkg/mime/multipart/Makefile +++ b/src/pkg/mime/multipart/Makefile @@ -8,5 +8,6 @@ TARG=mime/multipart GOFILES=\ formdata.go\ multipart.go\ + writer.go\ include ../../../Make.pkg diff --git a/src/pkg/mime/multipart/writer.go b/src/pkg/mime/multipart/writer.go new file mode 100644 index 00000000000..74aa7be1ca7 --- /dev/null +++ b/src/pkg/mime/multipart/writer.go @@ -0,0 +1,160 @@ +// Copyright 2011 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package multipart + +import ( + "bytes" + "fmt" + "io" + "net/textproto" + "os" + "rand" + "strings" +) + +// Writer is used to generate multipart messages. +type Writer struct { + // Boundary is the random boundary string between + // parts. NewWriter will generate this but it must + // not be changed after a part has been created. + // Setting this to an invalid value will generate + // malformed messages. + Boundary string + + w io.Writer + lastpart *part +} + +// NewWriter returns a new multipart Writer with a random boundary, +// writing to w. +func NewWriter(w io.Writer) *Writer { + return &Writer{ + w: w, + Boundary: randomBoundary(), + } +} + +// FormDataContentType returns the Content-Type for an HTTP +// multipart/form-data with this Writer's Boundary. +func (w *Writer) FormDataContentType() string { + return "multipart/form-data; boundary=" + w.Boundary +} + +const randChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + +func randomBoundary() string { + var buf [60]byte + for i := range buf { + buf[i] = randChars[rand.Intn(len(randChars))] + } + return string(buf[:]) +} + +// CreatePart creates a new multipart section with the provided +// header. The previous part, if still open, is closed. The body of +// the part should be written to the returned WriteCloser. Closing the +// returned WriteCloser after writing is optional. +func (w *Writer) CreatePart(header textproto.MIMEHeader) (io.WriteCloser, os.Error) { + if w.lastpart != nil { + if err := w.lastpart.Close(); err != nil { + return nil, err + } + } + var b bytes.Buffer + fmt.Fprintf(&b, "\r\n--%s\r\n", w.Boundary) + // TODO(bradfitz): move this to textproto.MimeHeader.Write(w), have it sort + // and clean, like http.Header.Write(w) does. + for k, vv := range header { + for _, v := range vv { + fmt.Fprintf(&b, "%s: %s\r\n", k, v) + } + } + fmt.Fprintf(&b, "\r\n") + _, err := io.Copy(w.w, &b) + if err != nil { + return nil, err + } + p := &part{ + mw: w, + } + w.lastpart = p + return p, nil +} + +func escapeQuotes(s string) string { + s = strings.Replace(s, "\\", "\\\\", -1) + s = strings.Replace(s, "\"", "\\\"", -1) + return s +} + +// CreateFormFile is a convenience wrapper around CreatePart. It creates +// a new form-data header with the provided field name and file name. +func (w *Writer) CreateFormFile(fieldname, filename string) (io.WriteCloser, os.Error) { + h := make(textproto.MIMEHeader) + h.Set("Content-Disposition", + fmt.Sprintf(`form-data; name="%s"; filename="%s"`, + escapeQuotes(fieldname), escapeQuotes(filename))) + h.Set("Content-Type", "application/octet-stream") + return w.CreatePart(h) +} + +// CreateFormField is a convenience wrapper around CreatePart. It creates +// a new form-data header with the provided field name. +func (w *Writer) CreateFormField(fieldname string) (io.WriteCloser, os.Error) { + h := make(textproto.MIMEHeader) + h.Set("Content-Disposition", + fmt.Sprintf(`form-data; name="%s"`, escapeQuotes(fieldname))) + return w.CreatePart(h) +} + +// WriteField is a convenience wrapper around CreateFormField. It creates and +// writes a part with the provided name and value. +func (w *Writer) WriteField(fieldname, value string) os.Error { + p, err := w.CreateFormField(fieldname) + if err != nil { + return err + } + _, err = p.Write([]byte(value)) + if err != nil { + return err + } + return p.Close() +} + +// Close finishes the multipart message. It closes the previous part, +// if still open, and writes the trailing boundary end line to the +// output. +func (w *Writer) Close() os.Error { + if w.lastpart != nil { + if err := w.lastpart.Close(); err != nil { + return err + } + w.lastpart = nil + } + _, err := fmt.Fprintf(w.w, "\r\n--%s--\r\n", w.Boundary) + return err +} + +type part struct { + mw *Writer + closed bool + we os.Error // last error that occurred writing +} + +func (p *part) Close() os.Error { + p.closed = true + return p.we +} + +func (p *part) Write(d []byte) (n int, err os.Error) { + if p.closed { + return 0, os.NewError("multipart: Write after Close") + } + n, err = p.mw.w.Write(d) + if err != nil { + p.we = err + } + return +} diff --git a/src/pkg/mime/multipart/writer_test.go b/src/pkg/mime/multipart/writer_test.go new file mode 100644 index 00000000000..b85fbf87716 --- /dev/null +++ b/src/pkg/mime/multipart/writer_test.go @@ -0,0 +1,71 @@ +// Copyright 2011 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package multipart + +import ( + "bytes" + "io/ioutil" + "testing" +) + +func TestWriter(t *testing.T) { + fileContents := []byte("my file contents") + + var b bytes.Buffer + w := NewWriter(&b) + { + part, err := w.CreateFormFile("myfile", "my-file.txt") + if err != nil { + t.Fatalf("CreateFormFile: %v", err) + } + part.Write(fileContents) + err = w.WriteField("key", "val") + if err != nil { + t.Fatalf("CreateFormFieldValue: %v", err) + } + part.Write([]byte("val")) + err = w.Close() + if err != nil { + t.Fatalf("Close: %v", err) + } + } + + r := NewReader(&b, w.Boundary) + + part, err := r.NextPart() + if err != nil { + t.Fatalf("part 1: %v", err) + } + if g, e := part.FormName(), "myfile"; g != e { + t.Errorf("part 1: want form name %q, got %q", e, g) + } + slurp, err := ioutil.ReadAll(part) + if err != nil { + t.Fatalf("part 1: ReadAll: %v", err) + } + if e, g := string(fileContents), string(slurp); e != g { + t.Errorf("part 1: want contents %q, got %q", e, g) + } + + part, err = r.NextPart() + if err != nil { + t.Fatalf("part 2: %v", err) + } + if g, e := part.FormName(), "key"; g != e { + t.Errorf("part 2: want form name %q, got %q", e, g) + } + slurp, err = ioutil.ReadAll(part) + if err != nil { + t.Fatalf("part 2: ReadAll: %v", err) + } + if e, g := "val", string(slurp); e != g { + t.Errorf("part 2: want contents %q, got %q", e, g) + } + + part, err = r.NextPart() + if part != nil || err == nil { + t.Fatalf("expected end of parts; got %v, %v", part, err) + } +}