From c3ed4478d981216e9c10f43b0f4b2c73d57a8f63 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Sat, 28 Nov 2020 22:12:54 +0800
Subject: [PATCH] simple http handler

---
 README.md   |  6 +++++
 bin/main.go |  9 ++++---
 proxy.go    | 71 +++++++++++++++++++++++++++++++++++------------------
 3 files changed, 59 insertions(+), 27 deletions(-)
 create mode 100644 README.md

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..b673390
--- /dev/null
+++ b/README.md
@@ -0,0 +1,6 @@
+# go-mitmproxy
+
+## TODO
+
+- [x] http handler
+- [ ] https handler
diff --git a/bin/main.go b/bin/main.go
index 42b79ca..b086a7a 100644
--- a/bin/main.go
+++ b/bin/main.go
@@ -1,8 +1,11 @@
 package main
 
-import proxy "github.com/lqqyt2423/go-mitmproxy"
+import (
+	"log"
+
+	proxy "github.com/lqqyt2423/go-mitmproxy"
+)
 
 func main() {
-	proxy.Create()
-	proxy.Init()
+	log.Fatal(proxy.NewProxy().Start())
 }
diff --git a/proxy.go b/proxy.go
index be7faf6..b9da61b 100644
--- a/proxy.go
+++ b/proxy.go
@@ -7,36 +7,59 @@ import (
 	"time"
 )
 
-func Create() {
-	http.HandleFunc("/", func(res http.ResponseWriter, req *http.Request) {
-		if !req.URL.IsAbs() || req.URL.Host == "" {
-			res.WriteHeader(400)
-			io.WriteString(res, "此为代理服务器，不能直接发起请求")
-			return
+type Proxy struct {
+	Server *http.Server
+}
+
+func (proxy *Proxy) Start() error {
+	log.Printf("Proxy start listen at :8080")
+	return proxy.Server.ListenAndServe()
+}
+
+func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
+	if !req.URL.IsAbs() || req.URL.Host == "" {
+		res.WriteHeader(400)
+		_, err := io.WriteString(res, "此为代理服务器，不能直接发起请求")
+		if err != nil {
+			log.Printf("error: %v", err)
 		}
+		return
+	}
 
-		start := time.Now()
+	start := time.Now()
 
-		proxyReq, _ := http.NewRequest(req.Method, req.URL.String(), req.Body)
+	proxyReq, _ := http.NewRequest(req.Method, req.URL.String(), req.Body)
 
-		// TODO: handle Proxy- header
-		for key, value := range req.Header {
-			proxyReq.Header[key] = value
-		}
-		proxyRes, _ := http.DefaultClient.Do(proxyReq)
-		defer proxyRes.Body.Close()
+	// TODO: handle Proxy- header
+	for key, value := range req.Header {
+		proxyReq.Header[key] = value
+	}
+	proxyRes, _ := http.DefaultClient.Do(proxyReq)
 
-		for key, value := range proxyRes.Header {
-			res.Header()[key] = value
-		}
-		res.WriteHeader(proxyRes.StatusCode)
-		io.Copy(res, proxyRes.Body)
+	for key, value := range proxyRes.Header {
+		res.Header()[key] = value
+	}
+	res.WriteHeader(proxyRes.StatusCode)
+	_, err := io.Copy(res, proxyRes.Body)
+	if err != nil {
+		log.Printf("error: %v", err)
+		return
+	}
+
+	err = proxyRes.Body.Close()
+	if err != nil {
+		log.Printf("error: %v", err)
+		return
+	}
 
-		log.Printf("%v %v %v - %v ms", req.Method, req.URL.String(), proxyRes.StatusCode, time.Since(start))
-	})
+	log.Printf("%v %v %v - %v ms", req.Method, req.URL.String(), proxyRes.StatusCode, time.Since(start).Milliseconds())
 }
 
-func Init() {
-	log.Println("server begin listen at :8000")
-	log.Fatal(http.ListenAndServe(":8000", nil))
+func NewProxy() *Proxy {
+	proxy := new(Proxy)
+	proxy.Server = &http.Server{
+		Addr:    ":8080",
+		Handler: proxy,
+	}
+	return proxy
 }