proxy.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806
  1. // Copyright 2017 fatedier, [email protected]
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package proxy
  15. import (
  16. "bytes"
  17. "context"
  18. "fmt"
  19. "io"
  20. "net"
  21. "strconv"
  22. "strings"
  23. "sync"
  24. "time"
  25. "github.com/fatedier/frp/pkg/config"
  26. "github.com/fatedier/frp/pkg/msg"
  27. plugin "github.com/fatedier/frp/pkg/plugin/client"
  28. "github.com/fatedier/frp/pkg/proto/udp"
  29. "github.com/fatedier/frp/pkg/util/limit"
  30. frpNet "github.com/fatedier/frp/pkg/util/net"
  31. "github.com/fatedier/frp/pkg/util/xlog"
  32. "github.com/fatedier/golib/errors"
  33. frpIo "github.com/fatedier/golib/io"
  34. libdial "github.com/fatedier/golib/net/dial"
  35. "github.com/fatedier/golib/pool"
  36. fmux "github.com/hashicorp/yamux"
  37. pp "github.com/pires/go-proxyproto"
  38. "golang.org/x/time/rate"
  39. )
  40. // Proxy defines how to handle work connections for different proxy type.
  41. type Proxy interface {
  42. Run() error
  43. // InWorkConn accept work connections registered to server.
  44. InWorkConn(net.Conn, *msg.StartWorkConn)
  45. Close()
  46. }
  47. func NewProxy(ctx context.Context, pxyConf config.ProxyConf, clientCfg config.ClientCommonConf, serverUDPPort int) (pxy Proxy) {
  48. var limiter *rate.Limiter
  49. limitBytes := pxyConf.GetBaseInfo().BandwidthLimit.Bytes()
  50. if limitBytes > 0 {
  51. limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
  52. }
  53. baseProxy := BaseProxy{
  54. clientCfg: clientCfg,
  55. serverUDPPort: serverUDPPort,
  56. limiter: limiter,
  57. xl: xlog.FromContextSafe(ctx),
  58. ctx: ctx,
  59. }
  60. switch cfg := pxyConf.(type) {
  61. case *config.TCPProxyConf:
  62. pxy = &TCPProxy{
  63. BaseProxy: &baseProxy,
  64. cfg: cfg,
  65. }
  66. case *config.TCPMuxProxyConf:
  67. pxy = &TCPMuxProxy{
  68. BaseProxy: &baseProxy,
  69. cfg: cfg,
  70. }
  71. case *config.UDPProxyConf:
  72. pxy = &UDPProxy{
  73. BaseProxy: &baseProxy,
  74. cfg: cfg,
  75. }
  76. case *config.HTTPProxyConf:
  77. pxy = &HTTPProxy{
  78. BaseProxy: &baseProxy,
  79. cfg: cfg,
  80. }
  81. case *config.HTTPSProxyConf:
  82. pxy = &HTTPSProxy{
  83. BaseProxy: &baseProxy,
  84. cfg: cfg,
  85. }
  86. case *config.STCPProxyConf:
  87. pxy = &STCPProxy{
  88. BaseProxy: &baseProxy,
  89. cfg: cfg,
  90. }
  91. case *config.XTCPProxyConf:
  92. pxy = &XTCPProxy{
  93. BaseProxy: &baseProxy,
  94. cfg: cfg,
  95. }
  96. case *config.SUDPProxyConf:
  97. pxy = &SUDPProxy{
  98. BaseProxy: &baseProxy,
  99. cfg: cfg,
  100. closeCh: make(chan struct{}),
  101. }
  102. }
  103. return
  104. }
  105. type BaseProxy struct {
  106. closed bool
  107. clientCfg config.ClientCommonConf
  108. serverUDPPort int
  109. limiter *rate.Limiter
  110. mu sync.RWMutex
  111. xl *xlog.Logger
  112. ctx context.Context
  113. }
  114. // TCP
  115. type TCPProxy struct {
  116. *BaseProxy
  117. cfg *config.TCPProxyConf
  118. proxyPlugin plugin.Plugin
  119. }
  120. func (pxy *TCPProxy) Run() (err error) {
  121. if pxy.cfg.Plugin != "" {
  122. pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
  123. if err != nil {
  124. return
  125. }
  126. }
  127. return
  128. }
  129. func (pxy *TCPProxy) Close() {
  130. if pxy.proxyPlugin != nil {
  131. pxy.proxyPlugin.Close()
  132. }
  133. }
  134. func (pxy *TCPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  135. HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
  136. conn, []byte(pxy.clientCfg.Token), m)
  137. }
  138. // TCP Multiplexer
  139. type TCPMuxProxy struct {
  140. *BaseProxy
  141. cfg *config.TCPMuxProxyConf
  142. proxyPlugin plugin.Plugin
  143. }
  144. func (pxy *TCPMuxProxy) Run() (err error) {
  145. if pxy.cfg.Plugin != "" {
  146. pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
  147. if err != nil {
  148. return
  149. }
  150. }
  151. return
  152. }
  153. func (pxy *TCPMuxProxy) Close() {
  154. if pxy.proxyPlugin != nil {
  155. pxy.proxyPlugin.Close()
  156. }
  157. }
  158. func (pxy *TCPMuxProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  159. HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
  160. conn, []byte(pxy.clientCfg.Token), m)
  161. }
  162. // HTTP
  163. type HTTPProxy struct {
  164. *BaseProxy
  165. cfg *config.HTTPProxyConf
  166. proxyPlugin plugin.Plugin
  167. }
  168. func (pxy *HTTPProxy) Run() (err error) {
  169. if pxy.cfg.Plugin != "" {
  170. pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
  171. if err != nil {
  172. return
  173. }
  174. }
  175. return
  176. }
  177. func (pxy *HTTPProxy) Close() {
  178. if pxy.proxyPlugin != nil {
  179. pxy.proxyPlugin.Close()
  180. }
  181. }
  182. func (pxy *HTTPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  183. HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
  184. conn, []byte(pxy.clientCfg.Token), m)
  185. }
  186. // HTTPS
  187. type HTTPSProxy struct {
  188. *BaseProxy
  189. cfg *config.HTTPSProxyConf
  190. proxyPlugin plugin.Plugin
  191. }
  192. func (pxy *HTTPSProxy) Run() (err error) {
  193. if pxy.cfg.Plugin != "" {
  194. pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
  195. if err != nil {
  196. return
  197. }
  198. }
  199. return
  200. }
  201. func (pxy *HTTPSProxy) Close() {
  202. if pxy.proxyPlugin != nil {
  203. pxy.proxyPlugin.Close()
  204. }
  205. }
  206. func (pxy *HTTPSProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  207. HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
  208. conn, []byte(pxy.clientCfg.Token), m)
  209. }
  210. // STCP
  211. type STCPProxy struct {
  212. *BaseProxy
  213. cfg *config.STCPProxyConf
  214. proxyPlugin plugin.Plugin
  215. }
  216. func (pxy *STCPProxy) Run() (err error) {
  217. if pxy.cfg.Plugin != "" {
  218. pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
  219. if err != nil {
  220. return
  221. }
  222. }
  223. return
  224. }
  225. func (pxy *STCPProxy) Close() {
  226. if pxy.proxyPlugin != nil {
  227. pxy.proxyPlugin.Close()
  228. }
  229. }
  230. func (pxy *STCPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  231. HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
  232. conn, []byte(pxy.clientCfg.Token), m)
  233. }
  234. // XTCP
  235. type XTCPProxy struct {
  236. *BaseProxy
  237. cfg *config.XTCPProxyConf
  238. proxyPlugin plugin.Plugin
  239. }
  240. func (pxy *XTCPProxy) Run() (err error) {
  241. if pxy.cfg.Plugin != "" {
  242. pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
  243. if err != nil {
  244. return
  245. }
  246. }
  247. return
  248. }
  249. func (pxy *XTCPProxy) Close() {
  250. if pxy.proxyPlugin != nil {
  251. pxy.proxyPlugin.Close()
  252. }
  253. }
  254. func (pxy *XTCPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  255. xl := pxy.xl
  256. defer conn.Close()
  257. var natHoleSidMsg msg.NatHoleSid
  258. err := msg.ReadMsgInto(conn, &natHoleSidMsg)
  259. if err != nil {
  260. xl.Error("xtcp read from workConn error: %v", err)
  261. return
  262. }
  263. natHoleClientMsg := &msg.NatHoleClient{
  264. ProxyName: pxy.cfg.ProxyName,
  265. Sid: natHoleSidMsg.Sid,
  266. }
  267. raddr, _ := net.ResolveUDPAddr("udp",
  268. fmt.Sprintf("%s:%d", pxy.clientCfg.ServerAddr, pxy.serverUDPPort))
  269. clientConn, err := net.DialUDP("udp", nil, raddr)
  270. if err != nil {
  271. xl.Error("dial server udp addr error: %v", err)
  272. return
  273. }
  274. defer clientConn.Close()
  275. err = msg.WriteMsg(clientConn, natHoleClientMsg)
  276. if err != nil {
  277. xl.Error("send natHoleClientMsg to server error: %v", err)
  278. return
  279. }
  280. // Wait for client address at most 5 seconds.
  281. var natHoleRespMsg msg.NatHoleResp
  282. clientConn.SetReadDeadline(time.Now().Add(5 * time.Second))
  283. buf := pool.GetBuf(1024)
  284. n, err := clientConn.Read(buf)
  285. if err != nil {
  286. xl.Error("get natHoleRespMsg error: %v", err)
  287. return
  288. }
  289. err = msg.ReadMsgInto(bytes.NewReader(buf[:n]), &natHoleRespMsg)
  290. if err != nil {
  291. xl.Error("get natHoleRespMsg error: %v", err)
  292. return
  293. }
  294. clientConn.SetReadDeadline(time.Time{})
  295. clientConn.Close()
  296. if natHoleRespMsg.Error != "" {
  297. xl.Error("natHoleRespMsg get error info: %s", natHoleRespMsg.Error)
  298. return
  299. }
  300. xl.Trace("get natHoleRespMsg, sid [%s], client address [%s] visitor address [%s]", natHoleRespMsg.Sid, natHoleRespMsg.ClientAddr, natHoleRespMsg.VisitorAddr)
  301. // Send detect message
  302. host, portStr, err := net.SplitHostPort(natHoleRespMsg.VisitorAddr)
  303. if err != nil {
  304. xl.Error("get NatHoleResp visitor address [%s] error: %v", natHoleRespMsg.VisitorAddr, err)
  305. }
  306. laddr, _ := net.ResolveUDPAddr("udp", clientConn.LocalAddr().String())
  307. port, err := strconv.ParseInt(portStr, 10, 64)
  308. if err != nil {
  309. xl.Error("get natHoleResp visitor address error: %v", natHoleRespMsg.VisitorAddr)
  310. return
  311. }
  312. pxy.sendDetectMsg(host, int(port), laddr, []byte(natHoleRespMsg.Sid))
  313. xl.Trace("send all detect msg done")
  314. msg.WriteMsg(conn, &msg.NatHoleClientDetectOK{})
  315. // Listen for clientConn's address and wait for visitor connection
  316. lConn, err := net.ListenUDP("udp", laddr)
  317. if err != nil {
  318. xl.Error("listen on visitorConn's local adress error: %v", err)
  319. return
  320. }
  321. defer lConn.Close()
  322. lConn.SetReadDeadline(time.Now().Add(8 * time.Second))
  323. sidBuf := pool.GetBuf(1024)
  324. var uAddr *net.UDPAddr
  325. n, uAddr, err = lConn.ReadFromUDP(sidBuf)
  326. if err != nil {
  327. xl.Warn("get sid from visitor error: %v", err)
  328. return
  329. }
  330. lConn.SetReadDeadline(time.Time{})
  331. if string(sidBuf[:n]) != natHoleRespMsg.Sid {
  332. xl.Warn("incorrect sid from visitor")
  333. return
  334. }
  335. pool.PutBuf(sidBuf)
  336. xl.Info("nat hole connection make success, sid [%s]", natHoleRespMsg.Sid)
  337. lConn.WriteToUDP(sidBuf[:n], uAddr)
  338. kcpConn, err := frpNet.NewKCPConnFromUDP(lConn, false, uAddr.String())
  339. if err != nil {
  340. xl.Error("create kcp connection from udp connection error: %v", err)
  341. return
  342. }
  343. fmuxCfg := fmux.DefaultConfig()
  344. fmuxCfg.KeepAliveInterval = 5 * time.Second
  345. fmuxCfg.LogOutput = io.Discard
  346. sess, err := fmux.Server(kcpConn, fmuxCfg)
  347. if err != nil {
  348. xl.Error("create yamux server from kcp connection error: %v", err)
  349. return
  350. }
  351. defer sess.Close()
  352. muxConn, err := sess.Accept()
  353. if err != nil {
  354. xl.Error("accept for yamux connection error: %v", err)
  355. return
  356. }
  357. HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
  358. muxConn, []byte(pxy.cfg.Sk), m)
  359. }
  360. func (pxy *XTCPProxy) sendDetectMsg(addr string, port int, laddr *net.UDPAddr, content []byte) (err error) {
  361. daddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", addr, port))
  362. if err != nil {
  363. return err
  364. }
  365. tConn, err := net.DialUDP("udp", laddr, daddr)
  366. if err != nil {
  367. return err
  368. }
  369. //uConn := ipv4.NewConn(tConn)
  370. //uConn.SetTTL(3)
  371. tConn.Write(content)
  372. tConn.Close()
  373. return nil
  374. }
  375. // UDP
  376. type UDPProxy struct {
  377. *BaseProxy
  378. cfg *config.UDPProxyConf
  379. localAddr *net.UDPAddr
  380. readCh chan *msg.UDPPacket
  381. // include msg.UDPPacket and msg.Ping
  382. sendCh chan msg.Message
  383. workConn net.Conn
  384. }
  385. func (pxy *UDPProxy) Run() (err error) {
  386. pxy.localAddr, err = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pxy.cfg.LocalIP, pxy.cfg.LocalPort))
  387. if err != nil {
  388. return
  389. }
  390. return
  391. }
  392. func (pxy *UDPProxy) Close() {
  393. pxy.mu.Lock()
  394. defer pxy.mu.Unlock()
  395. if !pxy.closed {
  396. pxy.closed = true
  397. if pxy.workConn != nil {
  398. pxy.workConn.Close()
  399. }
  400. if pxy.readCh != nil {
  401. close(pxy.readCh)
  402. }
  403. if pxy.sendCh != nil {
  404. close(pxy.sendCh)
  405. }
  406. }
  407. }
  408. func (pxy *UDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  409. xl := pxy.xl
  410. xl.Info("incoming a new work connection for udp proxy, %s", conn.RemoteAddr().String())
  411. // close resources releated with old workConn
  412. pxy.Close()
  413. var rwc io.ReadWriteCloser = conn
  414. var err error
  415. if pxy.limiter != nil {
  416. rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
  417. return conn.Close()
  418. })
  419. }
  420. if pxy.cfg.UseEncryption {
  421. rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
  422. if err != nil {
  423. conn.Close()
  424. xl.Error("create encryption stream error: %v", err)
  425. return
  426. }
  427. }
  428. if pxy.cfg.UseCompression {
  429. rwc = frpIo.WithCompression(rwc)
  430. }
  431. conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
  432. pxy.mu.Lock()
  433. pxy.workConn = conn
  434. pxy.readCh = make(chan *msg.UDPPacket, 1024)
  435. pxy.sendCh = make(chan msg.Message, 1024)
  436. pxy.closed = false
  437. pxy.mu.Unlock()
  438. workConnReaderFn := func(conn net.Conn, readCh chan *msg.UDPPacket) {
  439. for {
  440. var udpMsg msg.UDPPacket
  441. if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
  442. xl.Warn("read from workConn for udp error: %v", errRet)
  443. return
  444. }
  445. if errRet := errors.PanicToError(func() {
  446. xl.Trace("get udp package from workConn: %s", udpMsg.Content)
  447. readCh <- &udpMsg
  448. }); errRet != nil {
  449. xl.Info("reader goroutine for udp work connection closed: %v", errRet)
  450. return
  451. }
  452. }
  453. }
  454. workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
  455. defer func() {
  456. xl.Info("writer goroutine for udp work connection closed")
  457. }()
  458. var errRet error
  459. for rawMsg := range sendCh {
  460. switch m := rawMsg.(type) {
  461. case *msg.UDPPacket:
  462. xl.Trace("send udp package to workConn: %s", m.Content)
  463. case *msg.Ping:
  464. xl.Trace("send ping message to udp workConn")
  465. }
  466. if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
  467. xl.Error("udp work write error: %v", errRet)
  468. return
  469. }
  470. }
  471. }
  472. heartbeatFn := func(conn net.Conn, sendCh chan msg.Message) {
  473. var errRet error
  474. for {
  475. time.Sleep(time.Duration(30) * time.Second)
  476. if errRet = errors.PanicToError(func() {
  477. sendCh <- &msg.Ping{}
  478. }); errRet != nil {
  479. xl.Trace("heartbeat goroutine for udp work connection closed")
  480. break
  481. }
  482. }
  483. }
  484. go workConnSenderFn(pxy.workConn, pxy.sendCh)
  485. go workConnReaderFn(pxy.workConn, pxy.readCh)
  486. go heartbeatFn(pxy.workConn, pxy.sendCh)
  487. udp.Forwarder(pxy.localAddr, pxy.readCh, pxy.sendCh, int(pxy.clientCfg.UDPPacketSize))
  488. }
  489. type SUDPProxy struct {
  490. *BaseProxy
  491. cfg *config.SUDPProxyConf
  492. localAddr *net.UDPAddr
  493. closeCh chan struct{}
  494. }
  495. func (pxy *SUDPProxy) Run() (err error) {
  496. pxy.localAddr, err = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pxy.cfg.LocalIP, pxy.cfg.LocalPort))
  497. if err != nil {
  498. return
  499. }
  500. return
  501. }
  502. func (pxy *SUDPProxy) Close() {
  503. pxy.mu.Lock()
  504. defer pxy.mu.Unlock()
  505. select {
  506. case <-pxy.closeCh:
  507. return
  508. default:
  509. close(pxy.closeCh)
  510. }
  511. }
  512. func (pxy *SUDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  513. xl := pxy.xl
  514. xl.Info("incoming a new work connection for sudp proxy, %s", conn.RemoteAddr().String())
  515. var rwc io.ReadWriteCloser = conn
  516. var err error
  517. if pxy.limiter != nil {
  518. rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
  519. return conn.Close()
  520. })
  521. }
  522. if pxy.cfg.UseEncryption {
  523. rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
  524. if err != nil {
  525. conn.Close()
  526. xl.Error("create encryption stream error: %v", err)
  527. return
  528. }
  529. }
  530. if pxy.cfg.UseCompression {
  531. rwc = frpIo.WithCompression(rwc)
  532. }
  533. conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
  534. workConn := conn
  535. readCh := make(chan *msg.UDPPacket, 1024)
  536. sendCh := make(chan msg.Message, 1024)
  537. isClose := false
  538. mu := &sync.Mutex{}
  539. closeFn := func() {
  540. mu.Lock()
  541. defer mu.Unlock()
  542. if isClose {
  543. return
  544. }
  545. isClose = true
  546. if workConn != nil {
  547. workConn.Close()
  548. }
  549. close(readCh)
  550. close(sendCh)
  551. }
  552. // udp service <- frpc <- frps <- frpc visitor <- user
  553. workConnReaderFn := func(conn net.Conn, readCh chan *msg.UDPPacket) {
  554. defer closeFn()
  555. for {
  556. // first to check sudp proxy is closed or not
  557. select {
  558. case <-pxy.closeCh:
  559. xl.Trace("frpc sudp proxy is closed")
  560. return
  561. default:
  562. }
  563. var udpMsg msg.UDPPacket
  564. if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
  565. xl.Warn("read from workConn for sudp error: %v", errRet)
  566. return
  567. }
  568. if errRet := errors.PanicToError(func() {
  569. readCh <- &udpMsg
  570. }); errRet != nil {
  571. xl.Warn("reader goroutine for sudp work connection closed: %v", errRet)
  572. return
  573. }
  574. }
  575. }
  576. // udp service -> frpc -> frps -> frpc visitor -> user
  577. workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
  578. defer func() {
  579. closeFn()
  580. xl.Info("writer goroutine for sudp work connection closed")
  581. }()
  582. var errRet error
  583. for rawMsg := range sendCh {
  584. switch m := rawMsg.(type) {
  585. case *msg.UDPPacket:
  586. xl.Trace("frpc send udp package to frpc visitor, [udp local: %v, remote: %v], [tcp work conn local: %v, remote: %v]",
  587. m.LocalAddr.String(), m.RemoteAddr.String(), conn.LocalAddr().String(), conn.RemoteAddr().String())
  588. case *msg.Ping:
  589. xl.Trace("frpc send ping message to frpc visitor")
  590. }
  591. if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
  592. xl.Error("sudp work write error: %v", errRet)
  593. return
  594. }
  595. }
  596. }
  597. heartbeatFn := func(conn net.Conn, sendCh chan msg.Message) {
  598. ticker := time.NewTicker(30 * time.Second)
  599. defer func() {
  600. ticker.Stop()
  601. closeFn()
  602. }()
  603. var errRet error
  604. for {
  605. select {
  606. case <-ticker.C:
  607. if errRet = errors.PanicToError(func() {
  608. sendCh <- &msg.Ping{}
  609. }); errRet != nil {
  610. xl.Warn("heartbeat goroutine for sudp work connection closed")
  611. return
  612. }
  613. case <-pxy.closeCh:
  614. xl.Trace("frpc sudp proxy is closed")
  615. return
  616. }
  617. }
  618. }
  619. go workConnSenderFn(workConn, sendCh)
  620. go workConnReaderFn(workConn, readCh)
  621. go heartbeatFn(workConn, sendCh)
  622. udp.Forwarder(pxy.localAddr, readCh, sendCh, int(pxy.clientCfg.UDPPacketSize))
  623. }
  624. // Common handler for tcp work connections.
  625. func HandleTCPWorkConnection(ctx context.Context, localInfo *config.LocalSvrConf, proxyPlugin plugin.Plugin,
  626. baseInfo *config.BaseProxyConf, limiter *rate.Limiter, workConn net.Conn, encKey []byte, m *msg.StartWorkConn) {
  627. xl := xlog.FromContextSafe(ctx)
  628. var (
  629. remote io.ReadWriteCloser
  630. err error
  631. )
  632. remote = workConn
  633. if limiter != nil {
  634. remote = frpIo.WrapReadWriteCloser(limit.NewReader(workConn, limiter), limit.NewWriter(workConn, limiter), func() error {
  635. return workConn.Close()
  636. })
  637. }
  638. xl.Trace("handle tcp work connection, use_encryption: %t, use_compression: %t",
  639. baseInfo.UseEncryption, baseInfo.UseCompression)
  640. if baseInfo.UseEncryption {
  641. remote, err = frpIo.WithEncryption(remote, encKey)
  642. if err != nil {
  643. workConn.Close()
  644. xl.Error("create encryption stream error: %v", err)
  645. return
  646. }
  647. }
  648. if baseInfo.UseCompression {
  649. remote = frpIo.WithCompression(remote)
  650. }
  651. // check if we need to send proxy protocol info
  652. var extraInfo []byte
  653. if baseInfo.ProxyProtocolVersion != "" {
  654. if m.SrcAddr != "" && m.SrcPort != 0 {
  655. if m.DstAddr == "" {
  656. m.DstAddr = "127.0.0.1"
  657. }
  658. srcAddr, _ := net.ResolveTCPAddr("tcp", net.JoinHostPort(m.SrcAddr, strconv.Itoa(int(m.SrcPort))))
  659. dstAddr, _ := net.ResolveTCPAddr("tcp", net.JoinHostPort(m.DstAddr, strconv.Itoa(int(m.DstPort))))
  660. h := &pp.Header{
  661. Command: pp.PROXY,
  662. SourceAddr: srcAddr,
  663. DestinationAddr: dstAddr,
  664. }
  665. if strings.Contains(m.SrcAddr, ".") {
  666. h.TransportProtocol = pp.TCPv4
  667. } else {
  668. h.TransportProtocol = pp.TCPv6
  669. }
  670. if baseInfo.ProxyProtocolVersion == "v1" {
  671. h.Version = 1
  672. } else if baseInfo.ProxyProtocolVersion == "v2" {
  673. h.Version = 2
  674. }
  675. buf := bytes.NewBuffer(nil)
  676. h.WriteTo(buf)
  677. extraInfo = buf.Bytes()
  678. }
  679. }
  680. if proxyPlugin != nil {
  681. // if plugin is set, let plugin handle connections first
  682. xl.Debug("handle by plugin: %s", proxyPlugin.Name())
  683. proxyPlugin.Handle(remote, workConn, extraInfo)
  684. xl.Debug("handle by plugin finished")
  685. return
  686. }
  687. localConn, err := libdial.Dial(net.JoinHostPort(localInfo.LocalIP, strconv.Itoa(localInfo.LocalPort)))
  688. if err != nil {
  689. workConn.Close()
  690. xl.Error("connect to local service [%s:%d] error: %v", localInfo.LocalIP, localInfo.LocalPort, err)
  691. return
  692. }
  693. xl.Debug("join connections, localConn(l[%s] r[%s]) workConn(l[%s] r[%s])", localConn.LocalAddr().String(),
  694. localConn.RemoteAddr().String(), workConn.LocalAddr().String(), workConn.RemoteAddr().String())
  695. if len(extraInfo) > 0 {
  696. localConn.Write(extraInfo)
  697. }
  698. frpIo.Join(localConn, remote)
  699. xl.Debug("join connections closed")
  700. }