support protocol quic between frpc and frps (fatedier#3198)

client/control.go

@@ -16,24 +16,18 @@ package client
 
 import (
 	"context"
-	"crypto/tls"
 	"io"
 	"net"
 	"runtime/debug"
-	"strconv"
 	"time"
 
 	"github.com/fatedier/golib/control/shutdown"
 	"github.com/fatedier/golib/crypto"
-	libdial "github.com/fatedier/golib/net/dial"
-	fmux "github.com/hashicorp/yamux"
 
 	"github.com/fatedier/frp/client/proxy"
 	"github.com/fatedier/frp/pkg/auth"
 	"github.com/fatedier/frp/pkg/config"
 	"github.com/fatedier/frp/pkg/msg"
-	"github.com/fatedier/frp/pkg/transport"
-	frpNet "github.com/fatedier/frp/pkg/util/net"
 	"github.com/fatedier/frp/pkg/util/xlog"
 )
 
@@ -51,8 +45,7 @@ type Control struct {
 	// control connection
 	conn net.Conn
 
-	// tcp stream multiplexing, if enabled
-	session *fmux.Session
+	cm *ConnectionManager
 
 	// put a message in this channel to send it over control connection to server
 	sendCh chan (msg.Message)
@@ -87,7 +80,8 @@ type Control struct {
 	authSetter auth.Setter
 }
 
-func NewControl(ctx context.Context, runID string, conn net.Conn, session *fmux.Session,
+func NewControl(
+	ctx context.Context, runID string, conn net.Conn, cm *ConnectionManager,
 	clientCfg config.ClientCommonConf,
 	pxyCfgs map[string]config.ProxyConf,
 	visitorCfgs map[string]config.VisitorConf,
@@ -98,7 +92,7 @@ func NewControl(ctx context.Context, runID string, conn net.Conn, session *fmux.
 	ctl := &Control{
 		runID:              runID,
 		conn:               conn,
-		session:            session,
+		cm:                 cm,
 		pxyCfgs:            pxyCfgs,
 		sendCh:             make(chan msg.Message, 100),
 		readCh:             make(chan msg.Message, 100),
@@ -134,6 +128,7 @@ func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
 	xl := ctl.xl
 	workConn, err := ctl.connectServer()
 	if err != nil {
+		xl.Warn("start new connection to server error: %v", err)
 		return
 	}
 
@@ -189,9 +184,7 @@ func (ctl *Control) GracefulClose(d time.Duration) error {
 	time.Sleep(d)
 
 	ctl.conn.Close()
-	if ctl.session != nil {
-		ctl.session.Close()
-	}
+	ctl.cm.Close()
 	return nil
 }
 
@@ -202,70 +195,7 @@ func (ctl *Control) ClosedDoneCh() <-chan struct{} {
 
 // connectServer return a new connection to frps
 func (ctl *Control) connectServer() (conn net.Conn, err error) {
-	xl := ctl.xl
-	if ctl.clientCfg.TCPMux {
-		stream, errRet := ctl.session.OpenStream()
-		if errRet != nil {
-			err = errRet
-			xl.Warn("start new connection to server error: %v", err)
-			return
-		}
-		conn = stream
-	} else {
-		var tlsConfig *tls.Config
-		sn := ctl.clientCfg.TLSServerName
-		if sn == "" {
-			sn = ctl.clientCfg.ServerAddr
-		}
-
-		if ctl.clientCfg.TLSEnable {
-			tlsConfig, err = transport.NewClientTLSConfig(
-				ctl.clientCfg.TLSCertFile,
-				ctl.clientCfg.TLSKeyFile,
-				ctl.clientCfg.TLSTrustedCaFile,
-				sn)
-
-			if err != nil {
-				xl.Warn("fail to build tls configuration when connecting to server, err: %v", err)
-				return
-			}
-		}
-
-		proxyType, addr, auth, err := libdial.ParseProxyURL(ctl.clientCfg.HTTPProxy)
-		if err != nil {
-			xl.Error("fail to parse proxy url")
-			return nil, err
-		}
-		dialOptions := []libdial.DialOption{}
-		protocol := ctl.clientCfg.Protocol
-		if protocol == "websocket" {
-			protocol = "tcp"
-			dialOptions = append(dialOptions, libdial.WithAfterHook(libdial.AfterHook{Hook: frpNet.DialHookWebsocket()}))
-		}
-		if ctl.clientCfg.ConnectServerLocalIP != "" {
-			dialOptions = append(dialOptions, libdial.WithLocalAddr(ctl.clientCfg.ConnectServerLocalIP))
-		}
-		dialOptions = append(dialOptions,
-			libdial.WithProtocol(protocol),
-			libdial.WithTimeout(time.Duration(ctl.clientCfg.DialServerTimeout)*time.Second),
-			libdial.WithKeepAlive(time.Duration(ctl.clientCfg.DialServerKeepAlive)*time.Second),
-			libdial.WithProxy(proxyType, addr),
-			libdial.WithProxyAuth(auth),
-			libdial.WithTLSConfig(tlsConfig),
-			libdial.WithAfterHook(libdial.AfterHook{
-				Hook: frpNet.DialHookCustomTLSHeadByte(tlsConfig != nil, ctl.clientCfg.DisableCustomTLSFirstByte),
-			}),
-		)
-		conn, err = libdial.Dial(
-			net.JoinHostPort(ctl.clientCfg.ServerAddr, strconv.Itoa(ctl.clientCfg.ServerPort)),
-			dialOptions...,
-		)
-		if err != nil {
-			xl.Warn("start new connection to server error: %v", err)
-			return nil, err
-		}
-	}
-	return
+	return ctl.cm.Connect()
 }
 
 // reader read all messages from frps and send to readCh
@@ -409,9 +339,7 @@ func (ctl *Control) worker() {
 	ctl.vm.Close()
 
 	close(ctl.closedDoneCh)
-	if ctl.session != nil {
-		ctl.session.Close()
-	}
+	ctl.cm.Close()
 }
 
 func (ctl *Control) ReloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) error {