Support encoding for 'edit' command

Bugfix for truncating existing output file

Minor add log error in editing request

Change-Id: I4492dc7efa9fe2d722c81932c80afa2e54c6de4f
Reviewed-on: https://chromium-review.googlesource.com/951784
Commit-Queue: Oleg Neumyvakin <oneumyvakin@gmail.com>
Reviewed-by: Tom Bergan <tombergan@chromium.org>

web_page_replay_go/src/httparchive.go

@@ -26,6 +26,7 @@ const usage = "%s [ls|cat|edit] [options] archive_file [output_file]"
 
 type Config struct {
 	method, host, fullPath string
+	decodeResponseBody     bool
 }
 
 func (cfg *Config) Flags() []cli.Flag {
@@ -48,6 +49,11 @@ func (cfg *Config) Flags() []cli.Flag {
 			Usage:       "Only show URLs matching this full path.",
 			Destination: &cfg.fullPath,
 		},
+		cli.BoolFlag{
+			Name:        "decode_response_body",
+			Usage:       "Decode/encode response body according to Content-Encoding header.",
+			Destination: &cfg.decodeResponseBody,
+		},
 	}
 }
 
@@ -102,6 +108,11 @@ func edit(cfg *Config, a *webpagereplay.Archive, outfile string) {
 		if err := req.Write(w); err != nil {
 			return err
 		}
+		if cfg.decodeResponseBody {
+			if err := webpagereplay.DecompressResponse(resp); err != nil {
+				return fmt.Errorf("couldn't decompress body: %v", err)
+			}
+		}
 		return resp.Write(w)
 	}
 
@@ -118,6 +129,12 @@ func edit(cfg *Config, a *webpagereplay.Archive, outfile string) {
 			}
 			return nil, nil, fmt.Errorf("couldn't unmarshal response: %v", err)
 		}
+		if cfg.decodeResponseBody {
+			// Compress body back according to Content-Encoding
+			if err := compressResponse(resp); err != nil {
+				return nil, nil, fmt.Errorf("couldn't compress response: %v", err)
+			}
+		}
 		// Read resp.Body into a buffer since the tmpfile is about to be deleted.
 		body, err := ioutil.ReadAll(resp.Body)
 		resp.Body.Close()
@@ -165,7 +182,7 @@ func edit(cfg *Config, a *webpagereplay.Archive, outfile string) {
 			defer tmpf.Close()
 			newReq, newResp, err := unmarshalAfterEdit(tmpf)
 			if err != nil {
-				fmt.Printf("Error in editing request. Try again.\n")
+				fmt.Printf("Error in editing request. Try again: %v\n", err)
 				continue
 			}
 			return newReq, newResp, nil
@@ -176,7 +193,7 @@ func edit(cfg *Config, a *webpagereplay.Archive, outfile string) {
 		return
 	}
 
-	outf, err := os.OpenFile(outfile, os.O_WRONLY|os.O_CREATE, os.FileMode(0660))
+	outf, err := os.OpenFile(outfile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, os.FileMode(0660))
 	if err != nil {
 		fmt.Printf("Error opening output file %s: %v\n", outfile, err)
 		return
@@ -194,6 +211,29 @@ func edit(cfg *Config, a *webpagereplay.Archive, outfile string) {
 	fmt.Printf("Wrote edited archive to %s\n", outfile)
 }
 
+// compressResponse compresses resp.Body in place according to resp's Content-Encoding header.
+func compressResponse(resp *http.Response) error {
+	ce := strings.ToLower(resp.Header.Get("Content-Encoding"))
+	if ce == "" {
+		return nil
+	}
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return err
+	}
+	resp.Body.Close()
+
+	body, newCE, err := webpagereplay.CompressBody(ce, body)
+	if err != nil {
+		return err
+	}
+	if ce != newCE {
+		return fmt.Errorf("can't compress body to '%s' recieved Content-Encoding: '%s'", ce, newCE)
+	}
+	resp.Body = ioutil.NopCloser(bytes.NewReader(body))
+	return nil
+}
+
 func main() {
 	progName := filepath.Base(os.Args[0])
 	cfg := &Config{}

web_page_replay_go/src/webpagereplay/proxy.go

@@ -113,7 +113,7 @@ func (proxy *replayingProxy) ServeHTTP(w http.ResponseWriter, req *http.Request)
 			w.WriteHeader(http.StatusNotFound)
 			return
 		}
-		body, ce, err := compressBody(clientAE, body)
+		body, ce, err := CompressBody(clientAE, body)
 		if err != nil {
 			logf("error recompressing response body: %v", err)
 			w.WriteHeader(http.StatusNotFound)

web_page_replay_go/src/webpagereplay/transformers.go

@@ -81,7 +81,7 @@ func transformResponseBody(resp *http.Response, f func([]byte) []byte) error {
 	// Transform and recompress as needed.
 	body = f(body)
 	if isCompressed {
-		body, _, err = compressBody(ce, body)
+		body, _, err = CompressBody(ce, body)
 		if err != nil {
 			return failEarly(body, err)
 		}
@@ -136,9 +136,9 @@ func decompressBody(ce string, compressed []byte) ([]byte, error) {
 	return ioutil.ReadAll(r)
 }
 
-// compressBody reads a response body and compresses according to the given Accept-Encoding.
+// CompressBody reads a response body and compresses according to the given Accept-Encoding.
 // The chosen compressed encoding is returned along with the compressed body.
-func compressBody(ae string, uncompressed []byte) ([]byte, string, error) {
+func CompressBody(ae string, uncompressed []byte) ([]byte, string, error) {
 	var buf bytes.Buffer
 	var w io.WriteCloser
 	outCE := ""