From c38e1e0cfe9fd95f1cd2ed5f6c6cb10bf820b9b6 Mon Sep 17 00:00:00 2001
From: MHSanaei <ho3ein.sanaei@gmail.com>
Date: Tue, 23 May 2023 02:43:15 +0330
Subject: a lot of improvement

---
 util/sys/sys_linux.go   | 17 +++++++++--------
 util/sys/sys_windows.go | 20 +++++++++++++-------
 2 files changed, 22 insertions(+), 15 deletions(-)

(limited to 'util/sys')

diff --git a/util/sys/sys_linux.go b/util/sys/sys_linux.go
index bcb7046e..097ff7c7 100644
--- a/util/sys/sys_linux.go
+++ b/util/sys/sys_linux.go
@@ -24,8 +24,8 @@ func getLinesNum(filename string) (int, error) {
 
 		var buffPosition int
 		for {
-			i := bytes.IndexByte(buf[buffPosition:], '\n')
-			if i < 0 || n == buffPosition {
+			i := bytes.IndexByte(buf[buffPosition:n], '\n')
+			if i < 0 {
 				break
 			}
 			buffPosition += i + 1
@@ -33,11 +33,12 @@ func getLinesNum(filename string) (int, error) {
 		}
 
 		if err == io.EOF {
-			return sum, nil
+			break
 		} else if err != nil {
-			return sum, err
+			return 0, err
 		}
 	}
+	return sum, nil
 }
 
 func GetTCPCount() (int, error) {
@@ -45,11 +46,11 @@ func GetTCPCount() (int, error) {
 
 	tcp4, err := getLinesNum(fmt.Sprintf("%v/net/tcp", root))
 	if err != nil {
-		return tcp4, err
+		return 0, err
 	}
 	tcp6, err := getLinesNum(fmt.Sprintf("%v/net/tcp6", root))
 	if err != nil {
-		return tcp4 + tcp6, nil
+		return 0, err
 	}
 
 	return tcp4 + tcp6, nil
@@ -60,11 +61,11 @@ func GetUDPCount() (int, error) {
 
 	udp4, err := getLinesNum(fmt.Sprintf("%v/net/udp", root))
 	if err != nil {
-		return udp4, err
+		return 0, err
 	}
 	udp6, err := getLinesNum(fmt.Sprintf("%v/net/udp6", root))
 	if err != nil {
-		return udp4 + udp6, nil
+		return 0, err
 	}
 
 	return udp4 + udp6, nil
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")
 }
-- 
cgit v1.2.3