upd

This reverts commit 3803993529.

client/main.go

@@ -7,6 +7,7 @@ import (
 	"math/rand"
 	"net"
 	"os"
+	"runtime"
 	"time"
 
 	"golang.org/x/crypto/pbkdf2"
@@ -127,6 +128,11 @@ func main() {
 			Value: 1,
 			Usage: "set num of UDP connections to server",
 		},
+		cli.IntFlag{
+			Name:  "autoexpire",
+			Value: 0,
+			Usage: "set auto expiration time(in seconds) for a single UDP connection, 0 to disable",
+		},
 		cli.IntFlag{
 			Name:  "mtu",
 			Value: 1350,
@@ -254,6 +260,7 @@ func main() {
 		mtu, sndwnd, rcvwnd := c.Int("mtu"), c.Int("sndwnd"), c.Int("rcvwnd")
 		nocomp, acknodelay := c.Bool("nocomp"), c.Bool("acknodelay")
 		dscp, sockbuf, keepalive, conn := c.Int("dscp"), c.Int("sockbuf"), c.Int("keepalive"), c.Int("conn")
+		autoexpire := c.Int("autoexpire")
 
 		log.Println("listening on:", listener.Addr())
 		log.Println("encryption:", crypt)
@@ -268,12 +275,13 @@ func main() {
 		log.Println("sockbuf:", sockbuf)
 		log.Println("keepalive:", keepalive)
 		log.Println("conn:", conn)
+		log.Println("autoexpire:", autoexpire)
 
 		config := &yamux.Config{
 			AcceptBacklog:          256,
 			EnableKeepAlive:        true,
 			KeepAliveInterval:      30 * time.Second,
-			ConnectionWriteTimeout: 30 * time.Second,
+			ConnectionWriteTimeout: 10 * time.Second,
 			MaxStreamWindowSize:    uint32(sockbuf),
 			LogOutput:              os.Stderr,
 		}
@@ -305,27 +313,49 @@ func main() {
 				session, err = yamux.Client(newCompStream(kcpconn), config)
 			}
 			checkError(err)
+			runtime.SetFinalizer(session, func(s *yamux.Session) {
+				s.Close()
+			})
 			return session
 		}
 
 		numconn := uint16(conn)
-		var muxes []*yamux.Session
-		for i := uint16(0); i < numconn; i++ {
-			muxes = append(muxes, createConn())
+		muxes := make([]struct {
+			session *yamux.Session
+			ttl     time.Time
+		}, numconn)
+
+		for k := range muxes {
+			muxes[k].session = createConn()
+			muxes[k].ttl = time.Now().Add(time.Duration(autoexpire) * time.Second)
 		}
 
 		rr := uint16(0)
 		for {
 			p1, err := listener.AcceptTCP()
+			if err := p1.SetReadBuffer(sockbuf); err != nil {
+				log.Println("TCP SetReadBuffer:", err)
+			}
+			if err := p1.SetWriteBuffer(sockbuf); err != nil {
+				log.Println("TCP SetWriteBuffer:", err)
+			}
 			checkError(err)
-			mux := muxes[rr%numconn]
-			p2, err := mux.Open()
+			idx := rr % numconn
+
+		OPEN_P2:
+			// do auto expiration
+			if autoexpire > 0 && time.Now().After(muxes[idx].ttl) {
+				log.Println("autoexpired")
+				muxes[idx].session = createConn()
+				muxes[idx].ttl = time.Now().Add(time.Duration(autoexpire) * time.Second)
+			}
+
+			// do session open
+			p2, err := muxes[idx].session.Open()
 			if err != nil { // yamux failure
-				log.Println(err)
-				p1.Close()
-				muxes[rr%numconn] = createConn()
-				mux.Close()
-				continue
+				muxes[idx].session = createConn()
+				muxes[idx].ttl = time.Now().Add(time.Duration(autoexpire) * time.Second)
+				goto OPEN_P2
 			}
 			go handleClient(p1, p2)
 			rr++

server/main.go

@@ -61,18 +61,26 @@ func handleMux(conn io.ReadWriteCloser, target string, config *yamux.Config) {
 		return
 	}
 	defer mux.Close()
-
 	for {
 		p1, err := mux.Accept()
 		if err != nil {
 			log.Println(err)
 			return
 		}
+		sockbuf := int(config.MaxStreamWindowSize)
 		p2, err := net.DialTimeout("tcp", target, 5*time.Second)
 		if err != nil {
 			log.Println(err)
 			return
 		}
+
+		if err := p2.(*net.TCPConn).SetReadBuffer(sockbuf); err != nil {
+			log.Println("TCP SetReadBuffer:", err)
+		}
+		if err := p2.(*net.TCPConn).SetWriteBuffer(sockbuf); err != nil {
+			log.Println("TCP SetWriteBuffer:", err)
+		}
+
 		go handleClient(p1, p2)
 	}
 }
@@ -291,12 +299,12 @@ func main() {
 			AcceptBacklog:          256,
 			EnableKeepAlive:        true,
 			KeepAliveInterval:      30 * time.Second,
-			ConnectionWriteTimeout: 30 * time.Second,
+			ConnectionWriteTimeout: 10 * time.Second,
 			MaxStreamWindowSize:    uint32(sockbuf),
 			LogOutput:              os.Stderr,
 		}
 		for {
-			if conn, err := lis.Accept(); err == nil {
+			if conn, err := lis.AcceptKCP(); err == nil {
 				log.Println("remote address:", conn.RemoteAddr())
 				conn.SetStreamMode(true)
 				conn.SetNoDelay(nodelay, interval, resend, nc)