diff --git a/proxy/proxy_intercept.go b/proxy/proxy_intercept.go index 1b6ca46c..295bb68f 100644 --- a/proxy/proxy_intercept.go +++ b/proxy/proxy_intercept.go @@ -2,18 +2,31 @@ package proxy import ( "bufio" + "bytes" "errors" "fmt" "io" - "io/ioutil" "net" "net/http" "net/http/httputil" + "strconv" "github.com/fsouza/go-dockerclient" . "github.com/weaveworks/weave/common" ) +const ( + maxLineLength = 4096 // assumed <= bufio.defaultBufSize + maxChunkSize = bufio.MaxScanTokenSize +) + +var ( + ErrChunkTooLong = errors.New("chunk too long") + ErrInvalidChunkLength = errors.New("invalid byte in chunk length") + ErrLineTooLong = errors.New("header line too long") + ErrMalformedChunkEncoding = errors.New("malformed chunked encoding") +) + func (proxy *Proxy) Intercept(i interceptor, w http.ResponseWriter, r *http.Request) { if err := i.InterceptRequest(r); err != nil { switch err.(type) { @@ -71,7 +84,7 @@ func (proxy *Proxy) Intercept(i interceptor, w http.ResponseWriter, r *http.Requ } func doRawStream(w http.ResponseWriter, resp *http.Response, client *httputil.ClientConn) { - down, downBuf, up, rem, err := hijack(w, client) + down, downBuf, up, remaining, err := hijack(w, client) if err != nil { http.Error(w, "Unable to hijack connection for raw stream mode", http.StatusInternalServerError) return @@ -96,7 +109,7 @@ func doRawStream(w http.ResponseWriter, resp *http.Response, client *httputil.Cl upDone := make(chan struct{}) downDone := make(chan struct{}) - go copyStream(down, io.MultiReader(rem, up), upDone) + go copyStream(down, io.MultiReader(remaining, up), upDone) go copyStream(up, downBuf, downDone) <-upDone <-downDone @@ -116,33 +129,81 @@ func copyStream(dst io.Writer, src io.Reader, done chan struct{}) { } } -func doChunkedResponse(w http.ResponseWriter, resp *http.Response, client *httputil.ClientConn) { - // Because we can't go back to request/response after we - // hijack the connection, we need to close it and make the - // client open another. - w.Header().Add("Connection", "close") - w.WriteHeader(resp.StatusCode) - - down, _, up, rem, err := hijack(w, client) - if err != nil { - http.Error(w, "Unable to hijack response stream for chunked response", http.StatusInternalServerError) - return - } - defer up.Close() - defer down.Close() - // Copy the chunked response body to downstream, - // stopping at the end of the chunked section. - rawResponseBody := io.MultiReader(rem, up) - if _, err := io.Copy(ioutil.Discard, httputil.NewChunkedReader(io.TeeReader(rawResponseBody, down))); err != nil { - http.Error(w, "Error copying chunked response body", http.StatusInternalServerError) - return - } - resp.Trailer.Write(down) - // a chunked response ends with a CRLF - down.Write([]byte("\r\n")) +type writeFlusher interface { + io.Writer + http.Flusher } -func hijack(w http.ResponseWriter, client *httputil.ClientConn) (down net.Conn, downBuf *bufio.ReadWriter, up net.Conn, rem io.Reader, err error) { +func doChunkedResponse(w http.ResponseWriter, resp *http.Response, client *httputil.ClientConn) { + wf, ok := w.(writeFlusher) + if !ok { + http.Error(w, "Error forwarding chunked response body: flush not available", http.StatusInternalServerError) + return + } + + w.WriteHeader(resp.StatusCode) + + up, remaining := client.Hijack() + defer up.Close() + + var err error + chunks := bufio.NewScanner(io.MultiReader(remaining, up)) + chunks.Split(splitChunks) + for chunks.Scan() && err == nil { + _, err = wf.Write(chunks.Bytes()) + wf.Flush() + } + if err == nil { + err = chunks.Err() + } + if err != nil { + Log.Errorf("Error forwarding chunked response body: %s", err) + } +} + +// a bufio.SplitFunc for http chunks +func splitChunks(data []byte, atEOF bool) (advance int, token []byte, err error) { + if atEOF && len(data) == 0 { + return 0, nil, nil + } + + i := bytes.IndexByte(data, '\n') + if i < 0 { + return 0, nil, nil + } + if i > maxLineLength { + return 0, nil, ErrLineTooLong + } + + chunkSize64, err := strconv.ParseInt( + string(bytes.TrimRight(data[:i], " \t\r\n")), + 16, + 64, + ) + switch { + case err != nil: + return 0, nil, ErrInvalidChunkLength + case chunkSize64 > maxChunkSize: + return 0, nil, ErrChunkTooLong + case chunkSize64 == 0: + return 0, nil, io.EOF + } + chunkSize := int(chunkSize64) + + data = data[i+1:] + + if len(data) < chunkSize+2 { + return 0, nil, nil + } + + if data[chunkSize] != '\r' || data[chunkSize+1] != '\n' { + return 0, nil, ErrMalformedChunkEncoding + } + + return i + chunkSize + 3, data[:chunkSize], nil +} + +func hijack(w http.ResponseWriter, client *httputil.ClientConn) (down net.Conn, downBuf *bufio.ReadWriter, up net.Conn, remaining io.Reader, err error) { hj, ok := w.(http.Hijacker) if !ok { err = errors.New("Unable to cast to Hijack") @@ -152,6 +213,6 @@ func hijack(w http.ResponseWriter, client *httputil.ClientConn) (down net.Conn, if err != nil { return } - up, rem = client.Hijack() + up, remaining = client.Hijack() return } diff --git a/proxy/proxy_intercept_test.go b/proxy/proxy_intercept_test.go new file mode 100644 index 00000000..15f96230 --- /dev/null +++ b/proxy/proxy_intercept_test.go @@ -0,0 +1,136 @@ +// Based on net/http/internal +package proxy + +import ( + "bufio" + "bytes" + "strconv" + "strings" + "testing" +) + +func TestChunk(t *testing.T) { + r := bufio.NewScanner(bytes.NewBufferString( + "7\r\nhello, \r\n17\r\nworld! 0123456789abcdef\r\n0\r\n", + )) + r.Split(splitChunks) + + assertNextChunk(t, r, "hello, ") + assertNextChunk(t, r, "world! 0123456789abcdef") + assertNoMoreChunks(t, r) +} + +func TestMalformedChunks(t *testing.T) { + r := bufio.NewScanner(bytes.NewBufferString( + "7\r\nhello, GARBAGEBYTES17\r\nworld! 0123456789abcdef\r\n0\r\n", + )) + r.Split(splitChunks) + + // First chunk fails + { + if r.Scan() { + t.Errorf("Expected failure when reading chunks, but got one") + } + e := "malformed chunked encoding" + if r.Err() == nil || r.Err().Error() != e { + t.Errorf("chunk reader errored %q; want %q", r.Err(), e) + } + data := r.Bytes() + if len(data) != 0 { + t.Errorf("chunk should have been empty. got %q", data) + } + } + + if r.Scan() { + t.Errorf("Expected no more chunks, but found too many") + } +} + +func TestChunkTooLarge(t *testing.T) { + data := make([]byte, maxChunkSize+1) + r := bufio.NewScanner(bytes.NewBufferString(strings.Join( + []string{ + strconv.FormatInt(maxChunkSize+1, 16), string(data), + "0", "", + }, + "\r\n", + ))) + r.Split(splitChunks) + + // First chunk fails + { + if r.Scan() { + t.Errorf("Expected failure when reading chunks, but got one") + } + e := "chunk too long" + if r.Err() == nil || r.Err().Error() != e { + t.Errorf("chunk reader errored %q; want %q", r.Err(), e) + } + data := r.Bytes() + if len(data) != 0 { + t.Errorf("chunk should have been empty. got %q", data) + } + } + + if r.Scan() { + t.Errorf("Expected no more chunks, but found too many") + } +} + +func TestInvalidChunkSize(t *testing.T) { + r := bufio.NewScanner(bytes.NewBufferString( + "foobar\r\nhello, \r\n0\r\n", + )) + r.Split(splitChunks) + + // First chunk fails + { + if r.Scan() { + t.Errorf("Expected failure when reading chunks, but got one") + } + e := "invalid byte in chunk length" + if r.Err() == nil || r.Err().Error() != e { + t.Errorf("chunk reader errored %q; want %q", r.Err(), e) + } + data := r.Bytes() + if len(data) != 0 { + t.Errorf("chunk should have been empty. got %q", data) + } + } + + if r.Scan() { + t.Errorf("Expected no more chunks, but found too many") + } +} + +func TestBytesAfterLastChunkAreIgnored(t *testing.T) { + r := bufio.NewScanner(bytes.NewBufferString( + "7\r\nhello, \r\n0\r\nGARBAGEBYTES", + )) + r.Split(splitChunks) + + assertNextChunk(t, r, "hello, ") + assertNoMoreChunks(t, r) +} + +func assertNextChunk(t *testing.T, r *bufio.Scanner, expected string) { + if !r.Scan() { + t.Fatalf("Expected chunk, but ran out early: %v", r.Err()) + } + if r.Err() != nil { + t.Fatalf("Error reading chunk: %q", r.Err()) + } + data := r.Bytes() + if string(data) != expected { + t.Errorf("chunk reader read %q; want %q", data, expected) + } +} + +func assertNoMoreChunks(t *testing.T, r *bufio.Scanner) { + if r.Scan() { + t.Errorf("Expected no more chunks, but found too many") + } + if r.Err() != nil { + t.Errorf("Expected no error, but found: %q", r.Err()) + } +}