add MWSSSessionDeadLine

Ehco1996 · Jun 23, 2020 · 3695a0e · 3695a0e
1 parent 73d1006
commit 3695a0e
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 26 deletions.
diff --git a/cmd/main.go b/cmd/main.go
@@ -1,11 +1,12 @@
 package main
 
 import (
+	cli "github.com/urfave/cli/v2"
 	"log"
+	"net/http"
+	_ "net/http/pprof"
 	"os"
 
-	cli "github.com/urfave/cli/v2"
-
 	relay "github.com/Ehco1996/ehco/internal/relay"
 	web "github.com/Ehco1996/ehco/internal/web"
 )
@@ -15,12 +16,13 @@ var ListenType string
 var RemoteAddr string
 var TransportType string
 var ConfigPath string
+var PprofPort string
 
 func main() {
 	app := cli.NewApp()
 	app.Name = "ehco"
-	app.Version = "0.0.8"
-	app.Usage = "A proxy used to relay tcp/udp traffic to anywhere"
+	app.Version = "0.1.0"
+	app.Usage = "ehco is a network relay tool and a typo :)"
 	app.Flags = []cli.Flag{
 		&cli.StringFlag{
 			Name:        "l, local",
@@ -57,6 +59,12 @@ func main() {
 			Usage:       "配置文件地址",
 			Destination: &ConfigPath,
 		},
+		&cli.StringFlag{
+			Name:        "pport",
+			Usage:       "pprof监听端口",
+			EnvVars:     []string{"EHCO_PPROF_PORT"},
+			Destination: &PprofPort,
+		},
 	}
 
 	app.Action = start
@@ -81,6 +89,15 @@ func start(ctx *cli.Context) error {
 	} else {
 		go serveRelay(LocalAddr, ListenType, RemoteAddr, TransportType, ch)
 	}
+
+	if PprofPort != "" {
+		go func() {
+			pps := "0.0.0.0:" + PprofPort
+			log.Printf("[DEBUG] start pprof server at http://%s/debug/pprof/", pps)
+			log.Println(http.ListenAndServe(pps	, nil))
+		}()
+	}
+
 	return <-ch
 }
 

diff --git a/internal/relay/mwss.go b/internal/relay/mwss.go
@@ -93,7 +93,7 @@ func (tr *mwssTransporter) Dial(addr string) (conn net.Conn, err error) {
 
 	// 找到可以用的session
 	for _, s := range sessions {
-		if s.session.NumStreams() >= s.maxStreamCnt {
+		if s.NumStreams() >= s.maxStreamCnt {
 			ok = false
 		} else {
 			session = s
@@ -130,7 +130,7 @@ func (tr *mwssTransporter) Dial(addr string) (conn net.Conn, err error) {
 			return nil, err
 		}
 		conn.SetDeadline(time.Now().Add(WsDeadline))
-		defer conn.SetDeadline(time.Time{})
+
 		session, err = tr.initSession(addr, conn)
 		if err != nil {
 			conn.Close()
@@ -144,6 +144,9 @@ func (tr *mwssTransporter) Dial(addr string) (conn net.Conn, err error) {
 		session.Close()
 		return nil, err
 	}
+	// TODO 统一管理session的deadline
+	session.conn.SetDeadline(time.Now().Add(MWSSSessionDeadLine))
+	session.session.SetDeadline(time.Now().Add(MWSSSessionDeadLine))
 	tr.sessions[addr] = sessions
 	return cc, nil
 }
@@ -184,15 +187,15 @@ func (r *Relay) RunLocalMWSSServer() error {
 	}
 
 	mux := http.NewServeMux()
+	mux.Handle("/tcp/", http.HandlerFunc(s.upgrade))
+	// fake
+	mux.Handle("/", http.HandlerFunc(index))
 	server := &http.Server{
 		Addr:              r.LocalTCPAddr.String(),
 		Handler:           mux,
 		TLSConfig:         DefaultTLSConfig,
 		ReadHeaderTimeout: 30 * time.Second,
 	}
-	mux.Handle("/tcp/", http.HandlerFunc(s.upgrade))
-	// fake
-	mux.Handle("/", http.HandlerFunc(index))
 	s.server = server
 
 	ln, err := net.Listen("tcp", r.LocalTCPAddr.String())
@@ -305,6 +308,7 @@ func (r *Relay) handleTcpOverMWSS(c *net.TCPConn) error {
 		return err
 	}
 	defer wsc.Close()
+	log.Printf("handleTcpOverMWSS from:%s to:%s", c.RemoteAddr(), wsc.RemoteAddr())
 	transport(wsc, c)
 	return nil
 }

diff --git a/internal/relay/relay.go b/internal/relay/relay.go
@@ -4,19 +4,16 @@ import (
 	"io"
 	"log"
 	"net"
-	"net/http"
-	_ "net/http/pprof"
-	"os"
 	"time"
 )
 
 var (
-	TcpDeadline       = 60 * time.Second
-	UdpDeadline       = 6 * time.Second
-	WsDeadline        = 15 * time.Second
-	FastCloseDeadLine = 1 * time.Second
-	DEBUG             = os.Getenv("EHCO_DEBUG")
-	MaxMWSSStreamCnt  = 60
+	TcpDeadline         = 60 * time.Second
+	UdpDeadline         = 6 * time.Second
+	WsDeadline          = 15 * time.Second
+	FastCloseDeadLine   = 1 * time.Second
+	MaxMWSSStreamCnt    = 10
+	MWSSSessionDeadLine = 600 * time.Second
 )
 
 const (
@@ -67,12 +64,6 @@ func NewRelay(localAddr, listenType, remoteAddr, transportType string) (*Relay,
 
 		udpCache: make(map[string](*udpBufferCh)),
 	}
-	if DEBUG != "" {
-		go func() {
-			log.Printf("[DEBUG] start pprof server at http://127.0.0.1:6060/debug/pprof/")
-			log.Println(http.ListenAndServe("0.0.0.0:6060", nil))
-		}()
-	}
 
 	if listenType == Listen_WSS || transportType == Transport_WSS ||
 		listenType == Listen_MWSS || transportType == Transport_MWSS {
@@ -117,10 +108,9 @@ func (r *Relay) RunLocalTCPServer() error {
 	for {
 		c, err := r.TCPListener.AcceptTCP()
 		if err != nil {
+			log.Printf("accept tcp con error: %s", err)
 			return err
 		}
-		log.Printf("handle tcp con from: %s over: %s", c.RemoteAddr(), r.TransportType)
-
 		switch r.TransportType {
 		case Transport_WSS:
 			go func(c *net.TCPConn) {