udp.go 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226
  1. // Copyright 2019 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. "context"
  17. "fmt"
  18. "net"
  19. "time"
  20. "github.com/fatedier/frp/models/config"
  21. "github.com/fatedier/frp/models/msg"
  22. "github.com/fatedier/frp/models/proto/udp"
  23. "github.com/fatedier/frp/server/metrics"
  24. "github.com/fatedier/golib/errors"
  25. )
  26. type UDPProxy struct {
  27. *BaseProxy
  28. cfg *config.UDPProxyConf
  29. realPort int
  30. // udpConn is the listener of udp packages
  31. udpConn *net.UDPConn
  32. // there are always only one workConn at the same time
  33. // get another one if it closed
  34. workConn net.Conn
  35. // sendCh is used for sending packages to workConn
  36. sendCh chan *msg.UDPPacket
  37. // readCh is used for reading packages from workConn
  38. readCh chan *msg.UDPPacket
  39. // checkCloseCh is used for watching if workConn is closed
  40. checkCloseCh chan int
  41. isClosed bool
  42. }
  43. func (pxy *UDPProxy) Run() (remoteAddr string, err error) {
  44. xl := pxy.xl
  45. pxy.realPort, err = pxy.rc.UDPPortManager.Acquire(pxy.name, pxy.cfg.RemotePort)
  46. if err != nil {
  47. return
  48. }
  49. defer func() {
  50. if err != nil {
  51. pxy.rc.UDPPortManager.Release(pxy.realPort)
  52. }
  53. }()
  54. remoteAddr = fmt.Sprintf(":%d", pxy.realPort)
  55. pxy.cfg.RemotePort = pxy.realPort
  56. addr, errRet := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pxy.serverCfg.ProxyBindAddr, pxy.realPort))
  57. if errRet != nil {
  58. err = errRet
  59. return
  60. }
  61. udpConn, errRet := net.ListenUDP("udp", addr)
  62. if errRet != nil {
  63. err = errRet
  64. xl.Warn("listen udp port error: %v", err)
  65. return
  66. }
  67. xl.Info("udp proxy listen port [%d]", pxy.cfg.RemotePort)
  68. pxy.udpConn = udpConn
  69. pxy.sendCh = make(chan *msg.UDPPacket, 1024)
  70. pxy.readCh = make(chan *msg.UDPPacket, 1024)
  71. pxy.checkCloseCh = make(chan int)
  72. // read message from workConn, if it returns any error, notify proxy to start a new workConn
  73. workConnReaderFn := func(conn net.Conn) {
  74. for {
  75. var (
  76. rawMsg msg.Message
  77. errRet error
  78. )
  79. xl.Trace("loop waiting message from udp workConn")
  80. // client will send heartbeat in workConn for keeping alive
  81. conn.SetReadDeadline(time.Now().Add(time.Duration(60) * time.Second))
  82. if rawMsg, errRet = msg.ReadMsg(conn); errRet != nil {
  83. xl.Warn("read from workConn for udp error: %v", errRet)
  84. conn.Close()
  85. // notify proxy to start a new work connection
  86. // ignore error here, it means the proxy is closed
  87. errors.PanicToError(func() {
  88. pxy.checkCloseCh <- 1
  89. })
  90. return
  91. }
  92. conn.SetReadDeadline(time.Time{})
  93. switch m := rawMsg.(type) {
  94. case *msg.Ping:
  95. xl.Trace("udp work conn get ping message")
  96. continue
  97. case *msg.UDPPacket:
  98. if errRet := errors.PanicToError(func() {
  99. xl.Trace("get udp message from workConn: %s", m.Content)
  100. pxy.readCh <- m
  101. metrics.Server.AddTrafficOut(
  102. pxy.GetName(),
  103. pxy.GetConf().GetBaseInfo().ProxyType,
  104. int64(len(m.Content)),
  105. )
  106. }); errRet != nil {
  107. conn.Close()
  108. xl.Info("reader goroutine for udp work connection closed")
  109. return
  110. }
  111. }
  112. }
  113. }
  114. // send message to workConn
  115. workConnSenderFn := func(conn net.Conn, ctx context.Context) {
  116. var errRet error
  117. for {
  118. select {
  119. case udpMsg, ok := <-pxy.sendCh:
  120. if !ok {
  121. xl.Info("sender goroutine for udp work connection closed")
  122. return
  123. }
  124. if errRet = msg.WriteMsg(conn, udpMsg); errRet != nil {
  125. xl.Info("sender goroutine for udp work connection closed: %v", errRet)
  126. conn.Close()
  127. return
  128. }
  129. xl.Trace("send message to udp workConn: %s", udpMsg.Content)
  130. metrics.Server.AddTrafficIn(
  131. pxy.GetName(),
  132. pxy.GetConf().GetBaseInfo().ProxyType,
  133. int64(len(udpMsg.Content)),
  134. )
  135. continue
  136. case <-ctx.Done():
  137. xl.Info("sender goroutine for udp work connection closed")
  138. return
  139. }
  140. }
  141. }
  142. go func() {
  143. // Sleep a while for waiting control send the NewProxyResp to client.
  144. time.Sleep(500 * time.Millisecond)
  145. for {
  146. workConn, err := pxy.GetWorkConnFromPool(nil, nil)
  147. if err != nil {
  148. time.Sleep(1 * time.Second)
  149. // check if proxy is closed
  150. select {
  151. case _, ok := <-pxy.checkCloseCh:
  152. if !ok {
  153. return
  154. }
  155. default:
  156. }
  157. continue
  158. }
  159. // close the old workConn and replac it with a new one
  160. if pxy.workConn != nil {
  161. pxy.workConn.Close()
  162. }
  163. pxy.workConn = workConn
  164. ctx, cancel := context.WithCancel(context.Background())
  165. go workConnReaderFn(workConn)
  166. go workConnSenderFn(workConn, ctx)
  167. _, ok := <-pxy.checkCloseCh
  168. cancel()
  169. if !ok {
  170. return
  171. }
  172. }
  173. }()
  174. // Read from user connections and send wrapped udp message to sendCh (forwarded by workConn).
  175. // Client will transfor udp message to local udp service and waiting for response for a while.
  176. // Response will be wrapped to be forwarded by work connection to server.
  177. // Close readCh and sendCh at the end.
  178. go func() {
  179. udp.ForwardUserConn(udpConn, pxy.readCh, pxy.sendCh, int(pxy.serverCfg.UDPPacketSize))
  180. pxy.Close()
  181. }()
  182. return remoteAddr, nil
  183. }
  184. func (pxy *UDPProxy) GetConf() config.ProxyConf {
  185. return pxy.cfg
  186. }
  187. func (pxy *UDPProxy) Close() {
  188. pxy.mu.Lock()
  189. defer pxy.mu.Unlock()
  190. if !pxy.isClosed {
  191. pxy.isClosed = true
  192. pxy.BaseProxy.Close()
  193. if pxy.workConn != nil {
  194. pxy.workConn.Close()
  195. }
  196. pxy.udpConn.Close()
  197. // all channels only closed here
  198. close(pxy.checkCloseCh)
  199. close(pxy.readCh)
  200. close(pxy.sendCh)
  201. }
  202. pxy.rc.UDPPortManager.Release(pxy.realPort)
  203. }