proxy.go 20 KB

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