diff --git a/go.mod b/go.mod index ade92a0..9caf6d9 100644 --- a/go.mod +++ b/go.mod @@ -4,7 +4,9 @@ go 1.13 require ( github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d // indirect + github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5 github.com/go-ole/go-ole v1.2.4 // indirect github.com/kardianos/service v1.0.0 github.com/shirou/gopsutil v2.20.2+incompatible + github.com/stretchr/testify v1.6.1 // indirect ) diff --git a/go.sum b/go.sum index e9013f8..d58a102 100644 --- a/go.sum +++ b/go.sum @@ -1,10 +1,23 @@ github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d h1:G0m3OIz70MZUWq3EgK3CesDbo8upS2Vm9/P3FtgI+Jk= github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d/go.mod h1:3eOhrUMpNV+6aFIbp5/iudMxNCF27Vw2OZgy4xEx0Fg= +github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5 h1:BjkPE3785EwPhhyuFkbINB+2a1xATwk8SNDWnJiD41g= +github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5/go.mod h1:jtAfVaU/2cu1+wdSRPWE2c1N2qeAA3K4RH9pYgqwets= +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ole/go-ole v1.2.4 h1:nNBDSCOigTSiarFpYE9J/KtEA1IOW4CNeqT9TQDqCxI= github.com/go-ole/go-ole v1.2.4/go.mod h1:XCwSNxSkXRo4vlyPy93sltvi/qJq0jqQhjqQNIwKuxM= github.com/kardianos/service v1.0.0 h1:HgQS3mFfOlyntWX8Oke98JcJLqt1DBcHR4kxShpYef0= github.com/kardianos/service v1.0.0/go.mod h1:8CzDhVuCuugtsHyZoTvsOBuvonN/UDBvl0kH+BUxvbo= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/shirou/gopsutil v2.20.2+incompatible h1:ucK79BhBpgqQxPASyS2cu9HX8cfDVljBN1WWFvbNvgY= github.com/shirou/gopsutil v2.20.2+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= golang.org/x/sys v0.0.0-20190204203706-41f3e6584952 h1:FDfvYgoVsA7TTZSbgiqjAbfPbK47CNHdWl3h/PJtii0= golang.org/x/sys v0.0.0-20190204203706-41f3e6584952/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/src/config.go b/src/config.go index c9b8cac..ac70309 100644 --- a/src/config.go +++ b/src/config.go @@ -5,6 +5,7 @@ import ( "io/ioutil" "log" "os" + "path/filepath" "strconv" "strings" "time" @@ -67,7 +68,7 @@ type XMLConfig struct { } func readConfig() { - xmlFile, err := os.Open("C:/ProgramData/LoadBalancer.org/LoadBalancer/config.xml") + xmlFile, err := os.Open(filepath.Join(configDir, "config.xml")) if err != nil { panic(err) } @@ -84,7 +85,7 @@ func readConfig() { } func InitConfig() { - f, err := os.OpenFile("C:/ProgramData/LoadBalancer.org/LoadBalancer/lbfbalogfile", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) + f, err := os.OpenFile(filepath.Join(configDir, "lbfbalogfile"), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) if err != nil { log.Fatalf("error opening file: %v", err) } diff --git a/src/globals_linux.go b/src/globals_linux.go new file mode 100644 index 0000000..3b8abd5 --- /dev/null +++ b/src/globals_linux.go @@ -0,0 +1,7 @@ +package main + +import "path/filepath" + +var ( + configDir string = filepath.Join("/etc", "LoadBalancer.org", "LoadBalancer") +) diff --git a/src/globals_windows.go b/src/globals_windows.go new file mode 100644 index 0000000..c20234f --- /dev/null +++ b/src/globals_windows.go @@ -0,0 +1,10 @@ +package main + +import ( + "os" + "path/filepath" +) + +var ( + configDir string = filepath.Join(os.Getenv("ProgramData"), "LoadBalancer.org", "LoadBalancer") +) diff --git a/src/handler.go b/src/handler.go index 0aca120..0f8ad00 100644 --- a/src/handler.go +++ b/src/handler.go @@ -2,11 +2,13 @@ package main import ( "fmt" - "github.com/shirou/gopsutil/cpu" - "github.com/shirou/gopsutil/mem" "math" "net" - "strings" + "strconv" + + "github.com/cakturk/go-netstat/netstat" + "github.com/shirou/gopsutil/cpu" + "github.com/shirou/gopsutil/mem" ) const ( @@ -27,7 +29,6 @@ const ( func handleClient(conn net.Conn) { defer conn.Close() conn.Write(GetResponseForMode()) - conn.Close() } func GetResponseForMode() (response []byte) { @@ -63,8 +64,7 @@ func GetResponseForMode() (response []byte) { // If any resource is important and utilized 100% then everything else is not important if averageCpuLoad > cpuThresholdValue && cpuThresholdValue > 0 || (usedRam > ramThresholdValue && ramThresholdValue > 0) { - response = []byte("0%\n") - return + return []byte("0%\n") } utilization = utilization + averageCpuLoad*cpuImportance @@ -84,14 +84,13 @@ func GetResponseForMode() (response []byte) { sessionOccupied := GetSessionUtilized(tcpService.IPAddress.Value, tcpService.Port.Value, tcpService.MaxConnections.ToInt()) // Calculate utilization - utilization = utilization + sessionOccupied * tcpService.ImportanceFactor.ToFloat() + utilization = utilization + sessionOccupied*tcpService.ImportanceFactor.ToFloat() // increase our divider divider++ if sessionOccupied > 99 && tcpService.ImportanceFactor.ToFloat() == 1 { - response = []byte("0%\n") - return + return []byte("0%\n") } } } @@ -126,6 +125,7 @@ func GetResponseForMode() (response []byte) { default: response = []byte("error\n") } + return } @@ -138,13 +138,22 @@ func GetSessionUtilized(IPAddress, servicePort string, maxNumberOfSessionsPerSer } func getNumberOfLocalEstablishedConnections(ipAddress string, port string) int { - if ipAddress == "*" { - ipAddress = "" + p, err := strconv.Atoi(port) + if err != nil { + return 0 } - result := runcmd("netstat -nt | findstr " + ipAddress + ":" + port + " | findstr ESTABLISHED ") - count := len(strings.Split(result, "\n")) - if count == 0 { - return count + + // get slice of sockets based on match function + tabs, err := netstat.TCPSocks(func(s *netstat.SockTabEntry) bool { + if ipAddress == "*" { + return s.State == netstat.Established && s.LocalAddr.Port == uint16(p) + } + + return s.State == netstat.Established && s.LocalAddr.IP.String() == ipAddress && s.LocalAddr.Port == uint16(p) + }) + if err != nil { + return 0 } - return count - 1 + + return len(tabs) } diff --git a/src/run_command.go b/src/run_command.go deleted file mode 100644 index 2116f8c..0000000 --- a/src/run_command.go +++ /dev/null @@ -1,71 +0,0 @@ -package main - -import ( - "bytes" - "errors" - "log" - "os/exec" - "runtime" - "time" -) - -func runcmd(command string) (res string) { - var shell, flag string - if runtime.GOOS == "windows" { - shell = "cmd" - flag = "/c" - } else { - shell = "/bin/sh" - flag = "-c" - } - res, err := run(10, shell, flag, command) - if err != nil { - log.Println(err) - return - } - return -} - -func run(timeout int, command string, args ...string) (res string, err error) { - // instantiate new command - cmd := exec.Command(command, args...) - - // get pipe to standard output - stdout, err := cmd.StdoutPipe() - if err != nil { - return res, errors.New("cmd.StdoutPipe() error: " + err.Error()) - } - - // start process via command - if err = cmd.Start(); err != nil { - return res, errors.New("cmd.Start() error: " + err.Error()) - } - - // setup a buffer to capture standard output - var buf bytes.Buffer - - // create a channel to capture any errors from wait - done := make(chan error) - go func() { - if _, err := buf.ReadFrom(stdout); err != nil { - panic("buf.Read(stdout) error: " + err.Error()) - } - done <- cmd.Wait() - }() - - // block on select, and switch based on actions received - select { - case <-time.After(time.Duration(timeout) * time.Second): - if err := cmd.Process.Kill(); err != nil { - return res, errors.New("failed to kill: " + err.Error()) - } - return "", errors.New("command timed out") - case err = <-done: - if err != nil { - close(done) - return res, errors.New("process done, with error: " + err.Error()) - } - return buf.String(), nil - } - return "", nil -}