update

utils/util/util.go

@@ -19,6 +19,8 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
+	"strconv"
+	"strings"
 )
 
 // RandId return a rand string used in frp.
@@ -45,3 +47,66 @@ func GetAuthKey(token string, timestamp int64) (key string) {
 	data := md5Ctx.Sum(nil)
 	return hex.EncodeToString(data)
 }
+
+// for example: rangeStr is "1000-2000,2001,2002,3000-4000", return an array as port ranges.
+func GetPortRanges(rangeStr string) (portRanges [][2]int64, err error) {
+	// for example: 1000-2000,2001,2002,3000-4000
+	rangeArray := strings.Split(rangeStr, ",")
+	for _, portRangeStr := range rangeArray {
+		// 1000-2000 or 2001
+		portArray := strings.Split(portRangeStr, "-")
+		// length: only 1 or 2 is correct
+		rangeType := len(portArray)
+		if rangeType == 1 {
+			singlePort, err := strconv.ParseInt(portArray[0], 10, 64)
+			if err != nil {
+				return [][2]int64{}, err
+			}
+			portRanges = append(portRanges, [2]int64{singlePort, singlePort})
+		} else if rangeType == 2 {
+			min, err := strconv.ParseInt(portArray[0], 10, 64)
+			if err != nil {
+				return [][2]int64{}, err
+			}
+			max, err := strconv.ParseInt(portArray[1], 10, 64)
+			if err != nil {
+				return [][2]int64{}, err
+			}
+			if max < min {
+				return [][2]int64{}, fmt.Errorf("range incorrect")
+			}
+			portRanges = append(portRanges, [2]int64{min, max})
+		} else {
+			return [][2]int64{}, fmt.Errorf("format error")
+		}
+	}
+	return portRanges, nil
+}
+
+func ContainsPort(portRanges [][2]int64, port int64) bool {
+	for _, pr := range portRanges {
+		if port >= pr[0] && port <= pr[1] {
+			return true
+		}
+	}
+	return false
+}
+
+func PortRangesCut(portRanges [][2]int64, port int64) [][2]int64 {
+	var tmpRanges [][2]int64
+	for _, pr := range portRanges {
+		if port >= pr[0] && port <= pr[1] {
+			leftRange := [2]int64{pr[0], port - 1}
+			rightRange := [2]int64{port + 1, pr[1]}
+			if leftRange[0] <= leftRange[1] {
+				tmpRanges = append(tmpRanges, leftRange)
+			}
+			if rightRange[0] <= rightRange[1] {
+				tmpRanges = append(tmpRanges, rightRange)
+			}
+		} else {
+			tmpRanges = append(tmpRanges, pr)
+		}
+	}
+	return tmpRanges
+}

utils/util/util_test.go

@@ -20,3 +20,67 @@ func TestGetAuthKey(t *testing.T) {
 	t.Log(key)
 	assert.Equal("6df41a43725f0c770fd56379e12acf8c", key)
 }
+
+func TestGetPortRanges(t *testing.T) {
+	assert := assert.New(t)
+
+	rangesStr := "2000-3000,3001,4000-50000"
+	expect := [][2]int64{
+		[2]int64{2000, 3000},
+		[2]int64{3001, 3001},
+		[2]int64{4000, 50000},
+	}
+	actual, err := GetPortRanges(rangesStr)
+	assert.Nil(err)
+	t.Log(actual)
+	assert.Equal(expect, actual)
+}
+
+func TestContainsPort(t *testing.T) {
+	assert := assert.New(t)
+
+	rangesStr := "2000-3000,3001,4000-50000"
+	portRanges, err := GetPortRanges(rangesStr)
+	assert.Nil(err)
+
+	type Case struct {
+		Port   int64
+		Answer bool
+	}
+	cases := []Case{
+		Case{
+			Port:   3001,
+			Answer: true,
+		},
+		Case{
+			Port:   3002,
+			Answer: false,
+		},
+		Case{
+			Port:   44444,
+			Answer: true,
+		},
+	}
+	for _, elem := range cases {
+		ok := ContainsPort(portRanges, elem.Port)
+		assert.Equal(elem.Answer, ok)
+	}
+}
+
+func TestPortRangesCut(t *testing.T) {
+	assert := assert.New(t)
+
+	rangesStr := "2000-3000,3001,4000-50000"
+	portRanges, err := GetPortRanges(rangesStr)
+	assert.Nil(err)
+
+	expect := [][2]int64{
+		[2]int64{2000, 3000},
+		[2]int64{3001, 3001},
+		[2]int64{4000, 44443},
+		[2]int64{44445, 50000},
+	}
+	actual := PortRangesCut(portRanges, 44444)
+	t.Log(actual)
+	assert.Equal(expect, actual)
+}