If signer had to cap the AMP doc, don't sign and proxy full document. (…

…ampproject#444)
MichaelRybak · Jun 24, 2020 · c85b103 · c85b103
1 parent a224ce6
commit c85b103
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git a/packager/signer/signer.go b/packager/signer/signer.go
@@ -466,13 +466,29 @@ func formatLinkHeader(preloads []*rpb.Metadata_Preload) (string, error) {
 
 // serveSignedExchange does the actual work of transforming, packaging and signed and writing to the response.
 func (this *Signer) serveSignedExchange(resp http.ResponseWriter, fetchResp *http.Response, signURL *url.URL, act string, transformVersion int64) {
-	// After this, fetchResp.Body is consumed, and attempts to read or proxy it will result in an empty body.
+	// Cap with maxBodyLength to limit the memory that serveSignedExchange
+	// consumes per request. The serveSignedExchange implementation requires all
+	// the response body in memory, and doesn't work on streamed data.
 	fetchBody, err := ioutil.ReadAll(io.LimitReader(fetchResp.Body, maxBodyLength))
 	if err != nil {
 		util.NewHTTPError(http.StatusBadGateway, "Error reading body: ", err).LogAndRespond(resp)
 		return
 	}
 
+	// If the cap has been applied, it means the document is too large. Signer
+	// won't load all the body into memory, and therefore won't sign the
+	// document. But signer can still proxy the document unsigned: first write
+	// the capped part of the body that signer has already read, and then call
+	// proxy to write headers and stream the rest of the body. ResponseWriter
+	// will make sure the headers are eventually sent before the body.
+	if len(fetchBody) == maxBodyLength {
+		resp.Write(fetchBody)
+		proxy(resp, fetchResp, nil)
+		return
+	}
+
+	// Now that fetchResp.Body is consumed, attempts to read or proxy it will result in an empty body.
+
 	// Perform local transformations.
 	r := getTransformerRequest(this.rtvCache, string(fetchBody), signURL.String())
 	r.Version = transformVersion

diff --git a/packager/signer/signer_test.go b/packager/signer/signer_test.go
@@ -1035,3 +1035,23 @@ func (this *SignerSuite) TestPrometheusMetricGatewayRequestsLatency() {
 	}
 
 }
+
+func (this *SignerSuite) TestIfCappedDontSignAndProxyFullDocument() {
+	urlSets := []util.URLSet{{
+		Sign: &util.URLPattern{[]string{"https"}, "", this.httpsHost(), stringPtr("/amp/.*"), []string{}, stringPtr(""), false, 2000, nil}}}
+
+	const uncappedTailLength = 100
+	veryLongString := strings.Repeat("a", maxBodyLength+uncappedTailLength)
+	var customFakeBody = []byte("<html amp><body>" + veryLongString)
+
+	this.fakeHandler = func(resp http.ResponseWriter, req *http.Request) {
+		resp.Header().Set("Content-Type", "text/html")
+		resp.Write(customFakeBody)
+	}
+	resp := pkgt.NewRequest(this.T(), this.new(urlSets), "/priv/doc?sign="+url.QueryEscape(this.httpsURL()+fakePath)).SetHeaders("", header).Do()
+	this.Assert().Equal(http.StatusOK, resp.StatusCode, "incorrect status: %#v", resp)
+
+	body, err := ioutil.ReadAll(resp.Body)
+	this.Require().NoError(err)
+	this.Assert().Equal(customFakeBody, body, "incorrect body: %#v", resp)
+}