frp/test/e2e/pkg/port/port.go

116 lines
2.1 KiB
Go
Raw Normal View History

2020-06-02 22:48:55 +08:00
package port
import (
"fmt"
"net"
"strconv"
2021-03-31 16:57:39 +08:00
"sync"
2020-06-02 22:48:55 +08:00
"k8s.io/apimachinery/pkg/util/sets"
)
type Allocator struct {
2023-02-26 02:54:53 +08:00
reserved sets.Set[int]
used sets.Set[int]
2021-03-31 16:57:39 +08:00
mu sync.Mutex
2020-06-02 22:48:55 +08:00
}
// 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{
2023-02-26 02:54:53 +08:00
reserved: sets.New[int](),
used: sets.New[int](),
2020-06-02 22:48:55 +08:00
}
2021-06-18 16:48:36 +08:00
2020-06-02 22:48:55 +08:00
for i := from; i <= to; i++ {
if i%mod == index {
pa.reserved.Insert(i)
}
}
return pa
}
func (pa *Allocator) Get() int {
2021-03-31 16:57:39 +08:00
return pa.GetByName("")
}
func (pa *Allocator) GetByName(portName string) int {
var builder *nameBuilder
if portName == "" {
builder = &nameBuilder{}
} else {
var err error
builder, err = unmarshalFromName(portName)
if err != nil {
fmt.Println(err, portName)
return 0
}
}
pa.mu.Lock()
defer pa.mu.Unlock()
2021-06-18 16:48:36 +08:00
for i := 0; i < 20; i++ {
2021-03-31 16:57:39 +08:00
port := pa.getByRange(builder.rangePortFrom, builder.rangePortTo)
2020-06-02 22:48:55 +08:00
if port == 0 {
return 0
}
l, err := net.Listen("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)))
2020-06-02 22:48:55 +08:00
if err != nil {
// Maybe not controlled by us, mark it used.
pa.used.Insert(port)
continue
}
l.Close()
2021-03-31 16:57:39 +08:00
udpAddr, err := net.ResolveUDPAddr("udp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)))
2021-03-31 16:57:39 +08:00
if err != nil {
continue
}
udpConn, err := net.ListenUDP("udp", udpAddr)
if err != nil {
// Maybe not controlled by us, mark it used.
pa.used.Insert(port)
continue
}
udpConn.Close()
2020-06-02 22:48:55 +08:00
pa.used.Insert(port)
return port
}
return 0
}
2021-03-31 16:57:39 +08:00
func (pa *Allocator) getByRange(from, to int) int {
if from <= 0 {
port, _ := pa.reserved.PopAny()
return port
}
// choose a random port between from - to
ports := pa.reserved.UnsortedList()
for _, port := range ports {
if port >= from && port <= to {
return port
}
}
return 0
}
2020-06-02 22:48:55 +08:00
func (pa *Allocator) Release(port int) {
2021-03-31 16:57:39 +08:00
if port <= 0 {
return
}
pa.mu.Lock()
defer pa.mu.Unlock()
2020-06-02 22:48:55 +08:00
if pa.used.Has(port) {
pa.used.Delete(port)
pa.reserved.Insert(port)
}
}