From 3df27b9c043d435fca41cd09b7c2f8c8081898af Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Tue, 5 Mar 2019 11:18:17 +0800
Subject: [PATCH] xtcp: wrap yamux on kcp connections

---
 client/proxy/proxy.go | 19 ++++++++++++++++++-
 client/visitor.go     | 19 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/client/proxy/proxy.go b/client/proxy/proxy.go
index 4b721aa7..f08d3ead 100644
--- a/client/proxy/proxy.go
+++ b/client/proxy/proxy.go
@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net"
 	"sync"
 	"time"
@@ -33,6 +34,7 @@ import (
 	"github.com/fatedier/golib/errors"
 	frpIo "github.com/fatedier/golib/io"
 	"github.com/fatedier/golib/pool"
+	fmux "github.com/hashicorp/yamux"
 )
 
 // Proxy defines how to handle work connections for different proxy type.
@@ -302,8 +304,23 @@ func (pxy *XtcpProxy) InWorkConn(conn frpNet.Conn) {
 		return
 	}
 
+	fmuxCfg := fmux.DefaultConfig()
+	fmuxCfg.KeepAliveInterval = 5 * time.Second
+	fmuxCfg.LogOutput = ioutil.Discard
+	sess, err := fmux.Server(kcpConn, fmuxCfg)
+	if err != nil {
+		pxy.Error("create yamux server from kcp connection error: %v", err)
+		return
+	}
+	defer sess.Close()
+	muxConn, err := sess.Accept()
+	if err != nil {
+		pxy.Error("accept for yamux connection error: %v", err)
+		return
+	}
+
 	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf,
-		frpNet.WrapConn(kcpConn), []byte(pxy.cfg.Sk))
+		frpNet.WrapConn(muxConn), []byte(pxy.cfg.Sk))
 }
 
 // UDP
diff --git a/client/visitor.go b/client/visitor.go
index e81f7986..35abcf5a 100644
--- a/client/visitor.go
+++ b/client/visitor.go
@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net"
 	"strconv"
 	"strings"
@@ -35,6 +36,7 @@ import (
 
 	frpIo "github.com/fatedier/golib/io"
 	"github.com/fatedier/golib/pool"
+	fmux "github.com/hashicorp/yamux"
 )
 
 // Visitor is used for forward traffics from local port tot remote service.
@@ -316,7 +318,22 @@ func (sv *XtcpVisitor) handleConn(userConn frpNet.Conn) {
 		remote = frpIo.WithCompression(remote)
 	}
 
-	frpIo.Join(userConn, remote)
+	fmuxCfg := fmux.DefaultConfig()
+	fmuxCfg.KeepAliveInterval = 5 * time.Second
+	fmuxCfg.LogOutput = ioutil.Discard
+	sess, err := fmux.Client(remote, fmuxCfg)
+	if err != nil {
+		sv.Error("create yamux session error: %v", err)
+		return
+	}
+	defer sess.Close()
+	muxConn, err := sess.Open()
+	if err != nil {
+		sv.Error("open yamux stream error: %v", err)
+		return
+	}
+
+	frpIo.Join(userConn, muxConn)
 	sv.Debug("join connections closed")
 }