From 7faa3ce5ad3b197d96462c9f45734d60f3f1fa54 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Fri, 14 Jan 2022 21:12:02 +0800
Subject: [PATCH] try fix use too many port problem

---
 cmd/go-mitmproxy/main.go |  2 +-
 proxy/helper.go          | 11 ++---------
 proxy/middle.go          |  2 ++
 proxy/proxy.go           |  8 ++++++--
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/cmd/go-mitmproxy/main.go b/cmd/go-mitmproxy/main.go
index 8741ee7..f6de2fd 100644
--- a/cmd/go-mitmproxy/main.go
+++ b/cmd/go-mitmproxy/main.go
@@ -12,7 +12,7 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-const version = "0.1.3"
+const version = "0.1.4"
 
 type Config struct {
 	version bool
diff --git a/proxy/helper.go b/proxy/helper.go
index 4b97624..aa8e725 100644
--- a/proxy/helper.go
+++ b/proxy/helper.go
@@ -3,7 +3,6 @@ package proxy
 import (
 	"bytes"
 	"io"
-	"net"
 	"os"
 	"strings"
 	"sync"
@@ -40,11 +39,11 @@ func LogErr(log *_log.Entry, err error) (loged bool) {
 // 转发流量
 // Read a => Write b
 // Read b => Write a
-func Transfer(log *_log.Entry, a, b io.ReadWriteCloser) {
+func Transfer(log *_log.Entry, a, b io.ReadWriter) {
 	done := make(chan struct{})
 	defer close(done)
 
-	forward := func(dst io.WriteCloser, src io.Reader, ec chan<- error) {
+	forward := func(dst io.Writer, src io.Reader, ec chan<- error) {
 		_, err := io.Copy(dst, src)
 		if err != nil {
 			select {
@@ -55,12 +54,6 @@ func Transfer(log *_log.Entry, a, b io.ReadWriteCloser) {
 			}
 		}
 
-		if dstc, ok := dst.(*net.TCPConn); ok {
-			err = dstc.CloseWrite()
-		} else {
-			err = dst.Close()
-		}
-
 		select {
 		case <-done:
 			return
diff --git a/proxy/middle.go b/proxy/middle.go
index 26f656b..620c450 100644
--- a/proxy/middle.go
+++ b/proxy/middle.go
@@ -6,6 +6,7 @@ import (
 	"net"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/lqqyt2423/go-mitmproxy/cert"
 )
@@ -85,6 +86,7 @@ func NewMiddle(proxy *Proxy) (Interceptor, error) {
 
 	server := &http.Server{
 		Handler:      m,
+		IdleTimeout:  5 * time.Second,
 		TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)), // disable http2
 		TLSConfig: &tls.Config{
 			GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) {
diff --git a/proxy/proxy.go b/proxy/proxy.go
index e575d42..ba175c8 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -33,8 +33,9 @@ func NewProxy(opts *Options) (*Proxy, error) {
 	proxy := new(Proxy)
 
 	proxy.Server = &http.Server{
-		Addr:    opts.Addr,
-		Handler: proxy,
+		Addr:        opts.Addr,
+		Handler:     proxy,
+		IdleTimeout: 5 * time.Second,
 	}
 
 	proxy.Client = &http.Client{
@@ -278,6 +279,9 @@ func (proxy *Proxy) handleConnect(res http.ResponseWriter, req *http.Request) {
 		res.WriteHeader(502)
 		return
 	}
+
+	// send RST other than FIN when finished, to avoid TIME_WAIT state
+	cconn.(*net.TCPConn).SetLinger(0)
 	defer cconn.Close()
 
 	_, err = io.WriteString(cconn, "HTTP/1.1 200 Connection Established\r\n\r\n")