frp/client/control.go

296 lines
8.8 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"
"net"
"sync/atomic"
2017-03-09 02:03:47 +08:00
"time"
"github.com/samber/lo"
2022-08-29 01:02:53 +08:00
2018-12-09 22:06:22 +08:00
"github.com/fatedier/frp/client/proxy"
"github.com/fatedier/frp/client/visitor"
2020-09-23 13:49:14 +08:00
"github.com/fatedier/frp/pkg/auth"
v1 "github.com/fatedier/frp/pkg/config/v1"
2020-09-23 13:49:14 +08:00
"github.com/fatedier/frp/pkg/msg"
"github.com/fatedier/frp/pkg/transport"
2023-11-27 15:47:49 +08:00
netpkg "github.com/fatedier/frp/pkg/util/net"
"github.com/fatedier/frp/pkg/util/wait"
2020-09-23 13:49:14 +08:00
"github.com/fatedier/frp/pkg/util/xlog"
2017-03-09 02:03:47 +08:00
)
2023-11-27 15:47:49 +08:00
type SessionContext struct {
// The client common configuration.
Common *v1.ClientCommonConfig
// Unique ID obtained from frps.
// It should be attached to the login message when reconnecting.
RunID string
// Underlying control connection. Once conn is closed, the msgDispatcher and the entire Control will exit.
Conn net.Conn
// Indicates whether the connection is encrypted.
ConnEncrypted bool
// Sets authentication based on selected method
AuthSetter auth.Setter
// Connector is used to create new connections, which could be real TCP connections or virtual streams.
Connector Connector
}
2017-03-09 02:03:47 +08:00
type Control struct {
// service context
ctx context.Context
xl *xlog.Logger
2023-11-27 15:47:49 +08:00
// session context
sessionCtx *SessionContext
2017-03-09 02:03:47 +08:00
// manage all proxies
2023-11-27 15:47:49 +08:00
pm *proxy.Manager
2017-06-26 03:02:33 +08:00
// manage all visitors
vm *visitor.Manager
doneCh chan struct{}
2017-03-09 02:03:47 +08:00
// of time.Time, last time got the Pong message
lastPong atomic.Value
// The role of msgTransporter is similar to HTTP2.
// It allows multiple messages to be sent simultaneously on the same control connection.
// The server's response messages will be dispatched to the corresponding waiting goroutines based on the laneKey and message type.
msgTransporter transport.MessageTransporter
// msgDispatcher is a wrapper for control connection.
// It provides a channel for sending messages, and you can register handlers to process messages based on their respective types.
msgDispatcher *msg.Dispatcher
2017-03-09 02:03:47 +08:00
}
2023-11-27 15:47:49 +08:00
func NewControl(ctx context.Context, sessionCtx *SessionContext) (*Control, error) {
2019-10-12 20:13:12 +08:00
// new xlog instance
2018-01-17 01:09:33 +08:00
ctl := &Control{
ctx: ctx,
xl: xlog.FromContextSafe(ctx),
2023-11-27 15:47:49 +08:00
sessionCtx: sessionCtx,
doneCh: make(chan struct{}),
}
ctl.lastPong.Store(time.Now())
2023-11-27 15:47:49 +08:00
if sessionCtx.ConnEncrypted {
cryptoRW, err := netpkg.NewCryptoReadWriter(sessionCtx.Conn, []byte(sessionCtx.Common.Auth.Token))
if err != nil {
return nil, err
}
ctl.msgDispatcher = msg.NewDispatcher(cryptoRW)
} else {
ctl.msgDispatcher = msg.NewDispatcher(sessionCtx.Conn)
2017-03-09 02:03:47 +08:00
}
ctl.registerMsgHandlers()
ctl.msgTransporter = transport.NewMessageTransporter(ctl.msgDispatcher.SendChannel())
2023-11-27 15:47:49 +08:00
ctl.pm = proxy.NewManager(ctl.ctx, sessionCtx.Common, ctl.msgTransporter)
ctl.vm = visitor.NewManager(ctl.ctx, sessionCtx.RunID, sessionCtx.Common, ctl.connectServer, ctl.msgTransporter)
return ctl, nil
2017-03-09 02:03:47 +08:00
}
2023-11-27 15:47:49 +08:00
func (ctl *Control) Run(proxyCfgs []v1.ProxyConfigurer, visitorCfgs []v1.VisitorConfigurer) {
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
2023-11-27 15:47:49 +08:00
ctl.pm.UpdateAll(proxyCfgs)
2018-12-07 17:05:36 +08:00
// start all visitors
2023-11-27 15:47:49 +08:00
ctl.vm.UpdateAll(visitorCfgs)
2017-06-26 03:02:33 +08:00
}
2023-11-21 11:19:35 +08:00
func (ctl *Control) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
ctl.pm.SetInWorkConnCallback(cb)
}
func (ctl *Control) handleReqWorkConn(_ msg.Message) {
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 {
xl.Warn("start new connection to server error: %v", err)
2017-06-26 03:02:33 +08:00
return
2017-03-09 02:03:47 +08:00
}
m := &msg.NewWorkConn{
2023-11-27 15:47:49 +08:00
RunID: ctl.sessionCtx.RunID,
2017-03-09 02:03:47 +08:00
}
2023-11-27 15:47:49 +08:00
if err = ctl.sessionCtx.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 {
2022-12-18 18:43:42 +08:00
xl.Trace("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(m msg.Message) {
2019-10-12 20:13:12 +08:00
xl := ctl.xl
inMsg := m.(*msg.NewProxyResp)
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
}
}
func (ctl *Control) handleNatHoleResp(m msg.Message) {
xl := ctl.xl
inMsg := m.(*msg.NatHoleResp)
// Dispatch the NatHoleResp message to the related proxy.
ok := ctl.msgTransporter.DispatchWithType(inMsg, msg.TypeNameNatHoleResp, inMsg.TransactionID)
if !ok {
xl.Trace("dispatch NatHoleResp message to related proxy error")
}
}
func (ctl *Control) handlePong(m msg.Message) {
xl := ctl.xl
inMsg := m.(*msg.Pong)
if inMsg.Error != "" {
xl.Error("Pong message contains error: %s", inMsg.Error)
2023-11-27 15:47:49 +08:00
ctl.closeSession()
return
}
ctl.lastPong.Store(time.Now())
xl.Debug("receive heartbeat from server")
}
2023-11-27 15:47:49 +08:00
// closeSession closes the control connection.
func (ctl *Control) closeSession() {
ctl.sessionCtx.Conn.Close()
ctl.sessionCtx.Connector.Close()
}
2017-06-27 01:59:30 +08:00
func (ctl *Control) Close() error {
return ctl.GracefulClose(0)
}
2021-10-19 15:02:45 +08:00
func (ctl *Control) GracefulClose(d time.Duration) error {
2018-12-11 15:06:54 +08:00
ctl.pm.Close()
ctl.vm.Close()
2021-10-19 15:02:45 +08:00
time.Sleep(d)
2023-11-27 15:47:49 +08:00
ctl.closeSession()
2018-01-17 01:09:33 +08:00
return nil
2017-03-09 02:03:47 +08:00
}
// Done returns a channel that will be closed after all resources are released
func (ctl *Control) Done() <-chan struct{} {
return ctl.doneCh
2017-03-09 02:03:47 +08:00
}
2018-07-16 01:21:29 +08:00
// connectServer return a new connection to frps
2023-11-27 15:47:49 +08:00
func (ctl *Control) connectServer() (net.Conn, error) {
return ctl.sessionCtx.Connector.Connect()
2017-06-26 03:02:33 +08:00
}
func (ctl *Control) registerMsgHandlers() {
ctl.msgDispatcher.RegisterHandler(&msg.ReqWorkConn{}, msg.AsyncHandler(ctl.handleReqWorkConn))
ctl.msgDispatcher.RegisterHandler(&msg.NewProxyResp{}, ctl.handleNewProxyResp)
ctl.msgDispatcher.RegisterHandler(&msg.NatHoleResp{}, ctl.handleNatHoleResp)
ctl.msgDispatcher.RegisterHandler(&msg.Pong{}, ctl.handlePong)
2017-03-09 02:03:47 +08:00
}
// headerWorker sends heartbeat to server and check heartbeat timeout.
func (ctl *Control) heartbeatWorker() {
2019-10-12 20:13:12 +08:00
xl := ctl.xl
2020-05-24 17:48:37 +08:00
// TODO(fatedier): Change default value of HeartbeatInterval to -1 if tcpmux is enabled.
// Users can still enable heartbeat feature by setting HeartbeatInterval to a positive value.
2023-11-27 15:47:49 +08:00
if ctl.sessionCtx.Common.Transport.HeartbeatInterval > 0 {
// send heartbeat to server
sendHeartBeat := func() (bool, error) {
xl.Debug("send heartbeat to server")
pingMsg := &msg.Ping{}
2023-11-27 15:47:49 +08:00
if err := ctl.sessionCtx.AuthSetter.SetPing(pingMsg); err != nil {
xl.Warn("error during ping authentication: %v, skip sending ping message", err)
return false, err
}
_ = ctl.msgDispatcher.Send(pingMsg)
return false, nil
2017-03-09 02:03:47 +08:00
}
go wait.BackoffUntil(sendHeartBeat,
wait.NewFastBackoffManager(wait.FastBackoffOptions{
2023-11-27 15:47:49 +08:00
Duration: time.Duration(ctl.sessionCtx.Common.Transport.HeartbeatInterval) * time.Second,
InitDurationIfFail: time.Second,
Factor: 2.0,
Jitter: 0.1,
2023-11-27 15:47:49 +08:00
MaxDuration: time.Duration(ctl.sessionCtx.Common.Transport.HeartbeatInterval) * time.Second,
}),
true, ctl.doneCh,
)
}
2017-03-09 02:03:47 +08:00
// Check heartbeat timeout only if TCPMux is not enabled and users don't disable heartbeat feature.
2023-11-27 15:47:49 +08:00
if ctl.sessionCtx.Common.Transport.HeartbeatInterval > 0 && ctl.sessionCtx.Common.Transport.HeartbeatTimeout > 0 &&
!lo.FromPtr(ctl.sessionCtx.Common.Transport.TCPMux) {
2018-01-17 01:09:33 +08:00
go wait.Until(func() {
2023-11-27 15:47:49 +08:00
if time.Since(ctl.lastPong.Load().(time.Time)) > time.Duration(ctl.sessionCtx.Common.Transport.HeartbeatTimeout)*time.Second {
2019-10-12 20:13:12 +08:00
xl.Warn("heartbeat timeout")
2023-11-27 15:47:49 +08:00
ctl.closeSession()
2017-03-09 02:03:47 +08:00
return
}
}, time.Second, ctl.doneCh)
2017-03-09 02:03:47 +08:00
}
}
2018-01-17 01:09:33 +08:00
func (ctl *Control) worker() {
go ctl.heartbeatWorker()
go ctl.msgDispatcher.Run()
2018-01-17 01:09:33 +08:00
<-ctl.msgDispatcher.Done()
2023-11-27 15:47:49 +08:00
ctl.closeSession()
2018-11-06 18:35:05 +08:00
2022-08-29 01:02:53 +08:00
ctl.pm.Close()
ctl.vm.Close()
close(ctl.doneCh)
2017-03-09 02:03:47 +08:00
}
2017-06-27 23:30:09 +08:00
2023-11-27 15:47:49 +08:00
func (ctl *Control) UpdateAllConfigurer(proxyCfgs []v1.ProxyConfigurer, visitorCfgs []v1.VisitorConfigurer) error {
ctl.vm.UpdateAll(visitorCfgs)
ctl.pm.UpdateAll(proxyCfgs)
2018-12-07 17:05:36 +08:00
return nil
}