proxy.go 20 KB

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