litefs proxy_server.go with simple websocket support

proxy_server.go

package main

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

type ProxyServer struct {
	Target        string
	HTTPTransport *http.Transport
	WSProxy       *httputil.ReverseProxy
}

func NewProxyServer(target string) (*ProxyServer, error) {
	targetURL, err := url.Parse("http://" + target)
	if err != nil {
		return nil, err
	}

	return &ProxyServer{
		Target:        target,
		HTTPTransport: &http.Transport{},
		WSProxy:       httputil.NewSingleHostReverseProxy(targetURL),
	}, nil
}

func (s *ProxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.proxyToTarget(w, r, false)
}

func (s *ProxyServer) proxyToTarget(w http.ResponseWriter, r *http.Request, passthrough bool) {
	// If the request is a websocket request, proxy with the websocket proxy.
	if r.Header.Get("Upgrade") == "websocket" {
		s.WSProxy.ServeHTTP(w, r)
		return
	}

	// Update request URL to target server.
	r.URL.Scheme = "http"
	r.URL.Host = s.Target

	resp, err := s.HTTPTransport.RoundTrip(r)
	if err != nil {
		http.Error(w, "Proxy error: "+err.Error(), http.StatusBadGateway)
		return
	}
	defer func() { _ = resp.Body.Close() }()

	// Copy response headers
	for key, values := range resp.Header {
		for _, v := range values {
			w.Header().Add(key, v)
		}
	}

	// Set response code and copy the body.
	w.WriteHeader(resp.StatusCode)
	if err := copyAndFlush(w, resp.Body); err != nil {
		log.Printf("http: proxy response error: %s", err)
		return
	}
}

// copyAndFlush implements a basic io.Copy() but calls dst.Flush() after every write.
// dst must implement http.Flusher or else it will panic.
func copyAndFlush(dst io.Writer, src io.Reader) error {
	buf := make([]byte, 32*1024)

	for {
		n, err := src.Read(buf)
		if n > 0 {
			if _, e := dst.Write(buf[:n]); e != nil {
				return err
			}
			dst.(http.Flusher).Flush()
		}

		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
	}
}

func main() {
	// Hostport of application that is being proxied.
	target := "localhost:3008"
	proxy, err := NewProxyServer(target)
	if err != nil {
		log.Fatalf("Error creating proxy: %v", err)
	}

	log.Printf("Starting proxy server for %s", target)
	log.Fatal(http.ListenAndServe(":3009", proxy))
}