Skip to content

Instantly share code, notes, and snippets.

@bubbajoe
Created September 6, 2022 16:16
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save bubbajoe/838867c75fd55e2e68ad9fe3504ab326 to your computer and use it in GitHub Desktop.
Save bubbajoe/838867c75fd55e2e68ad9fe3504ab326 to your computer and use it in GitHub Desktop.
HTTP2 Golang Proxy w/ Load Balancer
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)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment