From 1d5123262cf2331f9e63196543c4fa268ec464b4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thomas=20SO=C3=8BTE?= <thomas.soete@corp.ovh.com>
Date: Fri, 21 Dec 2018 11:26:09 +0100
Subject: [PATCH] Add HTTP proxy mode

---
 http_proxy.go | 71 +++++++++++++++++++++++++++++++++++++++++++++++++++
 main.go       | 60 ++++++++++++++++++++++++++++++-------------
 2 files changed, 113 insertions(+), 18 deletions(-)
 create mode 100644 http_proxy.go

diff --git a/http_proxy.go b/http_proxy.go
new file mode 100644
index 0000000..0ddffc1
--- /dev/null
+++ b/http_proxy.go
@@ -0,0 +1,71 @@
+package main
+
+import (
+	"context"
+	"io"
+	"log"
+	"net"
+	"net/http"
+	"time"
+)
+
+type NameResolver interface {
+	Resolve(ctx context.Context, name string) (context.Context, net.IP, error)
+}
+
+type httpProxy struct {
+	Resolver NameResolver
+	Verbose  bool
+	Logger   *log.Logger
+}
+
+func (p *httpProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	_, ip, err := p.Resolver.Resolve(r.Context(), r.URL.Hostname())
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusServiceUnavailable)
+		return
+	}
+	p.Logger.Printf("Command: %s, FQDN: %q, IP: %q, port: %s", r.Method, r.URL.Hostname(), ip.String(), r.URL.Port())
+
+	if r.Method == http.MethodConnect {
+		dest_conn, err := net.DialTimeout("tcp", ip.String()+":"+r.URL.Port(), 10*time.Second)
+		if err != nil {
+			log.Printf("Faild to dial: %s", err.Error())
+			http.Error(w, err.Error(), http.StatusServiceUnavailable)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+		hijacker, ok := w.(http.Hijacker)
+		if !ok {
+			http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
+			return
+		}
+		client_conn, _, err := hijacker.Hijack()
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusServiceUnavailable)
+		}
+		go func() {
+			defer dest_conn.Close()
+			io.Copy(dest_conn, client_conn)
+		}()
+		go func() {
+			defer client_conn.Close()
+			io.Copy(client_conn, dest_conn)
+		}()
+	} else {
+		r.URL.Host = ip.String()
+		resp, err := http.DefaultTransport.RoundTrip(r)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusServiceUnavailable)
+			return
+		}
+		defer resp.Body.Close()
+		for k, vv := range resp.Header {
+			for _, v := range vv {
+				w.Header().Add(k, v)
+			}
+		}
+		w.WriteHeader(resp.StatusCode)
+		io.Copy(w, resp.Body)
+	}
+}
diff --git a/main.go b/main.go
index 2d39bcd..af9345d 100644
--- a/main.go
+++ b/main.go
@@ -1,9 +1,11 @@
 package main
 
 import (
+	"crypto/tls"
 	"flag"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
 	"strings"
 
@@ -19,6 +21,7 @@ var buildDate = "unknown"
 var listen = "127.0.0.1:8000"
 var dnsManglingList stringArray
 var verbose = false
+var mode string
 
 func init() {
 	flag.Usage = func() {
@@ -26,7 +29,8 @@ func init() {
 		flag.PrintDefaults()
 	}
 	flag.StringVar(&listen, "listen", listen, "Address on which the server will listen")
-	flag.Var(&dnsManglingList, "resolve", "Provide a custom address for a specific host and port pair. This option can be used many times")
+	flag.Var(&dnsManglingList, "resolve", "Provide a custom IP address for a specific host. This option can be used many times")
+	flag.StringVar(&mode, "mode", "socks5", "Proxy mode, allowed values: http, socks5")
 	flag.BoolVar(&verbose, "verbose", verbose, "Display access logs")
 	flag.Parse()
 
@@ -63,24 +67,44 @@ func main() {
 		log.Printf(" - Will resolve %q to %q", hostAndIP[0], hostAndIP[1])
 	}
 
-	var rewriter socks5.AddressRewriter
-	if verbose {
-		rewriter = NewRewriterLogger(log.New(os.Stderr, "", log.LstdFlags))
-	}
+	logger := log.New(os.Stderr, "", log.LstdFlags)
 
-	// Instanciate socks proxy
-	conf := &socks5.Config{
-		Resolver: dnsMangler,
-		Rewriter: rewriter,
-	}
-	server, err := socks5.New(conf)
-	if err != nil {
-		log.Fatal(err)
-	}
+	if mode == "socks5" {
+		var rewriter socks5.AddressRewriter
+		if verbose {
+			rewriter = NewRewriterLogger(logger)
+		}
 
-	// Start server
-	log.Printf("Start listening on %q", listen)
-	if err := server.ListenAndServe("tcp", listen); err != nil {
-		log.Fatal(err)
+		// Instanciate socks proxy
+		conf := &socks5.Config{
+			Resolver: dnsMangler,
+			Rewriter: rewriter,
+		}
+		server, err := socks5.New(conf)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		// Start server
+		log.Printf("Start listening on %q - mode: SOCKS5", listen)
+		if err := server.ListenAndServe("tcp", listen); err != nil {
+			log.Fatal(err)
+		}
+	} else if mode == "http" {
+		server := &http.Server{
+			Addr: listen,
+			Handler: &httpProxy{
+				Resolver: dnsMangler,
+				Verbose:  verbose,
+				Logger:   logger,
+			},
+			TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)),
+		}
+		log.Printf("Start listening on %q - mode: HTTP", listen)
+		if err := server.ListenAndServe(); err != nil {
+			log.Fatal(err)
+		}
+	} else {
+		log.Fatalf("Wrong mode: %s", mode)
 	}
 }