diff --git a/client/control.go b/client/control.go index 9cc139a..1f5e588 100644 --- a/client/control.go +++ b/client/control.go @@ -23,6 +23,7 @@ import ( "github.com/fatedier/frp/models/config" "github.com/fatedier/frp/models/msg" + "github.com/fatedier/frp/utils/crypto" "github.com/fatedier/frp/utils/log" "github.com/fatedier/frp/utils/net" "github.com/fatedier/frp/utils/util" @@ -135,7 +136,7 @@ func (ctl *Control) NewWorkConn() { var startMsg msg.StartWorkConn if err = msg.ReadMsgInto(workConn, &startMsg); err != nil { - ctl.Error("work connection closed and no response from server, %v", err) + ctl.Error("work connection closed, %v", err) workConn.Close() return } @@ -205,12 +206,17 @@ func (ctl *Control) reader() { ctl.Error("panic error: %v", err) } }() + defer close(ctl.closedCh) + encReader, err := crypto.NewReader(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken)) + if err != nil { + ctl.conn.Error("crypto new reader error: %v", err) + return + } for { - if m, err := msg.ReadMsg(ctl.conn); err != nil { + if m, err := msg.ReadMsg(encReader); err != nil { if err == io.EOF { ctl.Debug("read from control connection EOF") - close(ctl.closedCh) return } else { ctl.Warn("read error: %v", err) @@ -223,12 +229,18 @@ func (ctl *Control) reader() { } func (ctl *Control) writer() { + encWriter, err := crypto.NewWriter(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken)) + if err != nil { + ctl.conn.Error("crypto new writer error: %v", err) + ctl.conn.Close() + return + } for { if m, ok := <-ctl.sendCh; !ok { ctl.Info("control writer is closing") return } else { - if err := msg.WriteMsg(ctl.conn, m); err != nil { + if err := msg.WriteMsg(encWriter, m); err != nil { ctl.Warn("write message to control connection error: %v", err) return } diff --git a/client/proxy.go b/client/proxy.go index 5467bc3..a5a400b 100644 --- a/client/proxy.go +++ b/client/proxy.go @@ -23,8 +23,11 @@ import ( "github.com/fatedier/frp/utils/net" ) +// Proxy defines how to work for different proxy type. type Proxy interface { Run() + + // InWorkConn accept work connections registered to server. InWorkConn(conn net.Conn) Close() } diff --git a/server/control.go b/server/control.go index e3027ba..0c2d45c 100644 --- a/server/control.go +++ b/server/control.go @@ -9,6 +9,7 @@ import ( "github.com/fatedier/frp/models/config" "github.com/fatedier/frp/models/consts" "github.com/fatedier/frp/models/msg" + "github.com/fatedier/frp/utils/crypto" "github.com/fatedier/frp/utils/errors" "github.com/fatedier/frp/utils/net" "github.com/fatedier/frp/utils/shutdown" @@ -81,14 +82,14 @@ func NewControl(svr *Service, ctlConn net.Conn, loginMsg *msg.Login) *Control { // Start send a login success message to client and start working. func (ctl *Control) Start() { - go ctl.writer() - - ctl.sendCh <- &msg.LoginResp{ + loginRespMsg := &msg.LoginResp{ Version: version.Full(), RunId: ctl.runId, Error: "", } + msg.WriteMsg(ctl.conn, loginRespMsg) + go ctl.writer() for i := 0; i < ctl.poolCount; i++ { ctl.sendCh <- &msg.ReqWorkConn{} } @@ -182,12 +183,18 @@ func (ctl *Control) writer() { defer ctl.allShutdown.Start() defer ctl.writerShutdown.Done() + encWriter, err := crypto.NewWriter(ctl.conn, []byte(config.ServerCommonCfg.PrivilegeToken)) + if err != nil { + ctl.conn.Error("crypto new writer error: %v", err) + ctl.allShutdown.Start() + return + } for { if m, ok := <-ctl.sendCh; !ok { ctl.conn.Info("control writer is closing") return } else { - if err := msg.WriteMsg(ctl.conn, m); err != nil { + if err := msg.WriteMsg(encWriter, m); err != nil { ctl.conn.Warn("write message to control connection error: %v", err) return } @@ -205,8 +212,13 @@ func (ctl *Control) reader() { defer ctl.allShutdown.Start() defer ctl.readerShutdown.Done() + encReader, err := crypto.NewReader(ctl.conn, []byte(config.ServerCommonCfg.PrivilegeToken)) + if err != nil { + ctl.conn.Error("crypto new reader error: %v", err) + return + } for { - if m, err := msg.ReadMsg(ctl.conn); err != nil { + if m, err := msg.ReadMsg(encReader); err != nil { if err == io.EOF { ctl.conn.Debug("control connection closed") return diff --git a/server/proxy.go b/server/proxy.go index 108a8a9..e751789 100644 --- a/server/proxy.go +++ b/server/proxy.go @@ -165,6 +165,7 @@ func (pxy *UdpProxy) Close() { } // HandleUserTcpConnection is used for incoming tcp user connections. +// It can be used for tcp, http, https type. func HandleUserTcpConnection(pxy Proxy, userConn net.Conn) { defer userConn.Close() ctl := pxy.GetControl() diff --git a/utils/version/version_test.go b/utils/version/version_test.go index a77bf42..ce70dd9 100644 --- a/utils/version/version_test.go +++ b/utils/version/version_test.go @@ -61,5 +61,5 @@ func TestCompact(t *testing.T) { assert.True(ok) ok, _ = Compat("0.10.0") - assert.False(ok) + assert.True(ok) }