Some code changes

Author: nickname76

cmd/cli/main.go

@@ -80,6 +80,7 @@ func (uaf *uint16ArrFlags) Set(value string) error {
 
 var (
 	fwr          *p2pforwarder.Forwarder
+	fwrCancel    func()
 	connections  = make(map[string]func())
 	openTCPPorts = make(map[uint16]func())
 	openUDPPorts = make(map[uint16]func())
@@ -99,11 +100,9 @@ func main() {
 
 	zap.L().Info("Initialization...")
 
-	var (
-		cancel func()
-		err    error
-	)
-	fwr, cancel, err = p2pforwarder.NewForwarder()
+	var err error
+
+	fwr, fwrCancel, err = p2pforwarder.NewForwarder()
 	if err != nil {
 		zap.S().Error(err)
 	}
@@ -178,24 +177,26 @@ loop:
 		case str := <-cmdch:
 			executeCommand(str)
 		case <-termSignal:
-			zap.L().Info("Shutdown...")
-
-			cancel()
-
-			for _, cancel := range connections {
-				cancel()
-			}
-			for _, cancel := range openTCPPorts {
-				cancel()
-			}
-			for _, cancel := range openUDPPorts {
-				cancel()
-			}
-
+			shutdown()
 			break loop
 		}
 	}
 }
+func shutdown() {
+	zap.L().Info("Shutdown...")
+
+	fwrCancel()
+
+	for _, cancel := range connections {
+		cancel()
+	}
+	for _, cancel := range openTCPPorts {
+		cancel()
+	}
+	for _, cancel := range openUDPPorts {
+		cancel()
+	}
+}
 
 func parseArgs(argsStr string, n int) []string {
 	if n <= 0 {
@@ -233,8 +234,6 @@ func parseArgs(argsStr string, n int) []string {
 
 	if argStart != -1 {
 		args[a] = string(argsRunes[argStart:])
-		argStart = -1
-		a++
 	}
 
 	return args
@@ -247,96 +246,110 @@ func executeCommand(str string) {
 
 	switch cmd {
 	case "connect":
-		id := params[0]
-
-		zap.L().Info("Connecting to " + id)
+		cmdConnect(params)
+	case "disconnect":
+		cmdDisconnect(params)
+	case "open":
+		cmdOpen(params)
+	case "close":
+		cmdClose(params)
+	default:
+		zap.L().Info("")
+		zap.L().Info("Cli commands list:")
+		zap.L().Info("connect [ID_HERE]")
+		zap.L().Info("disconnect [ID_HERE]")
+		zap.L().Info("open [TCP_OR_UDP_HERE] [PORT_NUMBER_HERE]")
+		zap.L().Info("close [UDP_OR_UDP_HERE] [PORT_NUMBER_HERE]")
+		zap.L().Info("")
+	}
+}
 
-		listenip, cancel, err := fwr.Connect(id)
-		if err != nil {
-			zap.S().Error(err)
-			return
-		}
+func cmdConnect(params []string) {
+	id := params[0]
 
-		connections[id] = cancel
+	zap.L().Info("Connecting to " + id)
 
-		zap.L().Info("Connections to " + id + "'s ports are listened on " + listenip)
+	listenip, cancel, err := fwr.Connect(id)
+	if err != nil {
+		zap.S().Error(err)
+		return
+	}
 
-	case "disconnect":
-		id := params[0]
+	connections[id] = cancel
 
-		close := connections[id]
+	zap.L().Info("Connections to " + id + "'s ports are listened on " + listenip)
+}
 
-		if close == nil {
-			zap.L().Error("You are not connected to specified id")
-			return
-		}
+func cmdDisconnect(params []string) {
+	id := params[0]
 
-		zap.L().Info("Disconnecting from " + id)
+	close := connections[id]
 
-		close()
+	if close == nil {
+		zap.L().Error("You are not connected to specified id")
+		return
+	}
 
-		delete(connections, id)
+	zap.L().Info("Disconnecting from " + id)
 
-	case "open":
-		networkType := strings.ToLower(params[0])
+	close()
 
-		portStr := params[1]
-		portUint64, err := strconv.ParseUint(portStr, 10, 16)
-		if err != nil {
-			zap.S().Error(err)
-			return
-		}
-		port := uint16(portUint64)
+	delete(connections, id)
+}
 
-		zap.L().Info("Opening " + networkType + ":" + portStr)
+func cmdOpen(params []string) {
+	networkType := strings.ToLower(params[0])
 
-		cancel, err := fwr.OpenPort(networkType, port)
-		if err != nil {
-			zap.S().Error(err)
-			return
-		}
+	portStr := params[1]
+	portUint64, err := strconv.ParseUint(portStr, 10, 16)
+	if err != nil {
+		zap.S().Error(err)
+		return
+	}
+	port := uint16(portUint64)
 
-		switch networkType {
-		case "tcp":
-			openTCPPorts[port] = cancel
-		case "udp":
-			openUDPPorts[port] = cancel
-		}
-	case "close":
-		networkType := strings.ToLower(params[0])
-		portStr := params[1]
-		portUint64, err := strconv.ParseUint(portStr, 10, 16)
-		if err != nil {
-			zap.S().Error(err)
-			return
-		}
-		port := uint16(portUint64)
-
-		var cancel func()
-		switch networkType {
-		case "tcp":
-			cancel = openTCPPorts[port]
-		case "udp":
-			cancel = openUDPPorts[port]
-		}
+	zap.L().Info("Opening " + networkType + ":" + portStr)
 
-		if cancel == nil {
-			zap.L().Error("Specified port is not opened")
-			return
-		}
+	cancel, err := fwr.OpenPort(networkType, port)
+	if err != nil {
+		zap.S().Error(err)
+		return
+	}
 
-		zap.L().Info("Closing " + networkType + ":" + portStr)
+	switch networkType {
+	case "tcp":
+		openTCPPorts[port] = cancel
+	case "udp":
+		openUDPPorts[port] = cancel
+	}
+}
 
-		cancel()
+func cmdClose(params []string) {
+	networkType := strings.ToLower(params[0])
+	portStr := params[1]
+	portUint64, err := strconv.ParseUint(portStr, 10, 16)
+	if err != nil {
+		zap.S().Error(err)
+		return
+	}
+	port := uint16(portUint64)
+
+	var cancel func()
+	switch networkType {
+	case "tcp":
+		cancel = openTCPPorts[port]
+	case "udp":
+		cancel = openUDPPorts[port]
+	}
 
-		delete(openTCPPorts, port)
-	default:
-		zap.L().Info("")
-		zap.L().Info("Cli commands list:")
-		zap.L().Info("connect [ID_HERE]")
-		zap.L().Info("disconnect [ID_HERE]")
-		zap.L().Info("open [TCP_OR_UDP_HERE] [PORT_NUMBER_HERE]")
-		zap.L().Info("close [UDP_OR_UDP_HERE] [PORT_NUMBER_HERE]")
-		zap.L().Info("")
+	if cancel == nil {
+		zap.L().Error("Specified port is not opened")
+		return
 	}
+
+	zap.L().Info("Closing " + networkType + ":" + portStr)
+
+	cancel()
+
+	delete(openTCPPorts, port)
 }