fix: fix proxy unmodified response when body tag is missing (#611)

Co-authored-by: Neemias Almeida <ndalmeida@bol.com>

runner/proxy.go

@@ -55,28 +55,24 @@ func (p *Proxy) Reload() {
  p.stream.Reload()
 }
 
-func (p *Proxy) injectLiveReload(resp *http.Response) (page string, modified bool) {
- if !strings.Contains(resp.Header.Get("Content-Type"), "text/html") {
- return page, false
- }
-
+func (p *Proxy) injectLiveReload(resp *http.Response) (string, error) {
  buf := new(bytes.Buffer)
  if _, err := buf.ReadFrom(resp.Body); err != nil {
- return page, false
+ return "", fmt.Errorf("proxy inject: failed to read body from http response")
  }
- page = buf.String()
+ page := buf.String()
 
- // the script will be injected before the end of the body tag. In case the tag is missing, the injection will be skipped.
+ // the script will be injected before the end of the body tag. In case the tag is missing, the injection will be skipped with no error.
  body := strings.LastIndex(page, "</body>")
  if body == -1 {
- return page, false
+ return page, nil
  }
 
  script := fmt.Sprintf(
  `<script>new EventSource("http://localhost:%d/internal/reload").onmessage = () => { location.reload() }</script>`,
  p.config.ProxyPort,
  )
- return page[:body] + script + page[body:], true
+ return page[:body] + script + page[body:], nil
 }
 
 func (p *Proxy) proxyHandler(w http.ResponseWriter, r *http.Request) {
@@ -88,6 +84,7 @@ func (p *Proxy) proxyHandler(w http.ResponseWriter, r *http.Request) {
 
  if err := r.ParseForm(); err != nil {
  http.Error(w, "proxy handler: bad form", http.StatusInternalServerError)
+ return
  }
  var body io.Reader
  if len(r.Form) > 0 {
@@ -98,6 +95,7 @@ func (p *Proxy) proxyHandler(w http.ResponseWriter, r *http.Request) {
  req, err := http.NewRequest(r.Method, appURL.String(), body)
  if err != nil {
  http.Error(w, "proxy handler: unable to create request", http.StatusInternalServerError)
+ return
  }
 
  // Copy the headers from the original request
@@ -117,6 +115,7 @@ func (p *Proxy) proxyHandler(w http.ResponseWriter, r *http.Request) {
  }
  if !errors.Is(err, syscall.ECONNREFUSED) {
  http.Error(w, "proxy handler: unable to reach app", http.StatusInternalServerError)
+ return
  }
  time.Sleep(100 * time.Millisecond)
  }
@@ -133,16 +132,22 @@ func (p *Proxy) proxyHandler(w http.ResponseWriter, r *http.Request) {
  }
  w.WriteHeader(resp.StatusCode)
 
- page, modified := p.injectLiveReload(resp)
- if modified {
- w.Header().Set("Content-Length", strconv.Itoa((len([]byte(page)))))
- if _, err := io.WriteString(w, page); err != nil {
- http.Error(w, "proxy handler: unable to inject live reload script", http.StatusInternalServerError)
- }
- } else {
+ if !strings.Contains(resp.Header.Get("Content-Type"), "text/html") {
  w.Header().Set("Content-Length", resp.Header.Get("Content-Length"))
  if _, err := io.Copy(w, resp.Body); err != nil {
  http.Error(w, "proxy handler: failed to forward the response body", http.StatusInternalServerError)
+ return
+ }
+ } else {
+ page, err := p.injectLiveReload(resp)
+ if err != nil {
+ http.Error(w, err.Error(), http.StatusInternalServerError)
+ return
+ }
+ w.Header().Set("Content-Length", strconv.Itoa((len([]byte(page)))))
+ if _, err := io.WriteString(w, page); err != nil {
+ http.Error(w, "proxy handler: unable to inject live reload script", http.StatusInternalServerError)
+ return
  }
  }
 }