frp/client/control.go

384 lines
9.3 KiB
Go
Raw Normal View History

2017-03-09 02:03:47 +08:00
// Copyright 2017 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package client
import (
2019-10-12 20:13:12 +08:00
"context"
2019-03-11 14:14:31 +08:00
"crypto/tls"
2017-03-09 02:03:47 +08:00
"io"
2019-10-12 20:13:12 +08:00
"net"
2018-05-11 10:42:57 +08:00
"runtime/debug"
"strconv"
2017-03-09 02:03:47 +08:00
"sync"
"time"
2018-12-09 22:06:22 +08:00
"github.com/fatedier/frp/client/proxy"
2020-09-23 13:49:14 +08:00
"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"
2018-04-25 02:34:07 +08:00
2018-05-08 23:51:13 +08:00
"github.com/fatedier/golib/control/shutdown"
2018-05-08 02:13:30 +08:00
"github.com/fatedier/golib/crypto"
2018-04-25 02:34:07 +08:00
fmux "github.com/hashicorp/yamux"
2017-03-09 02:03:47 +08:00
)
type Control struct {
2018-11-06 18:35:05 +08:00
// uniq id got from frps, attach it in loginMsg
2020-05-24 17:48:37 +08:00
runID string
2017-03-09 02:03:47 +08:00
// manage all proxies
2018-12-07 17:05:36 +08:00
pxyCfgs map[string]config.ProxyConf
2020-05-24 17:48:37 +08:00
pm *proxy.Manager
2017-06-26 03:02:33 +08:00
// manage all visitors
vm *VisitorManager
2017-03-09 02:03:47 +08:00
// control connection
2019-10-12 20:13:12 +08:00
conn net.Conn
2017-03-09 02:03:47 +08:00
// tcp stream multiplexing, if enabled
2018-04-25 02:34:07 +08:00
session *fmux.Session
2017-03-09 02:03:47 +08:00
// put a message in this channel to send it over control connection to server
sendCh chan (msg.Message)
// read from this channel to get the next message sent by server
readCh chan (msg.Message)
// goroutines can block by reading from this channel, it will be closed only in reader() when control connection is closed
2018-11-06 18:35:05 +08:00
closedCh chan struct{}
closedDoneCh chan struct{}
2017-03-09 02:03:47 +08:00
// last time got the Pong message
lastPong time.Time
// The client configuration
clientCfg config.ClientCommonConf
2018-01-17 01:09:33 +08:00
readerShutdown *shutdown.Shutdown
writerShutdown *shutdown.Shutdown
msgHandlerShutdown *shutdown.Shutdown
// The UDP port that the server is listening on
serverUDPPort int
2017-03-09 02:03:47 +08:00
mu sync.RWMutex
2019-10-12 20:13:12 +08:00
xl *xlog.Logger
// service context
ctx context.Context
// sets authentication based on selected method
authSetter auth.Setter
2017-03-09 02:03:47 +08:00
}
2020-05-24 17:48:37 +08:00
func NewControl(ctx context.Context, runID string, conn net.Conn, session *fmux.Session,
2019-10-12 20:13:12 +08:00
clientCfg config.ClientCommonConf,
pxyCfgs map[string]config.ProxyConf,
visitorCfgs map[string]config.VisitorConf,
serverUDPPort int,
authSetter auth.Setter) *Control {
2019-10-12 20:13:12 +08:00
// new xlog instance
2018-01-17 01:09:33 +08:00
ctl := &Control{
2020-05-24 17:48:37 +08:00
runID: runID,
2018-11-06 18:35:05 +08:00
conn: conn,
session: session,
2018-12-07 17:05:36 +08:00
pxyCfgs: pxyCfgs,
sendCh: make(chan msg.Message, 100),
readCh: make(chan msg.Message, 100),
2018-11-06 18:35:05 +08:00
closedCh: make(chan struct{}),
closedDoneCh: make(chan struct{}),
clientCfg: clientCfg,
2018-01-17 01:09:33 +08:00
readerShutdown: shutdown.New(),
writerShutdown: shutdown.New(),
msgHandlerShutdown: shutdown.New(),
serverUDPPort: serverUDPPort,
2019-10-12 20:13:12 +08:00
xl: xlog.FromContextSafe(ctx),
ctx: ctx,
authSetter: authSetter,
2017-03-09 02:03:47 +08:00
}
2020-05-24 17:48:37 +08:00
ctl.pm = proxy.NewManager(ctl.ctx, ctl.sendCh, clientCfg, serverUDPPort)
2018-12-07 17:05:36 +08:00
2019-10-12 20:13:12 +08:00
ctl.vm = NewVisitorManager(ctl.ctx, ctl)
ctl.vm.Reload(visitorCfgs)
2018-01-17 01:09:33 +08:00
return ctl
2017-03-09 02:03:47 +08:00
}
2018-11-06 18:35:05 +08:00
func (ctl *Control) Run() {
2018-01-17 01:09:33 +08:00
go ctl.worker()
2017-06-27 23:30:09 +08:00
2018-12-07 17:05:36 +08:00
// start all proxies
ctl.pm.Reload(ctl.pxyCfgs)
2018-12-07 17:05:36 +08:00
// start all visitors
go ctl.vm.Run()
2018-11-06 18:35:05 +08:00
return
2017-06-26 03:02:33 +08:00
}
2018-01-17 01:09:33 +08:00
func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
2019-10-12 20:13:12 +08:00
xl := ctl.xl
2017-06-26 03:02:33 +08:00
workConn, err := ctl.connectServer()
if err != nil {
return
2017-03-09 02:03:47 +08:00
}
m := &msg.NewWorkConn{
2020-05-24 17:48:37 +08:00
RunID: ctl.runID,
2017-03-09 02:03:47 +08:00
}
if err = ctl.authSetter.SetNewWorkConn(m); err != nil {
xl.Warn("error during NewWorkConn authentication: %v", err)
return
}
2017-03-09 02:03:47 +08:00
if err = msg.WriteMsg(workConn, m); err != nil {
2019-10-12 20:13:12 +08:00
xl.Warn("work connection write to server error: %v", err)
2017-03-09 02:03:47 +08:00
workConn.Close()
return
}
var startMsg msg.StartWorkConn
if err = msg.ReadMsgInto(workConn, &startMsg); err != nil {
2019-10-12 20:13:12 +08:00
xl.Error("work connection closed before response StartWorkConn message: %v", err)
2017-03-09 02:03:47 +08:00
workConn.Close()
return
}
if startMsg.Error != "" {
xl.Error("StartWorkConn contains error: %s", startMsg.Error)
workConn.Close()
return
}
2017-03-09 02:03:47 +08:00
// dispatch this work connection to related proxy
2019-03-29 19:01:18 +08:00
ctl.pm.HandleWorkConn(startMsg.ProxyName, workConn, &startMsg)
2018-01-17 01:09:33 +08:00
}
func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
2019-10-12 20:13:12 +08:00
xl := ctl.xl
2018-01-17 01:09:33 +08:00
// Server will return NewProxyResp message to each NewProxy message.
// Start a new proxy handler if no error got
2018-01-17 14:40:08 +08:00
err := ctl.pm.StartProxy(inMsg.ProxyName, inMsg.RemoteAddr, inMsg.Error)
2018-01-17 01:09:33 +08:00
if err != nil {
2019-10-12 20:13:12 +08:00
xl.Warn("[%s] start error: %v", inMsg.ProxyName, err)
2017-03-13 02:44:47 +08:00
} else {
2019-10-12 20:13:12 +08:00
xl.Info("[%s] start proxy success", inMsg.ProxyName)
2017-03-09 02:03:47 +08:00
}
}
2017-06-27 01:59:30 +08:00
func (ctl *Control) Close() error {
2018-12-11 15:06:54 +08:00
ctl.pm.Close()
2018-12-07 17:05:36 +08:00
ctl.conn.Close()
ctl.vm.Close()
2019-04-08 15:39:14 +08:00
if ctl.session != nil {
ctl.session.Close()
}
2018-01-17 01:09:33 +08:00
return nil
2017-03-09 02:03:47 +08:00
}
2018-11-06 18:35:05 +08:00
// ClosedDoneCh returns a channel which will be closed after all resources are released
func (ctl *Control) ClosedDoneCh() <-chan struct{} {
return ctl.closedDoneCh
2017-03-09 02:03:47 +08:00
}
2018-07-16 01:21:29 +08:00
// connectServer return a new connection to frps
2019-10-12 20:13:12 +08:00
func (ctl *Control) connectServer() (conn net.Conn, err error) {
xl := ctl.xl
2020-05-24 17:48:37 +08:00
if ctl.clientCfg.TCPMux {
2017-06-26 03:02:33 +08:00
stream, errRet := ctl.session.OpenStream()
if errRet != nil {
err = errRet
2019-10-12 20:13:12 +08:00
xl.Warn("start new connection to server error: %v", err)
2017-06-26 03:02:33 +08:00
return
}
2019-10-12 20:13:12 +08:00
conn = stream
2017-06-26 03:02:33 +08:00
} else {
2019-03-11 14:14:31 +08:00
var tlsConfig *tls.Config
sn := ctl.clientCfg.TLSServerName
if sn == "" {
sn = ctl.clientCfg.ServerAddr
}
if ctl.clientCfg.TLSEnable {
2020-09-29 15:44:52 +08:00
tlsConfig, err = transport.NewClientTLSConfig(
ctl.clientCfg.TLSCertFile,
ctl.clientCfg.TLSKeyFile,
2020-09-29 15:44:52 +08:00
ctl.clientCfg.TLSTrustedCaFile,
sn)
if err != nil {
xl.Warn("fail to build tls configuration when connecting to server, err: %v", err)
return
2019-03-11 14:14:31 +08:00
}
}
address := net.JoinHostPort(ctl.clientCfg.ServerAddr, strconv.Itoa(ctl.clientCfg.ServerPort))
conn, err = frpNet.ConnectServerByProxyWithTLS(ctl.clientCfg.HTTPProxy, ctl.clientCfg.Protocol, address, tlsConfig)
2017-06-26 03:02:33 +08:00
if err != nil {
2019-10-12 20:13:12 +08:00
xl.Warn("start new connection to server error: %v", err)
2017-06-26 03:02:33 +08:00
return
}
}
return
}
2018-01-17 01:09:33 +08:00
// reader read all messages from frps and send to readCh
2017-03-09 02:03:47 +08:00
func (ctl *Control) reader() {
2019-10-12 20:13:12 +08:00
xl := ctl.xl
2017-03-09 02:03:47 +08:00
defer func() {
if err := recover(); err != nil {
2019-10-12 20:13:12 +08:00
xl.Error("panic error: %v", err)
xl.Error(string(debug.Stack()))
2017-03-09 02:03:47 +08:00
}
}()
2018-01-17 01:09:33 +08:00
defer ctl.readerShutdown.Done()
2017-03-10 01:42:06 +08:00
defer close(ctl.closedCh)
2017-03-09 02:03:47 +08:00
encReader := crypto.NewReader(ctl.conn, []byte(ctl.clientCfg.Token))
2017-03-09 02:03:47 +08:00
for {
2020-05-24 17:48:37 +08:00
m, err := msg.ReadMsg(encReader)
if err != nil {
2017-03-09 02:03:47 +08:00
if err == io.EOF {
2019-10-12 20:13:12 +08:00
xl.Debug("read from control connection EOF")
2017-03-09 02:03:47 +08:00
return
}
2020-05-24 17:48:37 +08:00
xl.Warn("read error: %v", err)
ctl.conn.Close()
return
2017-03-09 02:03:47 +08:00
}
2020-05-24 17:48:37 +08:00
ctl.readCh <- m
2017-03-09 02:03:47 +08:00
}
}
2018-01-17 01:09:33 +08:00
// writer writes messages got from sendCh to frps
2017-03-09 02:03:47 +08:00
func (ctl *Control) writer() {
2019-10-12 20:13:12 +08:00
xl := ctl.xl
2018-01-17 01:09:33 +08:00
defer ctl.writerShutdown.Done()
encWriter, err := crypto.NewWriter(ctl.conn, []byte(ctl.clientCfg.Token))
2017-03-10 01:42:06 +08:00
if err != nil {
2019-10-12 20:13:12 +08:00
xl.Error("crypto new writer error: %v", err)
2017-03-10 01:42:06 +08:00
ctl.conn.Close()
return
}
2017-03-09 02:03:47 +08:00
for {
2020-05-24 17:48:37 +08:00
m, ok := <-ctl.sendCh
if !ok {
2019-10-12 20:13:12 +08:00
xl.Info("control writer is closing")
2017-03-09 02:03:47 +08:00
return
2020-05-24 17:48:37 +08:00
}
if err := msg.WriteMsg(encWriter, m); err != nil {
xl.Warn("write message to control connection error: %v", err)
return
2017-03-09 02:03:47 +08:00
}
}
}
2018-01-17 01:09:33 +08:00
// msgHandler handles all channel events and do corresponding operations.
func (ctl *Control) msgHandler() {
2019-10-12 20:13:12 +08:00
xl := ctl.xl
2017-03-09 02:03:47 +08:00
defer func() {
if err := recover(); err != nil {
2019-10-12 20:13:12 +08:00
xl.Error("panic error: %v", err)
xl.Error(string(debug.Stack()))
2017-03-09 02:03:47 +08:00
}
}()
2018-01-17 01:09:33 +08:00
defer ctl.msgHandlerShutdown.Done()
2017-03-09 02:03:47 +08:00
2020-11-23 11:38:21 +08:00
hbSend := time.NewTicker(time.Duration(ctl.clientCfg.HeartbeatInterval) * time.Second)
2017-03-09 02:03:47 +08:00
defer hbSend.Stop()
hbCheck := time.NewTicker(time.Second)
defer hbCheck.Stop()
2018-01-17 01:09:33 +08:00
ctl.lastPong = time.Now()
2017-03-09 02:03:47 +08:00
for {
select {
case <-hbSend.C:
// send heartbeat to server
2019-10-12 20:13:12 +08:00
xl.Debug("send heartbeat to server")
pingMsg := &msg.Ping{}
if err := ctl.authSetter.SetPing(pingMsg); err != nil {
xl.Warn("error during ping authentication: %v", err)
return
}
ctl.sendCh <- pingMsg
2017-03-09 02:03:47 +08:00
case <-hbCheck.C:
2020-11-23 11:38:21 +08:00
if time.Since(ctl.lastPong) > time.Duration(ctl.clientCfg.HeartbeatTimeout)*time.Second {
2019-10-12 20:13:12 +08:00
xl.Warn("heartbeat timeout")
2017-03-27 17:25:25 +08:00
// let reader() stop
ctl.conn.Close()
2017-03-09 02:03:47 +08:00
return
}
case rawMsg, ok := <-ctl.readCh:
if !ok {
return
}
switch m := rawMsg.(type) {
case *msg.ReqWorkConn:
2018-01-17 01:09:33 +08:00
go ctl.HandleReqWorkConn(m)
2017-03-09 02:03:47 +08:00
case *msg.NewProxyResp:
2018-01-17 01:09:33 +08:00
ctl.HandleNewProxyResp(m)
2017-03-09 02:03:47 +08:00
case *msg.Pong:
if m.Error != "" {
xl.Error("Pong contains error: %s", m.Error)
ctl.conn.Close()
return
}
2017-03-09 02:03:47 +08:00
ctl.lastPong = time.Now()
2019-10-12 20:13:12 +08:00
xl.Debug("receive heartbeat from server")
2017-03-09 02:03:47 +08:00
}
}
}
}
2018-11-06 18:35:05 +08:00
// If controler is notified by closedCh, reader and writer and handler will exit
2018-01-17 01:09:33 +08:00
func (ctl *Control) worker() {
go ctl.msgHandler()
go ctl.reader()
go ctl.writer()
2018-01-17 01:09:33 +08:00
2018-12-07 17:05:36 +08:00
select {
case <-ctl.closedCh:
// close related channels and wait until other goroutines done
close(ctl.readCh)
ctl.readerShutdown.WaitDone()
ctl.msgHandlerShutdown.WaitDone()
2018-11-06 18:35:05 +08:00
2018-12-07 17:05:36 +08:00
close(ctl.sendCh)
ctl.writerShutdown.WaitDone()
2018-11-06 18:35:05 +08:00
2018-12-07 17:05:36 +08:00
ctl.pm.Close()
ctl.vm.Close()
2018-11-06 18:35:05 +08:00
2018-12-07 17:05:36 +08:00
close(ctl.closedDoneCh)
2019-04-08 15:39:14 +08:00
if ctl.session != nil {
ctl.session.Close()
}
2018-12-07 17:05:36 +08:00
return
2017-03-09 02:03:47 +08:00
}
}
2017-06-27 23:30:09 +08:00
2018-11-06 18:35:05 +08:00
func (ctl *Control) ReloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) error {
ctl.vm.Reload(visitorCfgs)
2018-12-07 17:05:36 +08:00
ctl.pm.Reload(pxyCfgs)
return nil
}