mirror of
https://github.com/fatedier/frp.git
synced 2026-01-11 22:23:12 +00:00
add e2e tests (#2334)
This commit is contained in:
@@ -3,6 +3,7 @@ package port
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"k8s.io/apimachinery/pkg/util/sets"
|
||||
)
|
||||
@@ -10,6 +11,7 @@ import (
|
||||
type Allocator struct {
|
||||
reserved sets.Int
|
||||
used sets.Int
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewAllocator return a port allocator for testing.
|
||||
@@ -29,8 +31,27 @@ func NewAllocator(from int, to int, mod int, index int) *Allocator {
|
||||
}
|
||||
|
||||
func (pa *Allocator) Get() int {
|
||||
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()
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
port, _ := pa.reserved.PopAny()
|
||||
port := pa.getByRange(builder.rangePortFrom, builder.rangePortTo)
|
||||
if port == 0 {
|
||||
return 0
|
||||
}
|
||||
@@ -43,13 +64,49 @@ func (pa *Allocator) Get() int {
|
||||
continue
|
||||
}
|
||||
l.Close()
|
||||
|
||||
udpAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("127.0.0.1:%d", port))
|
||||
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()
|
||||
|
||||
pa.used.Insert(port)
|
||||
return port
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
func (pa *Allocator) Release(port int) {
|
||||
if port <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
pa.mu.Lock()
|
||||
defer pa.mu.Unlock()
|
||||
|
||||
if pa.used.Has(port) {
|
||||
pa.used.Delete(port)
|
||||
pa.reserved.Insert(port)
|
||||
|
||||
67
test/e2e/pkg/port/util.go
Normal file
67
test/e2e/pkg/port/util.go
Normal file
@@ -0,0 +1,67 @@
|
||||
package port
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
NameDelimiter = "_"
|
||||
)
|
||||
|
||||
type NameOption func(*nameBuilder) *nameBuilder
|
||||
|
||||
type nameBuilder struct {
|
||||
name string
|
||||
rangePortFrom int
|
||||
rangePortTo int
|
||||
}
|
||||
|
||||
func unmarshalFromName(name string) (*nameBuilder, error) {
|
||||
var builder nameBuilder
|
||||
arrs := strings.Split(name, NameDelimiter)
|
||||
switch len(arrs) {
|
||||
case 2:
|
||||
builder.name = arrs[1]
|
||||
case 4:
|
||||
builder.name = arrs[1]
|
||||
if fromPort, err := strconv.Atoi(arrs[2]); err != nil {
|
||||
return nil, fmt.Errorf("error range port from")
|
||||
} else {
|
||||
builder.rangePortFrom = fromPort
|
||||
}
|
||||
if toPort, err := strconv.Atoi(arrs[3]); err != nil {
|
||||
return nil, fmt.Errorf("error range port to")
|
||||
} else {
|
||||
builder.rangePortTo = toPort
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("error port name format")
|
||||
}
|
||||
return &builder, nil
|
||||
}
|
||||
|
||||
func (builder *nameBuilder) String() string {
|
||||
name := fmt.Sprintf("Port%s%s", NameDelimiter, builder.name)
|
||||
if builder.rangePortFrom > 0 && builder.rangePortTo > 0 && builder.rangePortTo > builder.rangePortFrom {
|
||||
name += fmt.Sprintf("%s%d%s%d", NameDelimiter, builder.rangePortFrom, NameDelimiter, builder.rangePortTo)
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
func WithRangePorts(from, to int) NameOption {
|
||||
return func(builder *nameBuilder) *nameBuilder {
|
||||
builder.rangePortFrom = from
|
||||
builder.rangePortTo = to
|
||||
return builder
|
||||
}
|
||||
}
|
||||
|
||||
func GenName(name string, options ...NameOption) string {
|
||||
builder := &nameBuilder{name: name}
|
||||
for _, option := range options {
|
||||
option(builder)
|
||||
}
|
||||
return builder.String()
|
||||
}
|
||||
Reference in New Issue
Block a user