client: report sensible errors if response is excessively large or empty

omaha/client/error.go

@@ -15,9 +15,25 @@
 package client
 
 import (
+	"encoding/xml"
+	"errors"
+	"io"
 	"net/http"
 )
 
+var (
+	bodySizeError  = errors.New("http response exceeded 1MB")
+	bodyEmptyError = errors.New("http response was empty")
+)
+
+// xml doesn't return the standard io.ErrUnexpectedEOF so check for both.
+func isUnexpectedEOF(err error) bool {
+	if xerr, ok := err.(*xml.SyntaxError); ok {
+		return xerr.Msg == "unexpected EOF"
+	}
+	return err == io.ErrUnexpectedEOF
+}
+
 // httpError implements error and net.Error for http responses.
 type httpError struct {
 	*http.Response

omaha/client/http.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"encoding/xml"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"time"
@@ -50,8 +51,17 @@ func (hc *httpClient) doPost(url string, reqBody []byte) (*omaha.Response, error
 	}
 	defer resp.Body.Close()
 
+	// A response over 1M in size is certainly bogus.
+	respBody := &io.LimitedReader{R: resp.Body, N: 1024 * 1024}
 	contentType := resp.Header.Get("Content-Type")
-	omahaResp, err := omaha.ParseResponse(contentType, resp.Body)
+	omahaResp, err := omaha.ParseResponse(contentType, respBody)
+
+	// Report a more sensible error if we truncated the body.
+	if isUnexpectedEOF(err) && respBody.N <= 0 {
+		err = bodySizeError
+	} else if err == io.EOF {
+		err = bodyEmptyError
+	}
 
 	// Prefer reporting HTTP errors over XML parsing errors.
 	if resp.StatusCode != http.StatusOK {

omaha/client/http_test.go

@@ -15,6 +15,7 @@
 package client
 
 import (
+	"bytes"
 	"net"
 	"net/http"
 	"strings"
@@ -164,3 +165,48 @@ func TestHTTPClientRetry(t *testing.T) {
 		t.Fatalf("Server received %d requests, not 2", f.h.reqs)
 	}
 }
+
+// should result in an unexected EOF
+func largeHandler1(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "text/xml; charset=utf-8")
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte(`<?xml version="1.0" encoding="UTF-8"?><response protocol="3.0">`))
+	w.Write(bytes.Repeat([]byte{' '}, 2*1024*1024))
+	w.Write([]byte(`</response>`))
+}
+
+// should result in an EOF
+func largeHandler2(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "text/xml; charset=utf-8")
+	w.Write(bytes.Repeat([]byte{' '}, 2*1024*1024))
+}
+
+func TestHTTPClientLarge(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+
+	s := &http.Server{
+		Handler: http.HandlerFunc(largeHandler1),
+	}
+	go s.Serve(l)
+
+	c := newHTTPClient()
+	url := "http://" + l.Addr().String()
+
+	_, err = c.doPost(url, []byte(sampleRequest))
+	if err != bodySizeError {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	// switch to failing before XML is read instead of half-way
+	// through (which results in a different error internally)
+	s.Handler = http.HandlerFunc(largeHandler2)
+
+	_, err = c.doPost(url, []byte(sampleRequest))
+	if err != bodyEmptyError {
+		t.Errorf("Unexpected error: %v", err)
+	}
+}