fix content-length handling and multiple writes in http2 responsewriter

This commit is contained in:
Lucas Clemente 2016-04-19 14:35:02 +02:00
parent f2db542b3f
commit fbb89f5467

View file

@ -45,6 +45,9 @@ type responseWriter struct {
header http.Header header http.Header
headerWritten bool headerWritten bool
bytesWritten int
contentLength int
} }
func (w *responseWriter) Header() http.Header { func (w *responseWriter) Header() http.Header {
@ -64,13 +67,15 @@ func (w *responseWriter) WriteHeader(status int) {
enc.WriteField(hpack.HeaderField{Name: k, Value: v[0]}) enc.WriteField(hpack.HeaderField{Name: k, Value: v[0]})
} }
fmt.Printf("responding with %d %#v\n", status, w.header) fmt.Printf("Responding with %d %#v\n", status, w.header)
h2framer := http2.NewFramer(w.headerStream, nil) h2framer := http2.NewFramer(w.headerStream, nil)
h2framer.WriteHeaders(http2.HeadersFrameParam{ h2framer.WriteHeaders(http2.HeadersFrameParam{
StreamID: uint32(w.dataStreamID), StreamID: uint32(w.dataStreamID),
EndHeaders: true, EndHeaders: true,
BlockFragment: headers.Bytes(), BlockFragment: headers.Bytes(),
}) })
w.contentLength, _ = strconv.Atoi(w.header.Get("content-length"))
} }
func (w *responseWriter) Write(p []byte) (int, error) { func (w *responseWriter) Write(p []byte) (int, error) {
@ -79,12 +84,22 @@ func (w *responseWriter) Write(p []byte) (int, error) {
} }
if len(p) != 0 { if len(p) != 0 {
dataStream, err := w.session.NewStream(w.dataStreamID) if w.dataStream == nil {
if err != nil { var err error
return 0, fmt.Errorf("error creating data stream: %s\n", err.Error()) w.dataStream, err = w.session.NewStream(w.dataStreamID)
if err != nil {
return 0, fmt.Errorf("error creating data stream: %s\n", err.Error())
}
} }
defer dataStream.Close()
return dataStream.Write(p) n, err := w.dataStream.Write(p)
w.bytesWritten += n
if w.bytesWritten >= w.contentLength {
defer w.dataStream.Close()
}
return n, err
} }
return 0, nil return 0, nil