add e2e tests (#2334)

test/e2e/pkg/port/port.go

@@ -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)

test/e2e/pkg/port/util.go

@@ -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()
+}

test/e2e/pkg/request/request.go

@@ -4,12 +4,108 @@ import (
 	"fmt"
 	"net"
 	"time"
+
+	libnet "github.com/fatedier/golib/net"
 )
 
-func SendTCPRequest(port int, content []byte, timeout time.Duration) (string, error) {
+type Request struct {
+	protocol  string
+	addr      string
+	port      int
+	body      []byte
+	timeout   time.Duration
+	proxyURL  string
+	proxyHost string
+}
+
+func New() *Request {
+	return &Request{
+		protocol: "tcp",
+	}
+}
+
+func (r *Request) Protocol(protocol string) *Request {
+	r.protocol = protocol
+	return r
+}
+
+func (r *Request) TCP() *Request {
+	r.protocol = "tcp"
+	return r
+}
+
+func (r *Request) UDP() *Request {
+	r.protocol = "udp"
+	return r
+}
+
+func (r *Request) Proxy(url, host string) *Request {
+	r.proxyURL = url
+	r.proxyHost = host
+	return r
+}
+
+func (r *Request) Addr(addr string) *Request {
+	r.addr = addr
+	return r
+}
+
+func (r *Request) Port(port int) *Request {
+	r.port = port
+	return r
+}
+
+func (r *Request) Timeout(timeout time.Duration) *Request {
+	r.timeout = timeout
+	return r
+}
+
+func (r *Request) Body(content []byte) *Request {
+	r.body = content
+	return r
+}
+
+func (r *Request) Do() ([]byte, error) {
+	var (
+		conn net.Conn
+		err  error
+	)
+	if len(r.proxyURL) > 0 {
+		if r.protocol != "tcp" {
+			return nil, fmt.Errorf("only tcp protocol is allowed for proxy")
+		}
+		conn, err = libnet.DialTcpByProxy(r.proxyURL, r.proxyHost)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		if r.addr == "" {
+			r.addr = fmt.Sprintf("127.0.0.1:%d", r.port)
+		}
+		switch r.protocol {
+		case "tcp":
+			conn, err = net.Dial("tcp", r.addr)
+		case "udp":
+			conn, err = net.Dial("udp", r.addr)
+		default:
+			return nil, fmt.Errorf("invalid protocol")
+		}
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	defer conn.Close()
+	if r.timeout > 0 {
+		conn.SetDeadline(time.Now().Add(r.timeout))
+	}
+	return sendRequestByConn(conn, r.body)
+}
+
+func SendTCPRequest(port int, content []byte, timeout time.Duration) ([]byte, error) {
 	c, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", port))
 	if err != nil {
-		return "", fmt.Errorf("connect to tcp server error: %v", err)
+		return nil, fmt.Errorf("connect to tcp server error: %v", err)
 	}
 	defer c.Close()
 
@@ -17,10 +113,10 @@ func SendTCPRequest(port int, content []byte, timeout time.Duration) (string, er
 	return sendRequestByConn(c, content)
 }
 
-func SendUDPRequest(port int, content []byte, timeout time.Duration) (string, error) {
+func SendUDPRequest(port int, content []byte, timeout time.Duration) ([]byte, error) {
 	c, err := net.Dial("udp", fmt.Sprintf("127.0.0.1:%d", port))
 	if err != nil {
-		return "", fmt.Errorf("connect to udp server error:  %v", err)
+		return nil, fmt.Errorf("connect to udp server error:  %v", err)
 	}
 	defer c.Close()
 
@@ -28,13 +124,16 @@ func SendUDPRequest(port int, content []byte, timeout time.Duration) (string, er
 	return sendRequestByConn(c, content)
 }
 
-func sendRequestByConn(c net.Conn, content []byte) (string, error) {
-	c.Write(content)
+func sendRequestByConn(c net.Conn, content []byte) ([]byte, error) {
+	_, err := c.Write(content)
+	if err != nil {
+		return nil, fmt.Errorf("write error: %v", err)
+	}
 
 	buf := make([]byte, 2048)
 	n, err := c.Read(buf)
 	if err != nil {
-		return "", fmt.Errorf("read error: %v", err)
+		return nil, fmt.Errorf("read error: %v", err)
 	}
-	return string(buf[:n]), nil
+	return buf[:n], nil
 }