Reverse proxy assisting the backend

main.go

package main

import (
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"sync"
)

func main() {
	log.SetFlags(0)
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	ln, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		return err
	}
	defer ln.Close()
	const backendSpecialHeader = "Assist-Mode"

	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) { io.WriteString(w, "Hi from the backend\n") })
	mux.HandleFunc("/assisted", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set(backendSpecialHeader, "please")
		io.WriteString(w, "Hi from the backend\n")
	})
	backend := &http.Server{Handler: mux}
	defer backend.Close()
	var wg sync.WaitGroup
	defer wg.Wait()
	wg.Add(1)
	go func() { defer wg.Done(); log.Printf("backend: %v", backend.Serve(ln)) }()

	dstUrl := &url.URL{Scheme: "http", Host: ln.Addr().String()}
	proxy := &http.Server{
		Addr: "localhost:8080",
		Handler: &httputil.ReverseProxy{
			Rewrite: func(r *httputil.ProxyRequest) { r.SetURL(dstUrl) },
			ModifyResponse: func(r *http.Response) error {
				if r.Header.Get(backendSpecialHeader) != "please" {
					return nil
				}
				log.Println("found special backend header, replacing response")
				r.Header.Del("Content-Length")
				r.Header.Del(backendSpecialHeader)
				// consume the whole backend response
				if _, err := io.Copy(io.Discard, r.Body); err != nil {
					log.Println("consuming response:", err)
					return err
				}
				r.Body = io.NopCloser(strings.NewReader("This response was replaced by the proxy on behalf of the backend.\n"))
				return nil
			},
		},
	}
	log.Printf("listening on %s, try opening http://%[1]s/, and http://%[1]s/assisted", proxy.Addr)
	return proxy.ListenAndServe()
}