portallocator.go 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. package portallocator
  2. import (
  3. "bufio"
  4. "errors"
  5. "fmt"
  6. "net"
  7. "os"
  8. "sync"
  9. log "github.com/Sirupsen/logrus"
  10. )
  11. const (
  12. DefaultPortRangeStart = 49153
  13. DefaultPortRangeEnd = 65535
  14. )
  15. var (
  16. beginPortRange = DefaultPortRangeStart
  17. endPortRange = DefaultPortRangeEnd
  18. )
  19. type portMap struct {
  20. p map[int]struct{}
  21. last int
  22. }
  23. func newPortMap() *portMap {
  24. return &portMap{
  25. p: map[int]struct{}{},
  26. last: endPortRange,
  27. }
  28. }
  29. type protoMap map[string]*portMap
  30. func newProtoMap() protoMap {
  31. return protoMap{
  32. "tcp": newPortMap(),
  33. "udp": newPortMap(),
  34. }
  35. }
  36. type ipMapping map[string]protoMap
  37. var (
  38. ErrAllPortsAllocated = errors.New("all ports are allocated")
  39. ErrUnknownProtocol = errors.New("unknown protocol")
  40. )
  41. var (
  42. mutex sync.Mutex
  43. defaultIP = net.ParseIP("0.0.0.0")
  44. globalMap = ipMapping{}
  45. )
  46. type ErrPortAlreadyAllocated struct {
  47. ip string
  48. port int
  49. }
  50. func NewErrPortAlreadyAllocated(ip string, port int) ErrPortAlreadyAllocated {
  51. return ErrPortAlreadyAllocated{
  52. ip: ip,
  53. port: port,
  54. }
  55. }
  56. func init() {
  57. const portRangeKernelParam = "/proc/sys/net/ipv4/ip_local_port_range"
  58. portRangeFallback := fmt.Sprintf("using fallback port range %d-%d", beginPortRange, endPortRange)
  59. file, err := os.Open(portRangeKernelParam)
  60. if err != nil {
  61. log.Warnf("port allocator - %s due to error: %v", portRangeFallback, err)
  62. return
  63. }
  64. var start, end int
  65. n, err := fmt.Fscanf(bufio.NewReader(file), "%d\t%d", &start, &end)
  66. if n != 2 || err != nil {
  67. if err == nil {
  68. err = fmt.Errorf("unexpected count of parsed numbers (%d)", n)
  69. }
  70. log.Errorf("port allocator - failed to parse system ephemeral port range from %s - %s: %v", portRangeKernelParam, portRangeFallback, err)
  71. return
  72. }
  73. beginPortRange = start
  74. endPortRange = end
  75. }
  76. func PortRange() (int, int) {
  77. return beginPortRange, endPortRange
  78. }
  79. func (e ErrPortAlreadyAllocated) IP() string {
  80. return e.ip
  81. }
  82. func (e ErrPortAlreadyAllocated) Port() int {
  83. return e.port
  84. }
  85. func (e ErrPortAlreadyAllocated) IPPort() string {
  86. return fmt.Sprintf("%s:%d", e.ip, e.port)
  87. }
  88. func (e ErrPortAlreadyAllocated) Error() string {
  89. return fmt.Sprintf("Bind for %s:%d failed: port is already allocated", e.ip, e.port)
  90. }
  91. // RequestPort requests new port from global ports pool for specified ip and proto.
  92. // If port is 0 it returns first free port. Otherwise it cheks port availability
  93. // in pool and return that port or error if port is already busy.
  94. func RequestPort(ip net.IP, proto string, port int) (int, error) {
  95. mutex.Lock()
  96. defer mutex.Unlock()
  97. if proto != "tcp" && proto != "udp" {
  98. return 0, ErrUnknownProtocol
  99. }
  100. if ip == nil {
  101. ip = defaultIP
  102. }
  103. ipstr := ip.String()
  104. protomap, ok := globalMap[ipstr]
  105. if !ok {
  106. protomap = newProtoMap()
  107. globalMap[ipstr] = protomap
  108. }
  109. mapping := protomap[proto]
  110. if port > 0 {
  111. if _, ok := mapping.p[port]; !ok {
  112. mapping.p[port] = struct{}{}
  113. return port, nil
  114. }
  115. return 0, NewErrPortAlreadyAllocated(ipstr, port)
  116. }
  117. port, err := mapping.findPort()
  118. if err != nil {
  119. return 0, err
  120. }
  121. return port, nil
  122. }
  123. // ReleasePort releases port from global ports pool for specified ip and proto.
  124. func ReleasePort(ip net.IP, proto string, port int) error {
  125. mutex.Lock()
  126. defer mutex.Unlock()
  127. if ip == nil {
  128. ip = defaultIP
  129. }
  130. protomap, ok := globalMap[ip.String()]
  131. if !ok {
  132. return nil
  133. }
  134. delete(protomap[proto].p, port)
  135. return nil
  136. }
  137. // ReleaseAll releases all ports for all ips.
  138. func ReleaseAll() error {
  139. mutex.Lock()
  140. globalMap = ipMapping{}
  141. mutex.Unlock()
  142. return nil
  143. }
  144. func (pm *portMap) findPort() (int, error) {
  145. port := pm.last
  146. for i := 0; i <= endPortRange-beginPortRange; i++ {
  147. port++
  148. if port > endPortRange {
  149. port = beginPortRange
  150. }
  151. if _, ok := pm.p[port]; !ok {
  152. pm.p[port] = struct{}{}
  153. pm.last = port
  154. return port, nil
  155. }
  156. }
  157. return 0, ErrAllPortsAllocated
  158. }