Use httputil.ReverseProxy instead of low level TCP conn handling

README.md

@@ -37,5 +37,5 @@ go mod download
 Run the program
 
 ```shell
-go run main.go
+go run .
 ```

go.mod

@@ -3,6 +3,7 @@ module github.com/cyakimov/helios
 go 1.12
 
 require (
+	github.com/gorilla/mux v1.7.1
 	github.com/sirupsen/logrus v1.4.1
 	golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a // indirect
 )

go.sum

@@ -1,5 +1,7 @@
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/gorilla/mux v1.7.1 h1:Dw4jY2nghMMRsh1ol8dv1axHkDwMQK2DHerMNJsIpJU=
+github.com/gorilla/mux v1.7.1/go.mod h1:1lud6UwP+6orDFRuTfBEV8e9/aOM/c4fVVCaMa2zaAs=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1 h1:mweAR1A6xJ3oS2pRaGiHgQ4OO8tzTaLawm8vnODuwDk=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

main.go

@@ -1,13 +1,16 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"flag"
 	"fmt"
+	"github.com/gorilla/mux"
 	log "github.com/sirupsen/logrus"
-	"io"
-	"net"
+	"net/http"
+	"net/url"
 	"os"
+	"os/signal"
 	"time"
 )
 
@@ -26,7 +29,7 @@ func init() {
 	_ = os.Setenv("GODEBUG", os.Getenv("GODEBUG")+",tls13=1")
 
 	flag.IntVar(&port, "port", 443, "Listen port")
-	flag.StringVar(&ip, "ip", "", "Listen IP")
+	flag.StringVar(&ip, "ip", "0.0.0.0", "Listen IP")
 	flag.IntVar(&timeout, "timeout", 3, "Read/Write timeout")
 	flag.StringVar(&certificatePath, "certificate", "localhost.pem", "TLS certificate path")
 	flag.StringVar(&privateKeyPath, "key", "localhost-key.pem", "TLS private key path")
@@ -35,77 +38,53 @@ func init() {
 	log.SetLevel(log.DebugLevel)
 }
 
-func serveConn(conn net.Conn) {
-	log.Debugf("Got a new connection from %s", conn.RemoteAddr().String())
-	defer conn.Close()
-
-	// Set read/write timeout
-	err := conn.SetDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
-	// drop connection on error
-	if err != nil {
-		return
-	}
-
-	tlsConn := tls.Server(conn, tlsConfig)
-	if err := tlsConn.Handshake(); err != nil {
-		log.Errorf("TLS handshake error: %q", err)
-		return
-	}
-	proxy(tlsConn)
-}
-
-func proxy(conn net.Conn) {
-	remote, err := net.Dial("tcp", "httpbin.org:443")
-	if err != nil {
-		return
-	}
-
-	defer remote.Close()
-
-	tlsConn := tls.Client(remote, &tls.Config{
-		ServerName: "httpbin.org",
-	})
-	if err = tlsConn.Handshake(); err != nil {
-		log.Error(err)
-		return
-	}
-
-	if err = conn.SetDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
-		log.Debugf("Connection setTimeout error: %q", err)
-		return
-	}
-
-	go io.Copy(tlsConn, conn)
-	io.Copy(conn, tlsConn)
-}
-
 func main() {
-	// load TLS certificate
-	cert, err := tls.LoadX509KeyPair(certificatePath, privateKeyPath)
-	if err != nil {
-		log.Fatal(err)
-	}
+	var wait time.Duration
+
 	tlsConfig = &tls.Config{
-		MinVersion:   tls.VersionTLS12,
-		MaxVersion:   tls.VersionTLS13,
-		Certificates: []tls.Certificate{cert},
-		//NextProtos:   []string{"h2"},
+		MinVersion: tls.VersionTLS12,
+		MaxVersion: tls.VersionTLS13,
 	}
 
-	l, err := net.Listen("tcp", fmt.Sprintf("%s:%d", ip, port))
-	if err != nil {
-		log.Fatal(err)
+	upstream, _ := url.Parse("http://httpbin.org")
+	router := mux.NewRouter()
+	router.PathPrefix("/").Handler(NewProxy(upstream))
+
+	srv := &http.Server{
+		Addr:           fmt.Sprintf("%s:%d", ip, port),
+		WriteTimeout:   time.Second * 15,
+		ReadTimeout:    time.Second * 15,
+		IdleTimeout:    time.Second * 60,
+		TLSConfig:      tlsConfig,
+		MaxHeaderBytes: 1 << 20,
+		Handler:        router,
 	}
 
-	log.Infof("Listening on %s:%d", ip, port)
-
-	for {
-		conn, err := l.Accept()
-		if err != nil {
-			log.Error(err)
+	// Run our server in a goroutine so that it doesn't block.
+	go func() {
+		if err := srv.ListenAndServeTLS(certificatePath, privateKeyPath); err != nil {
+			log.Fatal(err)
 		}
+	}()
+	log.Infof("Listening on %s:%d", ip, port)
 
-		// TODO: limit open connections
-		go serveConn(conn)
-	}
+	c := make(chan os.Signal, 1)
+	// We'll accept graceful shutdowns when quit via SIGINT (Ctrl+C)
+	// SIGKILL, SIGQUIT or SIGTERM (Ctrl+/) will not be caught.
+	signal.Notify(c, os.Interrupt)
+
+	// Block until we receive our signal.
+	<-c
+
+	// Create a deadline to wait for.
+	ctx, cancel := context.WithTimeout(context.Background(), wait)
+	defer cancel()
+	// Doesn't block if no connections, but will otherwise wait
+	// until the timeout deadline.
+	srv.Shutdown(ctx)
+	// Optionally, you could run srv.Shutdown in a goroutine and block on
+	// <-ctx.Done() if your application should wait for other services
+	// to finalize based on context cancellation.
+	log.Info("shutting down")
+	os.Exit(0)
 }

proxy.go

@@ -0,0 +1,50 @@
+package main
+
+import (
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"strings"
+	"time"
+)
+
+func singleJoiningSlash(a, b string) string {
+	aslash := strings.HasSuffix(a, "/")
+	bslash := strings.HasPrefix(b, "/")
+	switch {
+	case aslash && bslash:
+		return a + b[1:]
+	case !aslash && !bslash:
+		return a + "/" + b
+	}
+	return a + b
+}
+
+func NewProxy(target *url.URL) http.Handler {
+	targetQuery := target.RawQuery
+	director := func(req *http.Request) {
+		req.URL.Scheme = target.Scheme
+		req.URL.Host = target.Host
+		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
+		if targetQuery == "" || req.URL.RawQuery == "" {
+			req.URL.RawQuery = targetQuery + req.URL.RawQuery
+		} else {
+			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
+		}
+		if _, ok := req.Header["User-Agent"]; !ok {
+			// explicitly disable User-Agent so it's not set to default value
+			req.Header.Set("User-Agent", "")
+		}
+	}
+
+	return &httputil.ReverseProxy{
+		FlushInterval: 200 * time.Millisecond,
+		Transport: &http.Transport{
+			TLSHandshakeTimeout:    10 * time.Second,
+			IdleConnTimeout:        30 * time.Second,
+			MaxResponseHeaderBytes: 1 << 20,
+			DisableCompression:     true,
+		},
+		Director: director,
+	}
+}