Refactor codebase

cmd/connect.go

@@ -30,8 +30,7 @@ func init() {
 }
 
 var connectCmd = &cobra.Command{
-	Use: "connect",
-
+	Use:   "connect",
 	Short: "Connect to a vpn server (survey selection appears if hostname is not provided)",
 	Long:  `Connect to a vpn from a list of relay servers`,
 	Args:  cobra.RangeArgs(0, 1),
@@ -39,49 +38,43 @@ var connectCmd = &cobra.Command{
 		vpnServers, err := vpn.GetList(flagProxy, flagSocks5Proxy)
 		if err != nil {
 			log.Fatal().Msg(err.Error())
-			os.Exit(1)
 		}
 
-		serverSelection := []string{}
-		serverSelected := vpn.Server{}
-
-		for _, s := range *vpnServers {
-			serverSelection = append(serverSelection, fmt.Sprintf("%s (%s)", s.HostName, s.CountryLong))
+		// Build server selection options and hostname lookup map
+		serverSelection := make([]string, len(*vpnServers))
+		serverMap := make(map[string]vpn.Server, len(*vpnServers))
+		for i, s := range *vpnServers {
+			serverSelection[i] = fmt.Sprintf("%s (%s)", s.HostName, s.CountryLong)
+			serverMap[s.HostName] = s
 		}
 
 		selection := ""
-		prompt := &survey.Select{
-			Message: "Choose a server:",
-			Options: serverSelection,
-		}
+		var serverSelected vpn.Server
 
 		if !flagRandom {
-
 			if len(args) > 0 {
 				selection = args[0]
 			} else {
+				prompt := &survey.Select{
+					Message: "Choose a server:",
+					Options: serverSelection,
+				}
 				err := survey.AskOne(prompt, &selection, survey.WithPageSize(10))
 				if err != nil {
-					log.Error().Msg("Unable to obtain hostname from survey")
-					os.Exit(1)
+					log.Fatal().Msg("Unable to obtain hostname from survey")
 				}
 			}
 
-			// Server lookup from selection could be more optimized with a hash map
-			for _, s := range *vpnServers {
-				if strings.Contains(selection, s.HostName) {
-					serverSelected = s
-				}
-			}
-
-			if serverSelected.HostName == "" {
+			// Lookup server from selection using map for O(1) lookup
+			hostname := extractHostname(selection)
+			if server, exists := serverMap[hostname]; exists {
+				serverSelected = server
+			} else {
 				log.Fatal().Msgf("Server '%s' was not found", selection)
-				os.Exit(1)
 			}
 		}
 
 		for {
-
 			if flagRandom {
 				// Select a random server
 				serverSelected = (*vpnServers)[rand.Intn(len(*vpnServers))]
@@ -90,23 +83,19 @@ var connectCmd = &cobra.Command{
 			decodedConfig, err := base64.StdEncoding.DecodeString(serverSelected.OpenVpnConfigData)
 			if err != nil {
 				log.Fatal().Msg(err.Error())
-				os.Exit(1)
 			}
 
 			tmpfile, err := os.CreateTemp("", "vpngate-openvpn-config-")
 			if err != nil {
 				log.Fatal().Msg(err.Error())
-				os.Exit(1)
 			}
 
 			if _, err := tmpfile.Write(decodedConfig); err != nil {
 				log.Fatal().Msg(err.Error())
-				os.Exit(1)
 			}
 
 			if err := tmpfile.Close(); err != nil {
 				log.Fatal().Msg(err.Error())
-				os.Exit(1)
 			}
 
 			log.Info().Msgf("Connecting to %s (%s) in %s", serverSelected.HostName, serverSelected.IPAddr, serverSelected.CountryLong)
@@ -114,16 +103,22 @@ var connectCmd = &cobra.Command{
 			err = vpn.Connect(tmpfile.Name())
 
 			if err != nil && !flagReconnect {
-				log.Fatal().Msg(err.Error())
-				os.Exit(1)
-			} else {
-				err = os.Remove(tmpfile.Name())
-				if err !=  nil {
-					log.Fatal().Msg(err.Error())
-					os.Exit(1)
-				}
+				// VPN connection failed and reconnect is disabled
+				_ = os.Remove(tmpfile.Name())
+				log.Fatal().Msg("VPN connection failed")
 			}
 
+			// Always try to clean up temporary file
+			_ = os.Remove(tmpfile.Name())
 		}
 	},
 }
+
+// extractHostname extracts the hostname from the selection string (format: "hostname (country)")
+func extractHostname(selection string) string {
+	parts := strings.Split(selection, " (")
+	if len(parts) > 0 {
+		return parts[0]
+	}
+	return selection
+}