HTTP2 Golang Proxy w/ Load Balancer

proxy.go

package main

// TARGET='http://localhost:3006/,http://localhost:3006/,http://localhost:3001' go run proxy.go

import (
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

func main() {
	t, exists := os.LookupEnv("TARGET")
	if !exists {
		log.Fatal("TARGET environment variable must be set")
	}
	targets := strings.Split(t, ",")
	urls := make([]*url.URL, len(targets))
	for i, target := range targets {
		tgt, err := url.Parse(target)
		urls[i] = tgt
		if err != nil {
			log.Println("Invalid target URL: ", target)
			log.Fatal(err)
		}
	}

	log.Println("Targets: ", urls)
	last := 0
	proxy := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		var resp *http.Response
		var err error = nil
		for i := last; i < len(urls)+last; i++ {
			resp, err = makeReq(req, urls[i%len(urls)])
			// fmt.Println("ff")
			// fmt.Println(i, last, len(urls))
			// fmt.Println(i, last, len(urls))
			if err != nil {
				log.Println("Error making request: ", err)
				continue
			} else {
				err = nil
				last = ((i + 1) % len(urls))
				break
			}
		}
		if resp == nil {
			last = 0
			if err != nil {
				http.Error(rw, err.Error(), http.StatusInternalServerError)
			} else {
				http.Error(rw, "No targets available", http.StatusServiceUnavailable)
			}
			// rw.WriteHeader(http.StatusInternalServerError)
			// fmt.Fprint(rw, "Proxy Error: "+err.Error())
			return
		}

		for key, values := range resp.Header {
			for _, value := range values {
				rw.Header().Set(key, value)
			}
		}

		done := make(chan bool)
		go func() {
			for {
				select {
				case <-time.Tick(time.Millisecond * 10):
					rw.(http.Flusher).Flush()
				case <-done:
					return
				}
			}
		}()

		var trailerKeys []string
		for key := range resp.Trailer {
			trailerKeys = append(trailerKeys, key)
		}

		rw.Header().Set("Trailer", strings.Join(trailerKeys, ","))

		rw.WriteHeader(resp.StatusCode)
		io.Copy(rw, resp.Body)
		close(done)

		for key, values := range resp.Trailer {
			for _, value := range values {
				rw.Header().Set(key, value)
			}
		}

	})
	http.ListenAndServe(":8080", proxy)
}

func makeReq(req *http.Request, url *url.URL) (*http.Response, error) {
	req.Host = url.Host
	req.URL.Host = url.Host
	req.URL.Scheme = url.Scheme
	req.RequestURI = ""

	log.Printf("%s %s %s\n", req.RemoteAddr, req.Method, req.URL)

	host, _, _ := net.SplitHostPort(req.RemoteAddr)
	req.Header.Set("X-Forwarded-For", host)

	return http.DefaultClient.Do(req)
}