feature: node-agent add port available check

Signed-off-by: wangdepeng <[email protected]>
(cherry picked from commit fa84fc1)

cmd/kubenest/node-agent/app/client/client.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strings"
 	"sync"
 
 	"github.com/gorilla/websocket"
@@ -51,6 +52,12 @@ var (
 		Long:  "execute command on remote server use pyt",
 		RunE:  cmdTtyRun,
 	}
+	checkCmd = &cobra.Command{
+		Use:   "check",
+		Short: "Check the port is open or not",
+		Long:  "Check the port can be assign or not",
+		RunE:  cmdCheckRun,
+	}
 	wg sync.WaitGroup
 
 	wsAddr    []string // websocket client connect address list
@@ -59,6 +66,35 @@ var (
 	params    []string // New slice to hold multiple command parameters
 	operation string   // operation for client to execute
 )
+
+func cmdCheckRun(cmd *cobra.Command, args []string) error {
+	if len(params) != 1 {
+		log.Errorf("port list is required and port list size must not be greater than 1")
+		return fmt.Errorf("port list is required and port list size must not be greater than 1")
+	}
+	auth, err := getAuth(cmd)
+	if err != nil {
+		return err
+	}
+	headers := http.Header{
+		"Authorization": {"Basic " + auth},
+	}
+	for _, addr := range wsAddr {
+		wg.Add(1)
+		go func(addr string) {
+			defer wg.Done()
+			wsURL := fmt.Sprintf("wss://%s/check/?port=%s", addr, params[0])
+			fmt.Println("Checking port:", wsURL)
+			err := connectAndHandleMessages(wsURL, headers)
+			if err != nil {
+				log.Errorf("failed to check port: %v on %s: %v\n", err, addr, strings.Join(params, "&"))
+			}
+		}(addr)
+	}
+	wg.Wait()
+	return nil
+}
+
 var uniqueValuesMap = make(map[string]bool)
 var dialer = websocket.DefaultDialer
 
@@ -91,6 +127,9 @@ func init() {
 	shCmd.Flags().StringVarP(&operation, "operation", "o", "", "Operation to perform")
 	_ = shCmd.MarkFlagRequired("addr")
 
+	checkCmd.Flags().StringArrayVarP(&params, "param", "r", []string{}, "Command parameters")
+	_ = checkCmd.MarkFlagRequired("addr")
+
 	uploadCmd.Flags().StringVarP(&fileName, "name", "n", "", "Name of the file to upload")
 	uploadCmd.Flags().StringVarP(&filePath, "path", "f", "", "Path to the file to upload")
 	// avoid can't show subcommand help and execute subcommand
@@ -105,6 +144,7 @@ func init() {
 	ClientCmd.AddCommand(shCmd)
 	ClientCmd.AddCommand(uploadCmd)
 	ClientCmd.AddCommand(ttyCmd)
+	ClientCmd.AddCommand(checkCmd)
 }
 
 func cmdTtyRun(cmd *cobra.Command, args []string) error {
@@ -247,7 +287,7 @@ func executeWebSocketCommand(auth string) error {
 	if len(params) > 1 {
 		paramsStr := "args="
 		for _, param := range params {
-			paramsStr += param + "&&args="
+			paramsStr += url.QueryEscape(param) + "&&args="
 		}
 		paramsStr = paramsStr[:len(paramsStr)-7]
 		cmdStr = fmt.Sprintf("command=%s&&%s", operation, paramsStr)

cmd/kubenest/node-agent/app/serve/serve.go

@@ -7,6 +7,7 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
@@ -118,6 +119,8 @@ func Start(addr, certFile, keyFile, user, password string) error {
 			handleScript(conn, queryParams, []string{"sh"})
 		case strings.HasPrefix(r.URL.Path, "/tty"):
 			handleTty(conn, queryParams)
+		case strings.HasPrefix(r.URL.Path, "/check"):
+			handleCheck(conn, queryParams)
 		default:
 			_ = conn.WriteMessage(websocket.TextMessage, []byte("Invalid path"))
 		}
@@ -141,6 +144,25 @@ func Start(addr, certFile, keyFile, user, password string) error {
 	return err
 }
 
+func handleCheck(conn *websocket.Conn, params url.Values) {
+	port := params.Get("port")
+	if len(port) == 0 {
+		log.Errorf("port is required")
+		return
+	}
+	log.Infof("Check port %s", port)
+	address := fmt.Sprintf(":%s", port)
+	listener, err := net.Listen("tcp", address)
+	if err != nil {
+		log.Infof("port not avalible %s %v", address, err)
+		_ = conn.WriteMessage(websocket.BinaryMessage, []byte("1"))
+		return
+	}
+	defer listener.Close()
+	log.Infof("port avalible %s", address)
+	_ = conn.WriteMessage(websocket.BinaryMessage, []byte("0"))
+}
+
 func handleTty(conn *websocket.Conn, queryParams url.Values) {
 	entrypoint := queryParams.Get("command")
 	if len(entrypoint) == 0 {