1
0

manager.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  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 server
  15. import (
  16. "fmt"
  17. "io"
  18. "sync"
  19. frpIo "github.com/fatedier/frp/utils/io"
  20. frpNet "github.com/fatedier/frp/utils/net"
  21. "github.com/fatedier/frp/utils/util"
  22. )
  23. type ControlManager struct {
  24. // controls indexed by run id
  25. ctlsByRunId map[string]*Control
  26. mu sync.RWMutex
  27. }
  28. func NewControlManager() *ControlManager {
  29. return &ControlManager{
  30. ctlsByRunId: make(map[string]*Control),
  31. }
  32. }
  33. func (cm *ControlManager) Add(runId string, ctl *Control) (oldCtl *Control) {
  34. cm.mu.Lock()
  35. defer cm.mu.Unlock()
  36. oldCtl, ok := cm.ctlsByRunId[runId]
  37. if ok {
  38. oldCtl.Replaced(ctl)
  39. }
  40. cm.ctlsByRunId[runId] = ctl
  41. return
  42. }
  43. func (cm *ControlManager) GetById(runId string) (ctl *Control, ok bool) {
  44. cm.mu.RLock()
  45. defer cm.mu.RUnlock()
  46. ctl, ok = cm.ctlsByRunId[runId]
  47. return
  48. }
  49. type ProxyManager struct {
  50. // proxies indexed by proxy name
  51. pxys map[string]Proxy
  52. mu sync.RWMutex
  53. }
  54. func NewProxyManager() *ProxyManager {
  55. return &ProxyManager{
  56. pxys: make(map[string]Proxy),
  57. }
  58. }
  59. func (pm *ProxyManager) Add(name string, pxy Proxy) error {
  60. pm.mu.Lock()
  61. defer pm.mu.Unlock()
  62. if _, ok := pm.pxys[name]; ok {
  63. return fmt.Errorf("proxy name [%s] is already in use", name)
  64. }
  65. pm.pxys[name] = pxy
  66. return nil
  67. }
  68. func (pm *ProxyManager) Del(name string) {
  69. pm.mu.Lock()
  70. defer pm.mu.Unlock()
  71. delete(pm.pxys, name)
  72. }
  73. func (pm *ProxyManager) GetByName(name string) (pxy Proxy, ok bool) {
  74. pm.mu.RLock()
  75. defer pm.mu.RUnlock()
  76. pxy, ok = pm.pxys[name]
  77. return
  78. }
  79. // Manager for visitor listeners.
  80. type VisitorManager struct {
  81. visitorListeners map[string]*frpNet.CustomListener
  82. skMap map[string]string
  83. mu sync.RWMutex
  84. }
  85. func NewVisitorManager() *VisitorManager {
  86. return &VisitorManager{
  87. visitorListeners: make(map[string]*frpNet.CustomListener),
  88. skMap: make(map[string]string),
  89. }
  90. }
  91. func (vm *VisitorManager) Listen(name string, sk string) (l *frpNet.CustomListener, err error) {
  92. vm.mu.Lock()
  93. defer vm.mu.Unlock()
  94. if _, ok := vm.visitorListeners[name]; ok {
  95. err = fmt.Errorf("custom listener for [%s] is repeated", name)
  96. return
  97. }
  98. l = frpNet.NewCustomListener()
  99. vm.visitorListeners[name] = l
  100. vm.skMap[name] = sk
  101. return
  102. }
  103. func (vm *VisitorManager) NewConn(name string, conn frpNet.Conn, timestamp int64, signKey string,
  104. useEncryption bool, useCompression bool) (err error) {
  105. vm.mu.RLock()
  106. defer vm.mu.RUnlock()
  107. if l, ok := vm.visitorListeners[name]; ok {
  108. var sk string
  109. if sk = vm.skMap[name]; util.GetAuthKey(sk, timestamp) != signKey {
  110. err = fmt.Errorf("visitor connection of [%s] auth failed", name)
  111. return
  112. }
  113. var rwc io.ReadWriteCloser = conn
  114. if useEncryption {
  115. if rwc, err = frpIo.WithEncryption(rwc, []byte(sk)); err != nil {
  116. err = fmt.Errorf("create encryption connection failed: %v", err)
  117. return
  118. }
  119. }
  120. if useCompression {
  121. rwc = frpIo.WithCompression(rwc)
  122. }
  123. err = l.PutConn(frpNet.WrapReadWriteCloserToConn(rwc, conn))
  124. } else {
  125. err = fmt.Errorf("custom listener for [%s] doesn't exist", name)
  126. return
  127. }
  128. return
  129. }
  130. func (vm *VisitorManager) CloseListener(name string) {
  131. vm.mu.Lock()
  132. defer vm.mu.Unlock()
  133. delete(vm.visitorListeners, name)
  134. delete(vm.skMap, name)
  135. }