转发https逻辑抽取出来

cmd/mitmproxy/main.go

@@ -12,5 +12,11 @@ func main() {
 	opts := &proxy.Options{
 		Addr: ":9080",
 	}
-	log.Fatal(proxy.NewProxy(opts).Start())
+
+	p, err := proxy.NewProxy(opts)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	log.Fatal(p.Start())
 }

proxy/mitm.go

@@ -0,0 +1,89 @@
+package proxy
+
+import (
+	"crypto/tls"
+	"log"
+	"net"
+	"net/http"
+
+	"github.com/lqqyt2423/go-mitmproxy/cert"
+)
+
+type Mitm interface {
+	Start() error
+	Dial(host string) (net.Conn, error)
+}
+
+// 直接转发 https 流量
+type MitmForward struct{}
+
+func (m *MitmForward) Start() error {
+	return nil
+}
+
+func (m *MitmForward) Dial(host string) (net.Conn, error) {
+	return net.Dial("tcp", host)
+}
+
+// 内部解析 https 流量
+type MitmServer struct {
+	Proxy    *Proxy
+	CA       *cert.CA
+	Listener net.Listener
+	Server   *http.Server
+}
+
+func NewMitmServer(proxy *Proxy) (Mitm, error) {
+	ca, err := cert.NewCA("")
+	if err != nil {
+		return nil, err
+	}
+
+	m := &MitmServer{
+		Proxy: proxy,
+		CA:    ca,
+	}
+
+	server := &http.Server{
+		Handler: m,
+		TLSConfig: &tls.Config{
+			GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) {
+				// log.Printf("MitmServer GetCertificate ServerName: %v\n", chi.ServerName)
+				return ca.DummyCert(chi.ServerName)
+			},
+		},
+	}
+
+	m.Server = server
+
+	return m, nil
+}
+
+func (m *MitmServer) Start() error {
+	ln, err := net.Listen("tcp", "127.0.0.1:") // port number is automatically chosen
+	if err != nil {
+		return err
+	}
+	m.Listener = ln
+	m.Server.Addr = ln.Addr().String()
+	log.Printf("MitmServer Server Addr is %v\n", m.Server.Addr)
+	defer ln.Close()
+
+	return m.Server.ServeTLS(ln, "", "")
+}
+
+func (m *MitmServer) Dial(host string) (net.Conn, error) {
+	return net.Dial("tcp", m.Server.Addr)
+}
+
+func (m *MitmServer) ServeHTTP(res http.ResponseWriter, req *http.Request) {
+	if req.URL.Scheme == "" {
+		req.URL.Scheme = "https"
+	}
+
+	if req.URL.Host == "" {
+		req.URL.Host = req.Host
+	}
+
+	m.Proxy.ServeHTTP(res, req)
+}

proxy/proxy.go

@@ -1,14 +1,10 @@
 package proxy
 
 import (
-	"crypto/tls"
 	"io"
 	"log"
-	"net"
 	"net/http"
 	"time"
-
-	"github.com/lqqyt2423/go-mitmproxy/cert"
 )
 
 type Options struct {
@@ -17,23 +13,16 @@ type Options struct {
 
 type Proxy struct {
 	Server *http.Server
-
-	ca               *cert.CA
-	extraNetListener net.Listener
-	extraServer      *http.Server
+	Mitm   Mitm
 }
 
 func (proxy *Proxy) Start() error {
-	ln, err := net.Listen("tcp", "127.0.0.1:") // port number is automatically chosen
-	if err != nil {
-		return err
-	}
-	proxy.extraNetListener = ln
-	proxy.extraServer.Addr = ln.Addr().String()
-	log.Printf("Proxy extraServer Addr is %v\n", proxy.extraServer.Addr)
 	go func() {
-		defer ln.Close()
-		log.Fatal(proxy.extraServer.ServeTLS(ln, "", ""))
+		err := proxy.Mitm.Start()
+		if err != nil {
+			// TODO
+			log.Fatal(err)
+		}
 	}()
 
 	log.Printf("Proxy start listen at %v\n", proxy.Server.Addr)
@@ -46,8 +35,6 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	log.Printf("url: %v\n", req.URL.String())
-
 	if !req.URL.IsAbs() || req.URL.Host == "" {
 		res.WriteHeader(400)
 		_, err := io.WriteString(res, "此为代理服务器，不能直接发起请求")
@@ -92,13 +79,9 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 }
 
 func (proxy *Proxy) handleConnect(res http.ResponseWriter, req *http.Request) {
-	log.Printf("CONNECT: %v\n", req.Host)
-
-	// 直接转发
-	// conn, err := net.Dial("tcp", req.Host)
+	// log.Printf("CONNECT: %v\n", req.Host)
 
-	// 内部解析 HTTPS
-	conn, err := net.Dial("tcp", proxy.extraServer.Addr)
+	conn, err := proxy.Mitm.Dial(req.Host)
 
 	if err != nil {
 		log.Printf("error: %v, host: %v\n", err, req.Host)
@@ -138,28 +121,19 @@ func (proxy *Proxy) handleConnect(res http.ResponseWriter, req *http.Request) {
 	<-ch
 }
 
-func NewProxy(opts *Options) *Proxy {
+func NewProxy(opts *Options) (*Proxy, error) {
 	proxy := new(Proxy)
 	proxy.Server = &http.Server{
 		Addr:    opts.Addr,
 		Handler: proxy,
 	}
 
-	ca, err := cert.NewCA("")
+	mitm, err := NewMitmServer(proxy)
 	if err != nil {
-		panic(err)
-	}
-	proxy.ca = ca
-	proxy.extraServer = &http.Server{
-		Handler: proxy,
-		TLSConfig: &tls.Config{
-			PreferServerCipherSuites: true,
-			GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) {
-				log.Printf("GetCertificate ServerName: %v\n", chi.ServerName)
-				return proxy.ca.DummyCert(chi.ServerName)
-			},
-		},
+		return nil, err
 	}
 
-	return proxy
+	proxy.Mitm = mitm
+
+	return proxy, nil
 }