a lot of improvement

1 files changed, 13 insertions, 7 deletions

diff --git a/util/sys/sys_windows.go b/util/sys/sys_windows.go
index ef9cdf54..34740ea6 100644
--- a/util/sys/sys_windows.go
+++ b/util/sys/sys_windows.go
@@ -4,21 +4,27 @@
 package sys
 
 import (
+	"errors"
+
 	"github.com/shirou/gopsutil/v3/net"
 )
 
-func GetTCPCount() (int, error) {
-	stats, err := net.Connections("tcp")
+func GetConnectionCount(proto string) (int, error) {
+	if proto != "tcp" && proto != "udp" {
+		return 0, errors.New("invalid protocol")
+	}
+
+	stats, err := net.Connections(proto)
 	if err != nil {
 		return 0, err
 	}
 	return len(stats), nil
 }
 
+func GetTCPCount() (int, error) {
+	return GetConnectionCount("tcp")
+}
+
 func GetUDPCount() (int, error) {
-	stats, err := net.Connections("udp")
-	if err != nil {
-		return 0, err
-	}
-	return len(stats), nil
+	return GetConnectionCount("udp")
 }