optimize timeout of http(s)\socks\nat

Signed-off-by: arraykeys@gmail.com <arraykeys@gmail.com>

services/http.go

@@ -94,7 +94,9 @@ func (s *HTTP) InitService() {
 			for {
 				conn, err := utils.ConnectHost(s.Resolve(*s.cfg.Parent), *s.cfg.Timeout*2)
 				if err == nil {
+					conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 					_, err = conn.Write([]byte{0})
+					conn.SetDeadline(time.Time{})
 				}
 				if err != nil {
 					if s.sshClient != nil {
@@ -215,7 +217,7 @@ func (s *HTTP) OutToTCP(useProxy bool, address string, inConn *net.Conn, req *ut
 			if *s.cfg.ParentType == "ssh" {
 				outConn, err = s.getSSHConn(address)
 			} else {
-				//log.Printf("%v", s.outPool)
+				// log.Printf("%v", s.outPool)
 				_outConn, err = s.outPool.Pool.Get()
 				if err == nil {
 					outConn = _outConn.(net.Conn)
@@ -237,16 +239,16 @@ func (s *HTTP) OutToTCP(useProxy bool, address string, inConn *net.Conn, req *ut
 		utils.CloseConn(inConn)
 		return
 	}
-
 	outAddr := outConn.RemoteAddr().String()
 	//outLocalAddr := outConn.LocalAddr().String()
-
 	if req.IsHTTPS() && (!useProxy || *s.cfg.ParentType == "ssh") {
 		//https无上级或者上级非代理,proxy需要响应connect请求,并直连目标
 		err = req.HTTPSReply()
 	} else {
 		//https或者http,上级是代理,proxy需要转发
+		outConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		_, err = outConn.Write(req.HeadBuf)
+		outConn.SetDeadline(time.Time{})
 		if err != nil {
 			log.Printf("write to %s , err:%s", *s.cfg.Parent, err)
 			utils.CloseConn(inConn)

services/mux_bridge.go

@@ -75,7 +75,9 @@ func (s *MuxBridge) handler(inConn net.Conn) {
 	var err error
 	var connType uint8
 	var key string
+	inConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 	err = utils.ReadPacket(reader, &connType, &key)
+	inConn.SetDeadline(time.Time{})
 	if err != nil {
 		log.Printf("read error,ERR:%s", err)
 		return
@@ -83,7 +85,9 @@ func (s *MuxBridge) handler(inConn net.Conn) {
 	switch connType {
 	case CONN_SERVER:
 		var serverID string
+		inConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		err = utils.ReadPacketData(reader, &serverID)
+		inConn.SetDeadline(time.Time{})
 		if err != nil {
 			log.Printf("read error,ERR:%s", err)
 			return

services/mux_client.go

@@ -68,7 +68,9 @@ func (s *MuxClient) Start(args interface{}) (err error) {
 					time.Sleep(time.Second * 3)
 					continue
 				}
+				conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 				_, err = conn.Write(utils.BuildPacket(CONN_CLIENT, fmt.Sprintf("%s-%d", *s.cfg.Key, i)))
+				conn.SetDeadline(time.Time{})
 				if err != nil {
 					conn.Close()
 					log.Printf("connection err: %s, retrying...", err)
@@ -98,7 +100,9 @@ func (s *MuxClient) Start(args interface{}) (err error) {
 							}
 						}()
 						var ID, clientLocalAddr, serverID string
+						stream.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 						err = utils.ReadPacketData(stream, &ID, &clientLocalAddr, &serverID)
+						stream.SetDeadline(time.Time{})
 						if err != nil {
 							log.Printf("read stream signal err: %s", err)
 							stream.Close()
@@ -140,7 +144,9 @@ func (s *MuxClient) getParentConn() (conn net.Conn, err error) {
 func (s *MuxClient) ServeUDP(inConn *smux.Stream, localAddr, ID string) {
 
 	for {
+		inConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		srcAddr, body, err := utils.ReadUDPPacket(inConn)
+		inConn.SetDeadline(time.Time{})
 		if err != nil {
 			log.Printf("udp packet revecived fail, err: %s", err)
 			log.Printf("connection %s released", ID)
@@ -169,13 +175,16 @@ func (s *MuxClient) processUDPPacket(inConn *smux.Stream, srcAddr, localAddr str
 	}
 	conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 	_, err = conn.Write(body)
+	conn.SetDeadline(time.Time{})
 	if err != nil {
 		log.Printf("send udp packet to %s fail,ERR:%s", dstAddr.String(), err)
 		return
 	}
 	//log.Printf("send udp packet to %s success", dstAddr.String())
 	buf := make([]byte, 1024)
+	conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 	length, _, err := conn.ReadFromUDP(buf)
+	conn.SetDeadline(time.Time{})
 	if err != nil {
 		log.Printf("read udp response from %s fail ,ERR:%s", dstAddr.String(), err)
 		return
@@ -183,7 +192,9 @@ func (s *MuxClient) processUDPPacket(inConn *smux.Stream, srcAddr, localAddr str
 	respBody := buf[0:length]
 	//log.Printf("revecived udp packet from %s , %v", dstAddr.String(), respBody)
 	bs := utils.UDPPacket(srcAddr, respBody)
+	(*inConn).SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 	_, err = (*inConn).Write(bs)
+	(*inConn).SetDeadline(time.Time{})
 	if err != nil {
 		log.Printf("send udp response fail ,ERR:%s", err)
 		inConn.Close()

services/mux_server.go

@@ -227,7 +227,9 @@ func (s *MuxServer) GetOutConn() (outConn net.Conn, ID string, err error) {
 		remoteAddr = "udp:" + *s.cfg.Remote
 	}
 	ID = utils.Uniqueid()
+	outConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 	_, err = outConn.Write(utils.BuildPacketData(ID, remoteAddr, s.cfg.Mgr.serverID))
+	outConn.SetDeadline(time.Time{})
 	if err != nil {
 		log.Printf("write stream data err: %s ,retrying...", err)
 		utils.CloseConn(&outConn)
@@ -253,7 +255,9 @@ func (s *MuxServer) GetConn(index string) (conn net.Conn, err error) {
 		if err != nil {
 			return
 		}
+		c.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		_, err = c.Write(utils.BuildPacket(CONN_SERVER, *s.cfg.Key, s.cfg.Mgr.serverID))
+		c.SetDeadline(time.Time{})
 		if err != nil {
 			c.Close()
 			return
@@ -327,7 +331,9 @@ func (s *MuxServer) UDPConnDeamon() {
 								// outConn.Close()
 							}()
 							for {
+								outConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 								srcAddrFromConn, body, err := utils.ReadUDPPacket(outConn)
+								outConn.SetDeadline(time.Time{})
 								if err != nil {
 									log.Printf("parse revecived udp packet fail, err: %s ,%v", err, body)
 									log.Printf("UDP deamon connection %s exited", ID)
@@ -341,7 +347,9 @@ func (s *MuxServer) UDPConnDeamon() {
 								}
 								port, _ := strconv.Atoi(_srcAddr[1])
 								dstAddr := &net.UDPAddr{IP: net.ParseIP(_srcAddr[0]), Port: port}
+								s.sc.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 								_, err = s.sc.UDPListener.WriteToUDP(body, dstAddr)
+								s.sc.UDPListener.SetDeadline(time.Time{})
 								if err != nil {
 									log.Printf("udp response to local %s fail,ERR:%s", srcAddrFromConn, err)
 									continue

services/socks.go

@@ -104,7 +104,9 @@ func (s *Socks) InitService() {
 			for {
 				conn, err := utils.ConnectHost(s.Resolve(*s.cfg.Parent), *s.cfg.Timeout*2)
 				if err == nil {
+					conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 					_, err = conn.Write([]byte{0})
+					conn.SetDeadline(time.Time{})
 				}
 				if err != nil {
 					if s.sshClient != nil {
@@ -216,6 +218,7 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 		}
 		conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout*5)))
 		_, err = conn.Write(rawB)
+		conn.SetDeadline(time.Time{})
 		log.Printf("udp request:%v", len(rawB))
 		if err != nil {
 			log.Printf("send udp packet to %s fail,ERR:%s", dstAddr.String(), err)
@@ -225,7 +228,9 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 
 		//log.Printf("send udp packet to %s success", dstAddr.String())
 		buf := make([]byte, 10*1024)
+		conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		length, _, err := conn.ReadFromUDP(buf)
+		conn.SetDeadline(time.Time{})
 		if err != nil {
 			log.Printf("read udp response from %s fail ,ERR:%s", dstAddr.String(), err)
 			conn.Close()
@@ -250,10 +255,14 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 				conn.Close()
 				return
 			}
+			s.udpSC.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			s.udpSC.UDPListener.WriteToUDP(d, srcAddr)
+			s.udpSC.UDPListener.SetDeadline(time.Time{})
 			log.Printf("udp reply:%v", len(d))
 		} else {
+			s.udpSC.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			s.udpSC.UDPListener.WriteToUDP(respBody, srcAddr)
+			s.udpSC.UDPListener.SetDeadline(time.Time{})
 			log.Printf("udp reply:%v", len(respBody))
 		}
 
@@ -272,6 +281,7 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 		}
 		conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout*3)))
 		_, err = conn.Write(p.Data())
+		conn.SetDeadline(time.Time{})
 		log.Printf("udp send:%v", len(p.Data()))
 		if err != nil {
 			log.Printf("send udp packet to %s fail,ERR:%s", dstAddr.String(), err)
@@ -280,7 +290,10 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 		}
 		//log.Printf("send udp packet to %s success", dstAddr.String())
 		buf := make([]byte, 10*1024)
+		conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		length, _, err := conn.ReadFromUDP(buf)
+		conn.SetDeadline(time.Time{})
+
 		if err != nil {
 			log.Printf("read udp response from %s fail ,ERR:%s", dstAddr.String(), err)
 			conn.Close()
@@ -297,9 +310,13 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 				conn.Close()
 				return
 			}
+			s.udpSC.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			s.udpSC.UDPListener.WriteToUDP(d, srcAddr)
+			s.udpSC.UDPListener.SetDeadline(time.Time{})
 		} else {
+			s.udpSC.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			s.udpSC.UDPListener.WriteToUDP(respPacket, srcAddr)
+			s.udpSC.UDPListener.SetDeadline(time.Time{})
 		}
 		log.Printf("udp reply:%v", len(respPacket))
 	}
@@ -371,9 +388,15 @@ func (s *Socks) socksConnCallback(inConn net.Conn) {
 		//auth
 		_addr := strings.Split(inConn.RemoteAddr().String(), ":")
 		if s.basicAuth.CheckUserPass(user, pass, _addr[0], "") {
+			inConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			inConn.Write([]byte{0x01, 0x00})
+			inConn.SetDeadline(time.Time{})
+
 		} else {
+			inConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			inConn.Write([]byte{0x01, 0x01})
+			inConn.SetDeadline(time.Time{})
+
 			utils.CloseConn(&inConn)
 			return
 		}
@@ -496,25 +519,32 @@ func (s *Socks) getOutConn(methodBytes, reqBytes []byte, host string) (outConn n
 		}
 		var buf = make([]byte, 1024)
 		//var n int
+		outConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		_, err = outConn.Write(methodBytes)
+		outConn.SetDeadline(time.Time{})
 		if err != nil {
 			err = fmt.Errorf("write method fail,%s", err)
 			return
 		}
+		outConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		_, err = outConn.Read(buf)
+		outConn.SetDeadline(time.Time{})
 		if err != nil {
 			err = fmt.Errorf("read method reply fail,%s", err)
 			return
 		}
 		//resp := buf[:n]
 		//log.Printf("resp:%v", resp)
-
+		outConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		_, err = outConn.Write(reqBytes)
+		outConn.SetDeadline(time.Time{})
 		if err != nil {
 			err = fmt.Errorf("write req detail fail,%s", err)
 			return
 		}
+		outConn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 		_, err = outConn.Read(buf)
+		outConn.SetDeadline(time.Time{})
 		if err != nil {
 			err = fmt.Errorf("read req reply fail,%s", err)
 			return