diff --git a/proxy/proxy_intercept.go b/proxy/proxy_intercept.go index 69f83c83..5dd109db 100644 --- a/proxy/proxy_intercept.go +++ b/proxy/proxy_intercept.go @@ -9,6 +9,7 @@ import ( "net" "net/http" "net/http/httputil" + "sync" "github.com/fsouza/go-dockerclient" . "github.com/weaveworks/weave/common" @@ -94,25 +95,30 @@ func doRawStream(w http.ResponseWriter, resp *http.Response, client *httputil.Cl return } - upDone := make(chan struct{}) - downDone := make(chan struct{}) - go copyStream(down, io.MultiReader(rem, up), upDone) - go copyStream(up, downBuf, downDone) - <-upDone - <-downDone + var wg sync.WaitGroup + wg.Add(2) + go copyStream(down, io.MultiReader(rem, up), &wg) + go copyStream(up, downBuf, &wg) + wg.Wait() } -func copyStream(dst io.Writer, src io.Reader, done chan struct{}) { - defer close(done) +type closeWriter interface { + CloseWrite() error +} + +func copyStream(dst io.WriteCloser, src io.Reader, wg *sync.WaitGroup) { + defer wg.Done() if _, err := io.Copy(dst, src); err != nil { Warning.Print(err) } - if c, ok := dst.(interface { - CloseWrite() error - }); ok { - if err := c.CloseWrite(); err != nil { - Warning.Printf("Error closing connection: %s", err) - } + var err error + if c, ok := dst.(closeWriter); ok { + err = c.CloseWrite() + } else { + err = dst.Close() + } + if err != nil { + Warning.Printf("Error closing connection: %s", err) } }