mirror of https://gitee.com/IrisVega/frp.git
58 lines
1.0 KiB
Go
58 lines
1.0 KiB
Go
package port
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
|
|
"k8s.io/apimachinery/pkg/util/sets"
|
|
)
|
|
|
|
type Allocator struct {
|
|
reserved sets.Int
|
|
used sets.Int
|
|
}
|
|
|
|
// NewAllocator return a port allocator for testing.
|
|
// Example: from: 10, to: 20, mod 4, index 1
|
|
// Reserved ports: 13, 17
|
|
func NewAllocator(from int, to int, mod int, index int) *Allocator {
|
|
pa := &Allocator{
|
|
reserved: sets.NewInt(),
|
|
used: sets.NewInt(),
|
|
}
|
|
for i := from; i <= to; i++ {
|
|
if i%mod == index {
|
|
pa.reserved.Insert(i)
|
|
}
|
|
}
|
|
return pa
|
|
}
|
|
|
|
func (pa *Allocator) Get() int {
|
|
for i := 0; i < 10; i++ {
|
|
port, _ := pa.reserved.PopAny()
|
|
if port == 0 {
|
|
return 0
|
|
}
|
|
|
|
// TODO: Distinguish between TCP and UDP
|
|
l, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
|
|
if err != nil {
|
|
// Maybe not controlled by us, mark it used.
|
|
pa.used.Insert(port)
|
|
continue
|
|
}
|
|
l.Close()
|
|
pa.used.Insert(port)
|
|
return port
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func (pa *Allocator) Release(port int) {
|
|
if pa.used.Has(port) {
|
|
pa.used.Delete(port)
|
|
pa.reserved.Insert(port)
|
|
}
|
|
}
|