frp/test/e2e/pkg/port/port.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)
	}
}