proxy.go 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806
  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. "io/ioutil"
  21. "net"
  22. "strconv"
  23. "strings"
  24. "sync"
  25. "time"
  26. "github.com/fatedier/frp/pkg/config"
  27. "github.com/fatedier/frp/pkg/msg"
  28. plugin "github.com/fatedier/frp/pkg/plugin/client"
  29. "github.com/fatedier/frp/pkg/proto/udp"
  30. "github.com/fatedier/frp/pkg/util/limit"
  31. frpNet "github.com/fatedier/frp/pkg/util/net"
  32. "github.com/fatedier/frp/pkg/util/xlog"
  33. "github.com/fatedier/golib/errors"
  34. frpIo "github.com/fatedier/golib/io"
  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.BaseProxyConf, 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.BaseProxyConf, 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.BaseProxyConf, 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.BaseProxyConf, 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.BaseProxyConf, 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. defer clientConn.Close()
  271. err = msg.WriteMsg(clientConn, natHoleClientMsg)
  272. if err != nil {
  273. xl.Error("send natHoleClientMsg to server error: %v", err)
  274. return
  275. }
  276. // Wait for client address at most 5 seconds.
  277. var natHoleRespMsg msg.NatHoleResp
  278. clientConn.SetReadDeadline(time.Now().Add(5 * time.Second))
  279. buf := pool.GetBuf(1024)
  280. n, err := clientConn.Read(buf)
  281. if err != nil {
  282. xl.Error("get natHoleRespMsg error: %v", err)
  283. return
  284. }
  285. err = msg.ReadMsgInto(bytes.NewReader(buf[:n]), &natHoleRespMsg)
  286. if err != nil {
  287. xl.Error("get natHoleRespMsg error: %v", err)
  288. return
  289. }
  290. clientConn.SetReadDeadline(time.Time{})
  291. clientConn.Close()
  292. if natHoleRespMsg.Error != "" {
  293. xl.Error("natHoleRespMsg get error info: %s", natHoleRespMsg.Error)
  294. return
  295. }
  296. xl.Trace("get natHoleRespMsg, sid [%s], client address [%s] visitor address [%s]", natHoleRespMsg.Sid, natHoleRespMsg.ClientAddr, natHoleRespMsg.VisitorAddr)
  297. // Send detect message
  298. array := strings.Split(natHoleRespMsg.VisitorAddr, ":")
  299. if len(array) <= 1 {
  300. xl.Error("get NatHoleResp visitor address error: %v", natHoleRespMsg.VisitorAddr)
  301. }
  302. laddr, _ := net.ResolveUDPAddr("udp", clientConn.LocalAddr().String())
  303. /*
  304. for i := 1000; i < 65000; i++ {
  305. pxy.sendDetectMsg(array[0], int64(i), laddr, "a")
  306. }
  307. */
  308. port, err := strconv.ParseInt(array[1], 10, 64)
  309. if err != nil {
  310. xl.Error("get natHoleResp visitor address error: %v", natHoleRespMsg.VisitorAddr)
  311. return
  312. }
  313. pxy.sendDetectMsg(array[0], int(port), laddr, []byte(natHoleRespMsg.Sid))
  314. xl.Trace("send all detect msg done")
  315. msg.WriteMsg(conn, &msg.NatHoleClientDetectOK{})
  316. // Listen for clientConn's address and wait for visitor connection
  317. lConn, err := net.ListenUDP("udp", laddr)
  318. if err != nil {
  319. xl.Error("listen on visitorConn's local adress error: %v", err)
  320. return
  321. }
  322. defer lConn.Close()
  323. lConn.SetReadDeadline(time.Now().Add(8 * time.Second))
  324. sidBuf := pool.GetBuf(1024)
  325. var uAddr *net.UDPAddr
  326. n, uAddr, err = lConn.ReadFromUDP(sidBuf)
  327. if err != nil {
  328. xl.Warn("get sid from visitor error: %v", err)
  329. return
  330. }
  331. lConn.SetReadDeadline(time.Time{})
  332. if string(sidBuf[:n]) != natHoleRespMsg.Sid {
  333. xl.Warn("incorrect sid from visitor")
  334. return
  335. }
  336. pool.PutBuf(sidBuf)
  337. xl.Info("nat hole connection make success, sid [%s]", natHoleRespMsg.Sid)
  338. lConn.WriteToUDP(sidBuf[:n], uAddr)
  339. kcpConn, err := frpNet.NewKCPConnFromUDP(lConn, false, uAddr.String())
  340. if err != nil {
  341. xl.Error("create kcp connection from udp connection error: %v", err)
  342. return
  343. }
  344. fmuxCfg := fmux.DefaultConfig()
  345. fmuxCfg.KeepAliveInterval = 5 * time.Second
  346. fmuxCfg.LogOutput = ioutil.Discard
  347. sess, err := fmux.Server(kcpConn, fmuxCfg)
  348. if err != nil {
  349. xl.Error("create yamux server from kcp connection error: %v", err)
  350. return
  351. }
  352. defer sess.Close()
  353. muxConn, err := sess.Accept()
  354. if err != nil {
  355. xl.Error("accept for yamux connection error: %v", err)
  356. return
  357. }
  358. HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, pxy.limiter,
  359. muxConn, []byte(pxy.cfg.Sk), m)
  360. }
  361. func (pxy *XTCPProxy) sendDetectMsg(addr string, port int, laddr *net.UDPAddr, content []byte) (err error) {
  362. daddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", addr, port))
  363. if err != nil {
  364. return err
  365. }
  366. tConn, err := net.DialUDP("udp", laddr, daddr)
  367. if err != nil {
  368. return err
  369. }
  370. //uConn := ipv4.NewConn(tConn)
  371. //uConn.SetTTL(3)
  372. tConn.Write(content)
  373. tConn.Close()
  374. return nil
  375. }
  376. // UDP
  377. type UDPProxy struct {
  378. *BaseProxy
  379. cfg *config.UDPProxyConf
  380. localAddr *net.UDPAddr
  381. readCh chan *msg.UDPPacket
  382. // include msg.UDPPacket and msg.Ping
  383. sendCh chan msg.Message
  384. workConn net.Conn
  385. }
  386. func (pxy *UDPProxy) Run() (err error) {
  387. pxy.localAddr, err = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pxy.cfg.LocalIP, pxy.cfg.LocalPort))
  388. if err != nil {
  389. return
  390. }
  391. return
  392. }
  393. func (pxy *UDPProxy) Close() {
  394. pxy.mu.Lock()
  395. defer pxy.mu.Unlock()
  396. if !pxy.closed {
  397. pxy.closed = true
  398. if pxy.workConn != nil {
  399. pxy.workConn.Close()
  400. }
  401. if pxy.readCh != nil {
  402. close(pxy.readCh)
  403. }
  404. if pxy.sendCh != nil {
  405. close(pxy.sendCh)
  406. }
  407. }
  408. }
  409. func (pxy *UDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  410. xl := pxy.xl
  411. xl.Info("incoming a new work connection for udp proxy, %s", conn.RemoteAddr().String())
  412. // close resources releated with old workConn
  413. pxy.Close()
  414. var rwc io.ReadWriteCloser = conn
  415. var err error
  416. if pxy.limiter != nil {
  417. rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
  418. return conn.Close()
  419. })
  420. }
  421. if pxy.cfg.UseEncryption {
  422. rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
  423. if err != nil {
  424. conn.Close()
  425. xl.Error("create encryption stream error: %v", err)
  426. return
  427. }
  428. }
  429. if pxy.cfg.UseCompression {
  430. rwc = frpIo.WithCompression(rwc)
  431. }
  432. conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
  433. pxy.mu.Lock()
  434. pxy.workConn = conn
  435. pxy.readCh = make(chan *msg.UDPPacket, 1024)
  436. pxy.sendCh = make(chan msg.Message, 1024)
  437. pxy.closed = false
  438. pxy.mu.Unlock()
  439. workConnReaderFn := func(conn net.Conn, readCh chan *msg.UDPPacket) {
  440. for {
  441. var udpMsg msg.UDPPacket
  442. if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
  443. xl.Warn("read from workConn for udp error: %v", errRet)
  444. return
  445. }
  446. if errRet := errors.PanicToError(func() {
  447. xl.Trace("get udp package from workConn: %s", udpMsg.Content)
  448. readCh <- &udpMsg
  449. }); errRet != nil {
  450. xl.Info("reader goroutine for udp work connection closed: %v", errRet)
  451. return
  452. }
  453. }
  454. }
  455. workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
  456. defer func() {
  457. xl.Info("writer goroutine for udp work connection closed")
  458. }()
  459. var errRet error
  460. for rawMsg := range sendCh {
  461. switch m := rawMsg.(type) {
  462. case *msg.UDPPacket:
  463. xl.Trace("send udp package to workConn: %s", m.Content)
  464. case *msg.Ping:
  465. xl.Trace("send ping message to udp workConn")
  466. }
  467. if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
  468. xl.Error("udp work write error: %v", errRet)
  469. return
  470. }
  471. }
  472. }
  473. heartbeatFn := func(conn net.Conn, sendCh chan msg.Message) {
  474. var errRet error
  475. for {
  476. time.Sleep(time.Duration(30) * time.Second)
  477. if errRet = errors.PanicToError(func() {
  478. sendCh <- &msg.Ping{}
  479. }); errRet != nil {
  480. xl.Trace("heartbeat goroutine for udp work connection closed")
  481. break
  482. }
  483. }
  484. }
  485. go workConnSenderFn(pxy.workConn, pxy.sendCh)
  486. go workConnReaderFn(pxy.workConn, pxy.readCh)
  487. go heartbeatFn(pxy.workConn, pxy.sendCh)
  488. udp.Forwarder(pxy.localAddr, pxy.readCh, pxy.sendCh, int(pxy.clientCfg.UDPPacketSize))
  489. }
  490. type SUDPProxy struct {
  491. *BaseProxy
  492. cfg *config.SUDPProxyConf
  493. localAddr *net.UDPAddr
  494. closeCh chan struct{}
  495. }
  496. func (pxy *SUDPProxy) Run() (err error) {
  497. pxy.localAddr, err = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pxy.cfg.LocalIP, pxy.cfg.LocalPort))
  498. if err != nil {
  499. return
  500. }
  501. return
  502. }
  503. func (pxy *SUDPProxy) Close() {
  504. pxy.mu.Lock()
  505. defer pxy.mu.Unlock()
  506. select {
  507. case <-pxy.closeCh:
  508. return
  509. default:
  510. close(pxy.closeCh)
  511. }
  512. }
  513. func (pxy *SUDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
  514. xl := pxy.xl
  515. xl.Info("incoming a new work connection for sudp proxy, %s", conn.RemoteAddr().String())
  516. var rwc io.ReadWriteCloser = conn
  517. var err error
  518. if pxy.limiter != nil {
  519. rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
  520. return conn.Close()
  521. })
  522. }
  523. if pxy.cfg.UseEncryption {
  524. rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
  525. if err != nil {
  526. conn.Close()
  527. xl.Error("create encryption stream error: %v", err)
  528. return
  529. }
  530. }
  531. if pxy.cfg.UseCompression {
  532. rwc = frpIo.WithCompression(rwc)
  533. }
  534. conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
  535. workConn := conn
  536. readCh := make(chan *msg.UDPPacket, 1024)
  537. sendCh := make(chan msg.Message, 1024)
  538. isClose := false
  539. mu := &sync.Mutex{}
  540. closeFn := func() {
  541. mu.Lock()
  542. defer mu.Unlock()
  543. if isClose {
  544. return
  545. }
  546. isClose = true
  547. if workConn != nil {
  548. workConn.Close()
  549. }
  550. close(readCh)
  551. close(sendCh)
  552. }
  553. // udp service <- frpc <- frps <- frpc visitor <- user
  554. workConnReaderFn := func(conn net.Conn, readCh chan *msg.UDPPacket) {
  555. defer closeFn()
  556. for {
  557. // first to check sudp proxy is closed or not
  558. select {
  559. case <-pxy.closeCh:
  560. xl.Trace("frpc sudp proxy is closed")
  561. return
  562. default:
  563. }
  564. var udpMsg msg.UDPPacket
  565. if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
  566. xl.Warn("read from workConn for sudp error: %v", errRet)
  567. return
  568. }
  569. if errRet := errors.PanicToError(func() {
  570. readCh <- &udpMsg
  571. }); errRet != nil {
  572. xl.Warn("reader goroutine for sudp work connection closed: %v", errRet)
  573. return
  574. }
  575. }
  576. }
  577. // udp service -> frpc -> frps -> frpc visitor -> user
  578. workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
  579. defer func() {
  580. closeFn()
  581. xl.Info("writer goroutine for sudp work connection closed")
  582. }()
  583. var errRet error
  584. for rawMsg := range sendCh {
  585. switch m := rawMsg.(type) {
  586. case *msg.UDPPacket:
  587. xl.Trace("frpc send udp package to frpc visitor, [udp local: %v, remote: %v], [tcp work conn local: %v, remote: %v]",
  588. m.LocalAddr.String(), m.RemoteAddr.String(), conn.LocalAddr().String(), conn.RemoteAddr().String())
  589. case *msg.Ping:
  590. xl.Trace("frpc send ping message to frpc visitor")
  591. }
  592. if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
  593. xl.Error("sudp work write error: %v", errRet)
  594. return
  595. }
  596. }
  597. }
  598. heartbeatFn := func(conn net.Conn, sendCh chan msg.Message) {
  599. ticker := time.NewTicker(30 * time.Second)
  600. defer func() {
  601. ticker.Stop()
  602. closeFn()
  603. }()
  604. var errRet error
  605. for {
  606. select {
  607. case <-ticker.C:
  608. if errRet = errors.PanicToError(func() {
  609. sendCh <- &msg.Ping{}
  610. }); errRet != nil {
  611. xl.Warn("heartbeat goroutine for sudp work connection closed")
  612. return
  613. }
  614. case <-pxy.closeCh:
  615. xl.Trace("frpc sudp proxy is closed")
  616. return
  617. }
  618. }
  619. }
  620. go workConnSenderFn(workConn, sendCh)
  621. go workConnReaderFn(workConn, readCh)
  622. go heartbeatFn(workConn, sendCh)
  623. udp.Forwarder(pxy.localAddr, readCh, sendCh, int(pxy.clientCfg.UDPPacketSize))
  624. }
  625. // Common handler for tcp work connections.
  626. func HandleTCPWorkConnection(ctx context.Context, localInfo *config.LocalSvrConf, proxyPlugin plugin.Plugin,
  627. baseInfo *config.BaseProxyConf, limiter *rate.Limiter, workConn net.Conn, encKey []byte, m *msg.StartWorkConn) {
  628. xl := xlog.FromContextSafe(ctx)
  629. var (
  630. remote io.ReadWriteCloser
  631. err error
  632. )
  633. remote = workConn
  634. if limiter != nil {
  635. remote = frpIo.WrapReadWriteCloser(limit.NewReader(workConn, limiter), limit.NewWriter(workConn, limiter), func() error {
  636. return workConn.Close()
  637. })
  638. }
  639. xl.Trace("handle tcp work connection, use_encryption: %t, use_compression: %t",
  640. baseInfo.UseEncryption, baseInfo.UseCompression)
  641. if baseInfo.UseEncryption {
  642. remote, err = frpIo.WithEncryption(remote, encKey)
  643. if err != nil {
  644. workConn.Close()
  645. xl.Error("create encryption stream error: %v", err)
  646. return
  647. }
  648. }
  649. if baseInfo.UseCompression {
  650. remote = frpIo.WithCompression(remote)
  651. }
  652. // check if we need to send proxy protocol info
  653. var extraInfo []byte
  654. if baseInfo.ProxyProtocolVersion != "" {
  655. if m.SrcAddr != "" && m.SrcPort != 0 {
  656. if m.DstAddr == "" {
  657. m.DstAddr = "127.0.0.1"
  658. }
  659. h := &pp.Header{
  660. Command: pp.PROXY,
  661. SourceAddress: net.ParseIP(m.SrcAddr),
  662. SourcePort: m.SrcPort,
  663. DestinationAddress: net.ParseIP(m.DstAddr),
  664. DestinationPort: m.DstPort,
  665. }
  666. if strings.Contains(m.SrcAddr, ".") {
  667. h.TransportProtocol = pp.TCPv4
  668. } else {
  669. h.TransportProtocol = pp.TCPv6
  670. }
  671. if baseInfo.ProxyProtocolVersion == "v1" {
  672. h.Version = 1
  673. } else if baseInfo.ProxyProtocolVersion == "v2" {
  674. h.Version = 2
  675. }
  676. buf := bytes.NewBuffer(nil)
  677. h.WriteTo(buf)
  678. extraInfo = buf.Bytes()
  679. }
  680. }
  681. if proxyPlugin != nil {
  682. // if plugin is set, let plugin handle connections first
  683. xl.Debug("handle by plugin: %s", proxyPlugin.Name())
  684. proxyPlugin.Handle(remote, workConn, extraInfo)
  685. xl.Debug("handle by plugin finished")
  686. return
  687. }
  688. localConn, err := frpNet.ConnectServer("tcp", fmt.Sprintf("%s:%d", localInfo.LocalIP, localInfo.LocalPort))
  689. if err != nil {
  690. workConn.Close()
  691. xl.Error("connect to local service [%s:%d] error: %v", localInfo.LocalIP, localInfo.LocalPort, err)
  692. return
  693. }
  694. xl.Debug("join connections, localConn(l[%s] r[%s]) workConn(l[%s] r[%s])", localConn.LocalAddr().String(),
  695. localConn.RemoteAddr().String(), workConn.LocalAddr().String(), workConn.RemoteAddr().String())
  696. if len(extraInfo) > 0 {
  697. localConn.Write(extraInfo)
  698. }
  699. frpIo.Join(localConn, remote)
  700. xl.Debug("join connections closed")
  701. }