portallocator.go 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  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. file, err := os.Open(portRangeKernelParam)
  59. if err != nil {
  60. log.Warnf("Failed to read %s kernel parameter: %v", portRangeKernelParam, err)
  61. return
  62. }
  63. var start, end int
  64. n, err := fmt.Fscanf(bufio.NewReader(file), "%d\t%d", &start, &end)
  65. if n != 2 || err != nil {
  66. if err == nil {
  67. err = fmt.Errorf("unexpected count of parsed numbers (%d)", n)
  68. }
  69. log.Errorf("Failed to parse port range from %s: %v", portRangeKernelParam, err)
  70. return
  71. }
  72. beginPortRange = start
  73. endPortRange = end
  74. }
  75. func PortRange() (int, int) {
  76. return beginPortRange, endPortRange
  77. }
  78. func (e ErrPortAlreadyAllocated) IP() string {
  79. return e.ip
  80. }
  81. func (e ErrPortAlreadyAllocated) Port() int {
  82. return e.port
  83. }
  84. func (e ErrPortAlreadyAllocated) IPPort() string {
  85. return fmt.Sprintf("%s:%d", e.ip, e.port)
  86. }
  87. func (e ErrPortAlreadyAllocated) Error() string {
  88. return fmt.Sprintf("Bind for %s:%d failed: port is already allocated", e.ip, e.port)
  89. }
  90. // RequestPort requests new port from global ports pool for specified ip and proto.
  91. // If port is 0 it returns first free port. Otherwise it cheks port availability
  92. // in pool and return that port or error if port is already busy.
  93. func RequestPort(ip net.IP, proto string, port int) (int, error) {
  94. mutex.Lock()
  95. defer mutex.Unlock()
  96. if proto != "tcp" && proto != "udp" {
  97. return 0, ErrUnknownProtocol
  98. }
  99. if ip == nil {
  100. ip = defaultIP
  101. }
  102. ipstr := ip.String()
  103. protomap, ok := globalMap[ipstr]
  104. if !ok {
  105. protomap = newProtoMap()
  106. globalMap[ipstr] = protomap
  107. }
  108. mapping := protomap[proto]
  109. if port > 0 {
  110. if _, ok := mapping.p[port]; !ok {
  111. mapping.p[port] = struct{}{}
  112. return port, nil
  113. }
  114. return 0, NewErrPortAlreadyAllocated(ipstr, port)
  115. }
  116. port, err := mapping.findPort()
  117. if err != nil {
  118. return 0, err
  119. }
  120. return port, nil
  121. }
  122. // ReleasePort releases port from global ports pool for specified ip and proto.
  123. func ReleasePort(ip net.IP, proto string, port int) error {
  124. mutex.Lock()
  125. defer mutex.Unlock()
  126. if ip == nil {
  127. ip = defaultIP
  128. }
  129. protomap, ok := globalMap[ip.String()]
  130. if !ok {
  131. return nil
  132. }
  133. delete(protomap[proto].p, port)
  134. return nil
  135. }
  136. // ReleaseAll releases all ports for all ips.
  137. func ReleaseAll() error {
  138. mutex.Lock()
  139. globalMap = ipMapping{}
  140. mutex.Unlock()
  141. return nil
  142. }
  143. func (pm *portMap) findPort() (int, error) {
  144. port := pm.last
  145. for i := 0; i <= endPortRange-beginPortRange; i++ {
  146. port++
  147. if port > endPortRange {
  148. port = beginPortRange
  149. }
  150. if _, ok := pm.p[port]; !ok {
  151. pm.p[port] = struct{}{}
  152. pm.last = port
  153. return port, nil
  154. }
  155. }
  156. return 0, ErrAllPortsAllocated
  157. }