support tcp stream multiplexing by smux

This commit is contained in:
fatedier
2017-05-17 17:47:20 +08:00
parent a5f06489cb
commit b600a07ec0
7 changed files with 153 additions and 50 deletions

View File

@ -28,6 +28,7 @@ import (
"github.com/fatedier/frp/utils/net"
"github.com/fatedier/frp/utils/util"
"github.com/fatedier/frp/utils/version"
"github.com/xtaci/smux"
)
const (
@ -50,6 +51,9 @@ type Control struct {
// control connection
conn net.Conn
// tcp stream multiplexing, if enabled
session *smux.Session
// put a message in this channel to send it over control connection to server
sendCh chan (msg.Message)
@ -122,11 +126,25 @@ func (ctl *Control) Run() error {
}
func (ctl *Control) NewWorkConn() {
workConn, err := net.ConnectTcpServerByHttpProxy(config.ClientCommonCfg.HttpProxy,
fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerPort))
if err != nil {
ctl.Warn("start new work connection error: %v", err)
return
var (
workConn net.Conn
err error
)
if config.ClientCommonCfg.TcpMux {
stream, err := ctl.session.OpenStream()
if err != nil {
ctl.Warn("start new work connection error: %v", err)
return
}
workConn = net.WrapConn(stream)
} else {
workConn, err = net.ConnectTcpServerByHttpProxy(config.ClientCommonCfg.HttpProxy,
fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerPort))
if err != nil {
ctl.Warn("start new work connection error: %v", err)
return
}
}
m := &msg.NewWorkConn{
@ -166,6 +184,10 @@ func (ctl *Control) login() (err error) {
if ctl.conn != nil {
ctl.conn.Close()
}
if ctl.session != nil {
ctl.session.Close()
}
conn, err := net.ConnectTcpServerByHttpProxy(config.ClientCommonCfg.HttpProxy,
fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerPort))
if err != nil {
@ -178,6 +200,20 @@ func (ctl *Control) login() (err error) {
}
}()
if config.ClientCommonCfg.TcpMux {
session, errRet := smux.Client(conn, nil)
if errRet != nil {
return errRet
}
stream, errRet := session.OpenStream()
if errRet != nil {
session.Close()
return errRet
}
conn = net.WrapConn(stream)
ctl.session = session
}
now := time.Now().Unix()
ctl.loginMsg.PrivilegeKey = util.GetAuthKey(config.ClientCommonCfg.PrivilegeToken, now)
ctl.loginMsg.Timestamp = now