Merge pull request #86 from shizunge/prometheus

Multiple ports and clean metrics

.gitignore

@@ -0,0 +1 @@
+endlessh-go

README.md

@@ -61,14 +61,20 @@ Usage of ./endlessh-go
         when logging hits line file:N, emit a stack trace
   -log_dir string
         If non-empty, write log files in this directory
+  -log_link string
+        If non-empty, add symbolic links in this directory to the log files
+  -logbuflevel int
+        Buffer log messages logged at this level or lower (-1 means don't buffer; 0 means buffer INFO only; ...). Has limited applicability on non-prod platforms.
   -logtostderr
         log to standard error instead of files
   -max_clients int
         Maximum number of clients (default 4096)
   -max_mind_db string
         Path to the MaxMind DB file.
-  -port string
+  -port value
-        SSH listening port (default "2222")
+        SSH listening port. You may provide multiple -port flags to listen to multiple ports. (default "2222")
+  -prometheus_clean_unseen_seconds int
+        Remove series if the IP is not seen for the given time. Set to 0 to disable. (default 0)
   -prometheus_entry string
         Entry point for prometheus (default "metrics")
   -prometheus_host string
@@ -76,7 +82,7 @@ Usage of ./endlessh-go
   -prometheus_port string
         The port for prometheus (default "2112")
   -stderrthreshold value
-        logs at or above this threshold go to stderr
+        logs at or above this threshold go to stderr (default 2)
   -v value
         log level for V logs
   -vmodule value
@@ -93,8 +99,8 @@ Endlessh-go exports the following Prometheus metrics.
 | endlessh_client_closed_count_total   | count | Total number of clients that stopped connecting to this host. |
 | endlessh_sent_bytes_total            | count | Total bytes sent to clients that tried to connect to this host. |
 | endlessh_trapped_time_seconds_total  | count | Total seconds clients spent on endlessh. |
-| endlessh_client_open_count           | count | Number of connections of clients. <br> Labels: <br> <ul><li> `ip`: IP of the client </li> <li>  `country`: Country of the IP </li> <li>  `location`: Country, Region, and City </li> <li>  `geohash`: Geohash of the location </li></ul> |
+| endlessh_client_open_count           | count | Number of connections of clients. <br> Labels: <br> <ul><li> `ip`: Remote IP of the client </li> <li> `local_port`: Local port the program listens to </li> <li> `country`: Country of the IP </li> <li> `location`: Country, Region, and City </li> <li> `geohash`: Geohash of the location </li></ul> |
-| endlessh_client_trapped_time_seconds | count | Seconds a client spends on endlessh. <br> Labels: <br> <ul><li>  `ip`: IP of the client </li></ul> |
+| endlessh_client_trapped_time_seconds | count | Seconds a client spends on endlessh. <br> Labels: <br> <ul><li> `ip`: Remote IP of the client </li> <li> `local_port`: Local port the program listens to </li></ul> |
 
 The metrics is off by default, you can turn it via the CLI argument `-enable_prometheus`.
 

client.go

@@ -1,111 +0,0 @@
-// Copyright (C) 2021 Shizun Ge
-//
-// This program is free software: you can redistribute it and/or modify
-// it under the terms of the GNU General Public License as published by
-// the Free Software Foundation, either version 3 of the License, or
-// (at your option) any later version.
-//
-// This program is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-// GNU General Public License for more details.
-//
-// You should have received a copy of the GNU General Public License
-// along with this program.  If not, see <https://www.gnu.org/licenses/>.
-//
-
-package main
-
-import (
-	"math/rand"
-	"net"
-	"sync/atomic"
-	"time"
-
-	"github.com/golang/glog"
-	"github.com/prometheus/client_golang/prometheus"
-)
-
-var letterBytes = []byte(" abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890!@#$%^&*()-=_+[]{}|;:',./<>?")
-
-func randStringBytes(n int64) []byte {
-	b := make([]byte, n+1)
-	for i := range b {
-		b[i] = letterBytes[rand.Intn(len(letterBytes))]
-	}
-	b[n] = '\n'
-	return b
-}
-
-type client struct {
-	conn              net.Conn
-	last              time.Time
-	next              time.Time
-	start             time.Time
-	interval          time.Duration
-	geoipSupplier     string
-	geohash           string
-	country           string
-	location          string
-	bytesSent         int
-	prometheusEnabled bool
-}
-
-func NewClient(conn net.Conn, interval time.Duration, maxClient int64, geoipSupplier string, prometheusEnabled bool) *client {
-	addr := conn.RemoteAddr().(*net.TCPAddr)
-	atomic.AddInt64(&numCurrentClients, 1)
-	atomic.AddInt64(&numTotalClients, 1)
-	geohash, country, location, err := geohashAndLocation(addr.IP.String(), geoipSupplier)
-	if err != nil {
-		glog.Warningf("Failed to obatin the geohash of %v: %v.", addr.IP, err)
-	}
-	if prometheusEnabled {
-		clientIP.With(prometheus.Labels{
-			"ip":       addr.IP.String(),
-			"geohash":  geohash,
-			"country":  country,
-			"location": location}).Inc()
-	}
-	glog.V(1).Infof("ACCEPT host=%v port=%v n=%v/%v\n", addr.IP, addr.Port, numCurrentClients, maxClient)
-	return &client{
-		conn:              conn,
-		last:              time.Now(),
-		next:              time.Now().Add(interval),
-		start:             time.Now(),
-		interval:          interval,
-		geohash:           geohash,
-		country:           country,
-		location:          location,
-		bytesSent:         0,
-		prometheusEnabled: prometheusEnabled,
-	}
-}
-
-func (c *client) Send(bannerMaxLength int64) error {
-	defer func(c *client) {
-		addr := c.conn.RemoteAddr().(*net.TCPAddr)
-		millisecondsSpent := time.Now().Sub(c.last).Milliseconds()
-		c.last = time.Now()
-		c.next = time.Now().Add(c.interval)
-		atomic.AddInt64(&numTotalMilliseconds, millisecondsSpent)
-		if c.prometheusEnabled {
-			clientSeconds.With(prometheus.Labels{"ip": addr.IP.String()}).Add(float64(millisecondsSpent) / 1000)
-		}
-	}(c)
-	length := rand.Int63n(bannerMaxLength)
-	bytesSent, err := c.conn.Write(randStringBytes(length))
-	if err != nil {
-		return err
-	}
-	c.bytesSent += bytesSent
-	atomic.AddInt64(&numTotalBytes, int64(bytesSent))
-	return nil
-}
-
-func (c *client) Close() {
-	addr := c.conn.RemoteAddr().(*net.TCPAddr)
-	atomic.AddInt64(&numCurrentClients, -1)
-	atomic.AddInt64(&numTotalClientsClosed, 1)
-	glog.V(1).Infof("CLOSE host=%v port=%v time=%v bytes=%v\n", addr.IP, addr.Port, time.Now().Sub(c.start).Seconds(), c.bytesSent)
-	c.conn.Close()
-}

client/client.go

@@ -0,0 +1,102 @@
+// Copyright (C) 2021-2024 Shizun Ge
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+//
+
+package client
+
+import (
+	"math/rand"
+	"net"
+	"strconv"
+	"sync/atomic"
+	"time"
+
+	"github.com/golang/glog"
+)
+
+var (
+	numCurrentClients int64
+	letterBytes       = []byte(" abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890!@#$%^&*()-=_+[]{}|;:',./<>?")
+)
+
+func randStringBytes(n int64) []byte {
+	b := make([]byte, n+1)
+	for i := range b {
+		b[i] = letterBytes[rand.Intn(len(letterBytes))]
+	}
+	b[n] = '\n'
+	return b
+}
+
+type Client struct {
+	conn      net.Conn
+	next      time.Time
+	start     time.Time
+	last      time.Time
+	interval  time.Duration
+	bytesSent int
+}
+
+func NewClient(conn net.Conn, interval time.Duration, maxClients int64) *Client {
+	for numCurrentClients >= maxClients {
+		time.Sleep(interval)
+	}
+	atomic.AddInt64(&numCurrentClients, 1)
+	addr := conn.RemoteAddr().(*net.TCPAddr)
+	glog.V(1).Infof("ACCEPT host=%v port=%v n=%v/%v\n", addr.IP, addr.Port, numCurrentClients, maxClients)
+	return &Client{
+		conn:      conn,
+		next:      time.Now().Add(interval),
+		start:     time.Now(),
+		last:      time.Now(),
+		interval:  interval,
+		bytesSent: 0,
+	}
+}
+
+func (c *Client) RemoteIpAddr() string {
+	return c.conn.RemoteAddr().(*net.TCPAddr).IP.String()
+}
+
+func (c *Client) LocalPort() string {
+	return strconv.Itoa(c.conn.LocalAddr().(*net.TCPAddr).Port)
+}
+
+func (c *Client) Send(bannerMaxLength int64) (int, error) {
+	if time.Now().Before(c.next) {
+		time.Sleep(c.next.Sub(time.Now()))
+	}
+	c.next = time.Now().Add(c.interval)
+	length := rand.Int63n(bannerMaxLength)
+	bytesSent, err := c.conn.Write(randStringBytes(length))
+	if err != nil {
+		return 0, err
+	}
+	c.bytesSent += bytesSent
+	return bytesSent, nil
+}
+
+func (c *Client) MillisecondsSinceLast() int64 {
+	millisecondsSpent := time.Now().Sub(c.last).Milliseconds()
+	c.last = time.Now()
+	return millisecondsSpent
+}
+
+func (c *Client) Close() {
+	addr := c.conn.RemoteAddr().(*net.TCPAddr)
+	glog.V(1).Infof("CLOSE host=%v port=%v time=%v bytes=%v\n", addr.IP, addr.Port, time.Now().Sub(c.start).Seconds(), c.bytesSent)
+	c.conn.Close()
+	atomic.AddInt64(&numCurrentClients, -1)
+}

dashboard/endlessh.json

@@ -76,7 +76,7 @@
       }
     ]
   },
-  "description": "Dashboard for endlessh (Load metrics at the first panel instead of the last panel)",
+  "description": "Dashboard for endlessh (Fix current connections)",
   "editable": false,
   "fiscalYearStartMonth": 0,
   "gnetId": 15156,
@@ -759,7 +759,7 @@
             "uid": "${DS_PROMETHEUS}"
           },
           "exemplar": true,
-          "expr": "sum((endlessh_client_open_count_total{instance=~\"$host\",job=~\"$job\"}) - (endlessh_client_closed_count_total{instance=~\"$host\",job=~\"$job\"} offset $__interval  or endlessh_client_open_count_total{instance=~\"$host\",job=~\"$job\"} * 0))",
+          "expr": "sum((endlessh_client_open_count_total{instance=~\"$host\",job=~\"$job\"}) - (endlessh_client_closed_count_total{instance=~\"$host\",job=~\"$job\"} or endlessh_client_open_count_total{instance=~\"$host\",job=~\"$job\"} * 0))",
           "instant": false,
           "interval": "",
           "legendFormat": "Open Connections",
@@ -1525,6 +1525,6 @@
   "timezone": "",
   "title": "Endlessh",
   "uid": "ATIxYkO7k",
-  "version": 6,
+  "version": 12,
   "weekStart": ""
 }

geoip/country.go

@@ -1,4 +1,4 @@
-// Copyright (C) 2023 Shizun Ge
+// Copyright (C) 2023-2024 Shizun Ge
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU General Public License as published by
@@ -14,16 +14,16 @@
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 //
 
-package coordinates
+package geoip
 
 // Map country's ISO to their capital's latitude and longitude.
 // Country's ISO see https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2
-type Location struct {
+type location struct {
 	Latitude  float64
 	Longitude float64
 }
 
-var Country = map[string]Location{
+var countryToLocation = map[string]location{
 	"AD": {42.5, 1.5},
 	"AE": {24.4511, 54.3969},
 	"AF": {34.5328, 69.1658},

geoip/geoip.go

@@ -1,4 +1,4 @@
-// Copyright (C) 2021-2023 Shizun Ge
+// Copyright (C) 2021-2024 Shizun Ge
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU General Public License as published by
@@ -14,7 +14,7 @@
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 //
 
-package main
+package geoip
 
 import (
 	"encoding/json"
@@ -24,15 +24,14 @@ import (
 	"net/http"
 	"strings"
 
-	"endlessh-go/coordinates"
-
 	"github.com/oschwald/geoip2-golang"
 	"github.com/pierrre/geohash"
 )
 
-var (
+type GeoOption struct {
-	maxMindDbFileName *string
+	GeoipSupplier     string
-)
+	MaxMindDbFileName string
+}
 
 func composeLocation(country string, region string, city string) string {
 	var locations []string
@@ -69,9 +68,9 @@ type ipapi struct {
 	Longitude   float64 `json:"lon"`
 }
 
-func geohashAndLocationFromIpapi(address string) (string, string, string, error) {
+func geohashAndLocationFromIpapi(ipAddr string) (string, string, string, error) {
 	var geo ipapi
-	response, err := http.Get("http://ip-api.com/json/" + address)
+	response, err := http.Get("http://ip-api.com/json/" + ipAddr)
 	if err != nil {
 		return "s000", "Unknown", "Unknown", err
 	}
@@ -88,7 +87,7 @@ func geohashAndLocationFromIpapi(address string) (string, string, string, error)
 	}
 
 	if geo.Status != "success" {
-		return "s000", "Unknown", "Unknown", fmt.Errorf("failed to query %v via ip-api: status: %v, message: %v", address, geo.Status, geo.Message)
+		return "s000", "Unknown", "Unknown", fmt.Errorf("failed to query %v via ip-api: status: %v, message: %v", ipAddr, geo.Status, geo.Message)
 	}
 
 	gh := geohash.EncodeAuto(geo.Latitude, geo.Longitude)
@@ -98,14 +97,14 @@ func geohashAndLocationFromIpapi(address string) (string, string, string, error)
 	return gh, country, location, nil
 }
 
-func geohashAndLocationFromMaxMindDb(address string) (string, string, string, error) {
+func geohashAndLocationFromMaxMindDb(ipAddr, maxMindDbFileName string) (string, string, string, error) {
-	db, err := geoip2.Open(*maxMindDbFileName)
+	db, err := geoip2.Open(maxMindDbFileName)
 	if err != nil {
 		return "s000", "Unknown", "Unknown", err
 	}
 	defer db.Close()
 	// If you are using strings that may be invalid, check that ip is not nil
-	ip := net.ParseIP(address)
+	ip := net.ParseIP(ipAddr)
 	cityRecord, err := db.City(ip)
 	if err != nil {
 		return "s000", "Unknown", "Unknown", err
@@ -117,7 +116,7 @@ func geohashAndLocationFromMaxMindDb(address string) (string, string, string, er
 	iso := cityRecord.Country.IsoCode
 	if latitude == 0 && longitude == 0 {
 		// In case of using Country DB, city is not available.
-		loc, ok := coordinates.Country[iso]
+		loc, ok := countryToLocation[iso]
 		if ok {
 			latitude = loc.Latitude
 			longitude = loc.Longitude
@@ -135,15 +134,15 @@ func geohashAndLocationFromMaxMindDb(address string) (string, string, string, er
 	return gh, country, location, nil
 }
 
-func geohashAndLocation(address string, geoipSupplier string) (string, string, string, error) {
+func GeohashAndLocation(ipAddr string, option GeoOption) (string, string, string, error) {
-	switch geoipSupplier {
+	switch option.GeoipSupplier {
 	case "off":
 		return "s000", "Geohash off", "Geohash off", nil
 	case "ip-api":
-		return geohashAndLocationFromIpapi(address)
+		return geohashAndLocationFromIpapi(ipAddr)
 	case "max-mind-db":
-		return geohashAndLocationFromMaxMindDb(address)
+		return geohashAndLocationFromMaxMindDb(ipAddr, option.MaxMindDbFileName)
 	default:
-		return "s000", "Unknown", "Unknown", fmt.Errorf("unknown geoipSupplier %v.", geoipSupplier)
+		return "s000", "Unknown", "Unknown", fmt.Errorf("unknown geoipSupplier %v.", option.GeoipSupplier)
 	}
 }

go.mod

@@ -1,6 +1,8 @@
 module endlessh-go
 
-go 1.20
+go 1.21.0
+
+toolchain go1.21.4
 
 require (
 	github.com/golang/glog v1.2.0

go.sum

@@ -1,31 +1,43 @@
 github.com/Codefor/geohash v0.0.0-20140723084247-1b41c28e3a9d h1:iG9B49Q218F/XxXNRM7k/vWf7MKmLIS8AcJV9cGN4nA=
+github.com/Codefor/geohash v0.0.0-20140723084247-1b41c28e3a9d/go.mod h1:RVnhzAX71far8Kc3TQeA0k/dcaEKUnTDSOyet/JCmGI=
 github.com/TomiHiltunen/geohash-golang v0.0.0-20150112065804-b3e4e625abfb h1:wumPkzt4zaxO4rHPBrjDK8iZMR41C1qs7njNqlacwQg=
+github.com/TomiHiltunen/geohash-golang v0.0.0-20150112065804-b3e4e625abfb/go.mod h1:QiYsIBRQEO+Z4Rz7GoI+dsHVneZNONvhczuA+llOZNM=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/broady/gogeohash v0.0.0-20120525094510-7b2c40d64042 h1:iEdmkrNMLXbM7ecffOAtZJQOQUTE4iMonxrb5opUgE4=
+github.com/broady/gogeohash v0.0.0-20120525094510-7b2c40d64042/go.mod h1:f1L9YvXvlt9JTa+A17trQjSMM6bV40f+tHjB+Pi+Fqk=
 github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
 github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/fanixk/geohash v0.0.0-20150324002647-c1f9b5fa157a h1:Fyfh/dsHFrC6nkX7H7+nFdTd1wROlX/FxEIWVpKYf1U=
+github.com/fanixk/geohash v0.0.0-20150324002647-c1f9b5fa157a/go.mod h1:UgNw+PTmmGN8rV7RvjvnBMsoTU8ZXXnaT3hYsDTBlgQ=
 github.com/golang/glog v1.2.0 h1:uCdmnmatrKCgMBlM4rMuJZWOkPDqdbZPnrMXDY4gI68=
 github.com/golang/glog v1.2.0/go.mod h1:6AhwSGph0fcJtXVM/PEHPqZlFeoLxhs7/t5UDAwmO+w=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 h1:jWpvCLoY8Z/e3VKvlsiIGKtc+UG6U5vzxaoagmhXfyg=
 github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0/go.mod h1:QUyp042oQthUoa9bqDv0ER0wrtXnBruoNd7aNjkbP+k=
 github.com/mmcloughlin/geohash v0.10.0 h1:9w1HchfDfdeLc+jFEf/04D27KP7E2QmpDu52wPbJWRE=
+github.com/mmcloughlin/geohash v0.10.0/go.mod h1:oNZxQo5yWJh0eMQEP/8hwQuVx9Z9tjwFUqcTB1SmG0c=
 github.com/oschwald/geoip2-golang v1.9.0 h1:uvD3O6fXAXs+usU+UGExshpdP13GAqp4GBrzN7IgKZc=
 github.com/oschwald/geoip2-golang v1.9.0/go.mod h1:BHK6TvDyATVQhKNbQBdrj9eAvuwOMi2zSFXizL3K81Y=
 github.com/oschwald/maxminddb-golang v1.11.0 h1:aSXMqYR/EPNjGE8epgqwDay+P30hCBZIveY0WZbAWh0=
 github.com/oschwald/maxminddb-golang v1.11.0/go.mod h1:YmVI+H0zh3ySFR3w+oz8PCfglAFj3PuCmui13+P9zDg=
 github.com/pierrre/assert v0.3.2 h1:wXdlkVN5FVSLEKl6pGijcCYkldgfjRgyheU3C1/by9Q=
+github.com/pierrre/assert v0.3.2/go.mod h1:zwOn9QE9/+eSgqR/4iCS9K9dUpkqyl0iih2APCI5d8E=
 github.com/pierrre/compare v1.4.2 h1:oabIiWclzAlXG7S/2MYSFDJ/vR34oa/MYrBZh5PNU80=
+github.com/pierrre/compare v1.4.2/go.mod h1:EBDtLjy0XYVBEFP4T3pWljpfTwL7X8DqPt9RIP1+svY=
 github.com/pierrre/geohash v1.1.1 h1:XCkvOyv/uesenMPhkvsCfIiUalBmGdHlFY0bIWTqb+s=
 github.com/pierrre/geohash v1.1.1/go.mod h1:ucAm7cbgGBoVr6cr1t+d3ea5XQ9P5zKHXfS1Qy2iKVY=
 github.com/pierrre/go-libs v0.2.14 h1:wAPoOrslKLnha6ow5EKkxxZpo76kOea57efs71A/ZnQ=
+github.com/pierrre/go-libs v0.2.14/go.mod h1:eA3pQD5LHZmavOpTpUfO8FszduBNHoFXDWrevDR6Dy8=
 github.com/pierrre/pretty v0.0.10 h1:Cb5som+1EpU+x7UA5AMy9I8AY2XkzMBywkLEAdo1JDg=
+github.com/pierrre/pretty v0.0.10/go.mod h1:F+Z4XV4T5GIvbr/swCAkuQ2ng81qMaQT9CfI8rKOLdY=
 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/prometheus/client_golang v1.18.0 h1:HzFfmkOzH5Q8L8G+kSJKUx5dtG87sewO+FoDDqP5Tbk=
 github.com/prometheus/client_golang v1.18.0/go.mod h1:T+GXkCk5wSJyOqMIzVgvvjFDlkOQntgjkJWKrN5txjA=
 github.com/prometheus/client_model v0.5.0 h1:VQw1hfvPvk3Uv6Qf29VrPF32JB6rtbgI6cYPYQjL0Qw=
@@ -35,7 +47,9 @@ github.com/prometheus/common v0.45.0/go.mod h1:YJmSTw9BoKxJplESWWxlbyttQR4uaEcGy
 github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k6Bo=
 github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/the42/cartconvert v1.0.0 h1:g8kt6ic2GEhdcZ61ZP9GsWwhosVo5nCnH1n2/oAQXUU=
+github.com/the42/cartconvert v1.0.0/go.mod h1:fWO/msnJVhHqN1yX6OBoxSyfj7TEj1hHiL8bJSQsK30=
 golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc=
 golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
@@ -43,3 +57,4 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
 google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -1,4 +1,4 @@
-// Copyright (C) 2021-2023 Shizun Ge
+// Copyright (C) 2021-2024 Shizun Ge
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU General Public License as published by
@@ -17,108 +17,113 @@
 package main
 
 import (
+	"endlessh-go/client"
+	"endlessh-go/geoip"
+	"endlessh-go/metrics"
 	"flag"
 	"fmt"
-	"math/rand"
 	"net"
-	"net/http"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/golang/glog"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
-var (
+func startSending(maxClients int64, bannerMaxLength int64, records chan<- metrics.RecordEntry) chan *client.Client {
-	numCurrentClients     int64
+	clients := make(chan *client.Client, maxClients)
-	numTotalClients       int64
-	numTotalClientsClosed int64
-	numTotalBytes         int64
-	numTotalMilliseconds  int64
-	totalClients          prometheus.CounterFunc
-	totalClientsClosed    prometheus.CounterFunc
-	totalBytes            prometheus.CounterFunc
-	totalSeconds          prometheus.CounterFunc
-	clientIP              *prometheus.CounterVec
-	clientSeconds         *prometheus.CounterVec
-)
-
-func initPrometheus(prometheusHost, prometheusPort, prometheusEntry string) {
-	totalClients = prometheus.NewCounterFunc(
-		prometheus.CounterOpts{
-			Name: "endlessh_client_open_count_total",
-			Help: "Total number of clients that tried to connect to this host.",
-		}, func() float64 {
-			return float64(numTotalClients)
-		},
-	)
-	totalClientsClosed = prometheus.NewCounterFunc(
-		prometheus.CounterOpts{
-			Name: "endlessh_client_closed_count_total",
-			Help: "Total number of clients that stopped connecting to this host.",
-		}, func() float64 {
-			return float64(numTotalClientsClosed)
-		},
-	)
-	totalBytes = prometheus.NewCounterFunc(
-		prometheus.CounterOpts{
-			Name: "endlessh_sent_bytes_total",
-			Help: "Total bytes sent to clients that tried to connect to this host.",
-		}, func() float64 {
-			return float64(numTotalBytes)
-		},
-	)
-	totalSeconds = prometheus.NewCounterFunc(
-		prometheus.CounterOpts{
-			Name: "endlessh_trapped_time_seconds_total",
-			Help: "Total seconds clients spent on endlessh.",
-		}, func() float64 {
-			return float64(numTotalMilliseconds) / 1000
-		},
-	)
-	clientIP = prometheus.NewCounterVec(
-		prometheus.CounterOpts{
-			Name: "endlessh_client_open_count",
-			Help: "Number of connections of clients.",
-		},
-		[]string{"ip", "geohash", "country", "location"},
-	)
-	clientSeconds = prometheus.NewCounterVec(
-		prometheus.CounterOpts{
-			Name: "endlessh_client_trapped_time_seconds",
-			Help: "Seconds a client spends on endlessh.",
-		},
-		[]string{"ip"},
-	)
-	promReg := prometheus.NewRegistry()
-	promReg.MustRegister(totalClients)
-	promReg.MustRegister(totalClientsClosed)
-	promReg.MustRegister(totalBytes)
-	promReg.MustRegister(totalSeconds)
-	promReg.MustRegister(clientIP)
-	promReg.MustRegister(clientSeconds)
-	handler := promhttp.HandlerFor(promReg, promhttp.HandlerOpts{EnableOpenMetrics: true})
-	http.Handle("/"+prometheusEntry, handler)
 	go func() {
-		glog.Infof("Starting Prometheus on %v:%v, entry point is /%v", prometheusHost, prometheusPort, prometheusEntry)
+		for {
-		http.ListenAndServe(prometheusHost+":"+prometheusPort, nil)
+			c, more := <-clients
+			if !more {
+				return
+			}
+			go func() {
+				bytesSent, err := c.Send(bannerMaxLength)
+				remoteIpAddr := c.RemoteIpAddr()
+				localPort := c.LocalPort()
+				if err != nil {
+					c.Close()
+					records <- metrics.RecordEntry{
+						RecordType: metrics.RecordEntryTypeStop,
+						IpAddr:     remoteIpAddr,
+						LocalPort:  localPort,
+					}
+					return
+				}
+				millisecondsSpent := c.MillisecondsSinceLast()
+				clients <- c
+				records <- metrics.RecordEntry{
+					RecordType:        metrics.RecordEntryTypeSend,
+					IpAddr:            remoteIpAddr,
+					LocalPort:         localPort,
+					BytesSent:         bytesSent,
+					MillisecondsSpent: millisecondsSpent,
+				}
+			}()
+		}
+	}()
+	return clients
+}
+
+func startAccepting(maxClients int64, connType, connHost, connPort string, interval time.Duration, clients chan<- *client.Client, records chan<- metrics.RecordEntry) {
+	go func() {
+		l, err := net.Listen(connType, connHost+":"+connPort)
+		if err != nil {
+			glog.Errorf("Error listening: %v", err)
+			os.Exit(1)
+		}
+		// Close the listener when the application closes.
+		defer l.Close()
+		glog.Infof("Listening on %v:%v", connHost, connPort)
+		for {
+			// Listen for an incoming connection.
+			conn, err := l.Accept()
+			if err != nil {
+				glog.Errorf("Error accepting connection from port %v: %v", connPort, err)
+				os.Exit(1)
+			}
+			c := client.NewClient(conn, interval, maxClients)
+			remoteIpAddr := c.RemoteIpAddr()
+			records <- metrics.RecordEntry{
+				RecordType: metrics.RecordEntryTypeStart,
+				IpAddr:     remoteIpAddr,
+				LocalPort:  connPort,
+			}
+			clients <- c
+		}
 	}()
 }
 
+type arrayStrings []string
+
+func (a *arrayStrings) String() string {
+	return strings.Join(*a, ", ")
+}
+
+func (a *arrayStrings) Set(value string) error {
+	*a = append(*a, value)
+	return nil
+}
+
+const defaultPort = "2222"
+
+var connPorts arrayStrings
+
 func main() {
 	intervalMs := flag.Int("interval_ms", 1000, "Message millisecond delay")
 	bannerMaxLength := flag.Int64("line_length", 32, "Maximum banner line length")
 	maxClients := flag.Int64("max_clients", 4096, "Maximum number of clients")
 	connType := flag.String("conn_type", "tcp", "Connection type. Possible values are tcp, tcp4, tcp6")
 	connHost := flag.String("host", "0.0.0.0", "SSH listening address")
-	connPort := flag.String("port", "2222", "SSH listening port")
+	flag.Var(&connPorts, "port", fmt.Sprintf("SSH listening port. You may provide multiple -port flags to listen to multiple ports. (default %q)", defaultPort))
 	prometheusEnabled := flag.Bool("enable_prometheus", false, "Enable prometheus")
 	prometheusHost := flag.String("prometheus_host", "0.0.0.0", "The address for prometheus")
 	prometheusPort := flag.String("prometheus_port", "2112", "The port for prometheus")
 	prometheusEntry := flag.String("prometheus_entry", "metrics", "Entry point for prometheus")
+	prometheusCleanUnseenSeconds := flag.Int("prometheus_clean_unseen_seconds", 0, "Remove series if the IP is not seen for the given time. Set to 0 to disable. (default 0)")
 	geoipSupplier := flag.String("geoip_supplier", "off", "Supplier to obtain Geohash of IPs. Possible values are \"off\", \"ip-api\", \"max-mind-db\"")
-	maxMindDbFileName = flag.String("max_mind_db", "", "Path to the MaxMind DB file.")
+	maxMindDbFileName := flag.String("max_mind_db", "", "Path to the MaxMind DB file.")
 
 	flag.Usage = func() {
 		fmt.Fprintf(flag.CommandLine.Output(), "Usage of %v \n", os.Args[0])
@@ -126,60 +131,39 @@ func main() {
 	}
 	flag.Parse()
 
-	if *connType == "tcp6" && *prometheusHost == "0.0.0.0" {
-		*prometheusHost = "[::]"
-	}
 	if *prometheusEnabled {
-		initPrometheus(*prometheusHost, *prometheusPort, *prometheusEntry)
+		if *connType == "tcp6" && *prometheusHost == "0.0.0.0" {
+			*prometheusHost = "[::]"
+		}
+		metrics.InitPrometheus(*prometheusHost, *prometheusPort, *prometheusEntry)
 	}
 
-	rand.Seed(time.Now().UnixNano())
+	records := metrics.StartRecording(*maxClients, *prometheusEnabled, *prometheusCleanUnseenSeconds,
+		geoip.GeoOption{
+			GeoipSupplier:     *geoipSupplier,
+			MaxMindDbFileName: *maxMindDbFileName,
+		})
+	clients := startSending(*maxClients, *bannerMaxLength, records)
+
 	interval := time.Duration(*intervalMs) * time.Millisecond
 	// Listen for incoming connections.
 	if *connType == "tcp6" && *connHost == "0.0.0.0" {
 		*connHost = "[::]"
 	}
-	l, err := net.Listen(*connType, *connHost+":"+*connPort)
+	if len(connPorts) == 0 {
-	if err != nil {
+		connPorts = append(connPorts, defaultPort)
-		glog.Errorf("Error listening: %v", err)
-		os.Exit(1)
 	}
-	// Close the listener when the application closes.
+	for _, connPort := range connPorts {
-	defer l.Close()
+		startAccepting(*maxClients, *connType, *connHost, connPort, interval, clients, records)
-	glog.Infof("Listening on %v:%v", *connHost, *connPort)
+	}
-
+	for {
-	clients := make(chan *client, *maxClients)
+		if *prometheusCleanUnseenSeconds <= 0 {
-	go func() {
+			time.Sleep(time.Duration(1<<63 - 1))
-		for {
+		} else {
-			c, more := <-clients
+			time.Sleep(time.Second * time.Duration(60))
-			if !more {
+			records <- metrics.RecordEntry{
-				return
+				RecordType: metrics.RecordEntryTypeClean,
 			}
-			if time.Now().Before(c.next) {
-				time.Sleep(c.next.Sub(time.Now()))
-			}
-			err := c.Send(*bannerMaxLength)
-			if err != nil {
-				c.Close()
-				continue
-			}
-			go func() { clients <- c }()
-		}
-	}()
-	listener := func() {
-		for {
-			// Listen for an incoming connection.
-			conn, err := l.Accept()
-			if err != nil {
-				glog.Errorf("Error accepting: %v", err)
-				os.Exit(1)
-			}
-			// Handle connections in a new goroutine.
-			for numCurrentClients >= *maxClients {
-				time.Sleep(interval)
-			}
-			clients <- NewClient(conn, interval, *maxClients, *geoipSupplier, *prometheusEnabled)
 		}
 	}
-	listener()
 }

metrics/metrics.go

@@ -0,0 +1,162 @@
+// Copyright (C) 2024 Shizun Ge
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+//
+
+package metrics
+
+import (
+	"endlessh-go/geoip"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/golang/glog"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+var (
+	pq                 *UpdatablePriorityQueue
+	totalClients       *prometheus.CounterVec
+	totalClientsClosed *prometheus.CounterVec
+	totalBytes         *prometheus.CounterVec
+	totalSeconds       *prometheus.CounterVec
+	clientIP           *prometheus.CounterVec
+	clientSeconds      *prometheus.CounterVec
+)
+
+func InitPrometheus(prometheusHost, prometheusPort, prometheusEntry string) {
+	pq = NewUpdatablePriorityQueue()
+	totalClients = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "endlessh_client_open_count_total",
+			Help: "Total number of clients that tried to connect to this host.",
+		}, []string{"local_port"},
+	)
+	totalClientsClosed = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "endlessh_client_closed_count_total",
+			Help: "Total number of clients that stopped connecting to this host.",
+		}, []string{"local_port"},
+	)
+	totalBytes = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "endlessh_sent_bytes_total",
+			Help: "Total bytes sent to clients that tried to connect to this host.",
+		}, []string{"local_port"},
+	)
+	totalSeconds = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "endlessh_trapped_time_seconds_total",
+			Help: "Total seconds clients spent on endlessh.",
+		}, []string{"local_port"},
+	)
+	clientIP = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "endlessh_client_open_count",
+			Help: "Number of connections of clients.",
+		},
+		[]string{"ip", "local_port", "geohash", "country", "location"},
+	)
+	clientSeconds = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "endlessh_client_trapped_time_seconds",
+			Help: "Seconds a client spends on endlessh.",
+		},
+		[]string{"ip", "local_port"},
+	)
+	promReg := prometheus.NewRegistry()
+	promReg.MustRegister(totalClients)
+	promReg.MustRegister(totalClientsClosed)
+	promReg.MustRegister(totalBytes)
+	promReg.MustRegister(totalSeconds)
+	promReg.MustRegister(clientIP)
+	promReg.MustRegister(clientSeconds)
+	handler := promhttp.HandlerFor(promReg, promhttp.HandlerOpts{EnableOpenMetrics: true})
+	http.Handle("/"+prometheusEntry, handler)
+	go func() {
+		glog.Infof("Starting Prometheus on %v:%v, entry point is /%v", prometheusHost, prometheusPort, prometheusEntry)
+		if err := http.ListenAndServe(prometheusHost+":"+prometheusPort, nil); err != nil {
+			glog.Errorf("Error starting Prometheus at port %v:%v: %v", prometheusHost, prometheusPort, err)
+			os.Exit(1)
+		}
+	}()
+}
+
+const (
+	RecordEntryTypeStart = iota
+	RecordEntryTypeSend  = iota
+	RecordEntryTypeStop  = iota
+	RecordEntryTypeClean = iota
+)
+
+type RecordEntry struct {
+	RecordType        int
+	IpAddr            string
+	LocalPort         string
+	BytesSent         int
+	MillisecondsSpent int64
+}
+
+func StartRecording(maxClients int64, prometheusEnabled bool, prometheusCleanUnseenSeconds int, geoOption geoip.GeoOption) chan RecordEntry {
+	records := make(chan RecordEntry, maxClients)
+	go func() {
+		for {
+			r, more := <-records
+			if !more {
+				return
+			}
+			if !prometheusEnabled {
+				continue
+			}
+			switch r.RecordType {
+			case RecordEntryTypeStart:
+				geohash, country, location, err := geoip.GeohashAndLocation(r.IpAddr, geoOption)
+				if err != nil {
+					glog.Warningf("Failed to obatin the geohash of %v: %v.", r.IpAddr, err)
+				}
+				clientIP.With(prometheus.Labels{
+					"ip":         r.IpAddr,
+					"local_port": r.LocalPort,
+					"geohash":    geohash,
+					"country":    country,
+					"location":   location}).Inc()
+				totalClients.With(prometheus.Labels{"local_port": r.LocalPort}).Inc()
+				pq.Update(r.IpAddr, time.Now())
+			case RecordEntryTypeSend:
+				secondsSpent := float64(r.MillisecondsSpent) / 1000
+				clientSeconds.With(prometheus.Labels{
+					"ip":         r.IpAddr,
+					"local_port": r.LocalPort}).Add(secondsSpent)
+				totalBytes.With(prometheus.Labels{"local_port": r.LocalPort}).Add(float64(r.BytesSent))
+				totalSeconds.With(prometheus.Labels{"local_port": r.LocalPort}).Add(secondsSpent)
+				pq.Update(r.IpAddr, time.Now())
+			case RecordEntryTypeStop:
+				totalClientsClosed.With(prometheus.Labels{"local_port": r.LocalPort}).Inc()
+				pq.Update(r.IpAddr, time.Now())
+			case RecordEntryTypeClean:
+				top := pq.Peek()
+				deadline := time.Now().Add(-time.Second * time.Duration(prometheusCleanUnseenSeconds))
+				for top != nil && top.Value.Before(deadline) {
+					clientIP.DeletePartialMatch(prometheus.Labels{"ip": top.Key})
+					clientSeconds.DeletePartialMatch(prometheus.Labels{"ip": top.Key})
+					pq.Pop()
+					top = pq.Peek()
+				}
+			}
+		}
+	}()
+	return records
+}

metrics/priority_queue.go

@@ -0,0 +1,94 @@
+package metrics
+
+import (
+	"container/heap"
+	"time"
+)
+
+// Pair represents a key-value pair with a timestamp
+type Pair struct {
+	Key     string
+	Value   time.Time
+	HeapIdx int // Index in the heap for efficient updates
+}
+
+// PriorityQueue is a min-heap implementation for Pairs
+type PriorityQueue []*Pair
+
+// Len returns the length of the priority queue
+func (pq PriorityQueue) Len() int { return len(pq) }
+
+// Less compares two pairs based on their values (timestamps)
+func (pq PriorityQueue) Less(i, j int) bool {
+	return pq[i].Value.Before(pq[j].Value)
+}
+
+// Swap swaps two pairs in the priority queue
+func (pq PriorityQueue) Swap(i, j int) {
+	pq[i], pq[j] = pq[j], pq[i]
+	pq[i].HeapIdx = i
+	pq[j].HeapIdx = j
+}
+
+// Push adds a pair to the priority queue
+func (pq *PriorityQueue) Push(x interface{}) {
+	pair := x.(*Pair)
+	pair.HeapIdx = len(*pq)
+	*pq = append(*pq, pair)
+}
+
+// Pop removes the pair with the minimum value (timestamp) from the priority queue
+func (pq *PriorityQueue) Pop() interface{} {
+	old := *pq
+	n := len(old)
+	pair := old[n-1]
+	pair.HeapIdx = -1 // for safety
+	*pq = old[0 : n-1]
+	return pair
+}
+
+// UpdatablePriorityQueue represents the data structure with the priority queue
+type UpdatablePriorityQueue struct {
+	pq     PriorityQueue
+	keyMap map[string]*Pair
+}
+
+// NewUpdatablePriorityQueue initializes a new UpdatablePriorityQueue
+func NewUpdatablePriorityQueue() *UpdatablePriorityQueue {
+	return &UpdatablePriorityQueue{
+		pq:     make(PriorityQueue, 0),
+		keyMap: make(map[string]*Pair),
+	}
+}
+
+// Update adds or updates a key-value pair in the data structure
+func (ds *UpdatablePriorityQueue) Update(key string, value time.Time) {
+	if pair, ok := ds.keyMap[key]; ok {
+		// Key exists, update the time
+		pair.Value = value
+		heap.Fix(&ds.pq, pair.HeapIdx)
+	} else {
+		// Key does not exist, create a new entry
+		pair := &Pair{Key: key, Value: value}
+		heap.Push(&ds.pq, pair)
+		ds.keyMap[key] = pair
+	}
+}
+
+// Peek returns the entry with the minimal time
+func (ds *UpdatablePriorityQueue) Peek() *Pair {
+	if ds.pq.Len() == 0 {
+		return nil
+	}
+	return ds.pq[0]
+}
+
+// Pop removes the entry with the minimal time
+func (ds *UpdatablePriorityQueue) Pop() *Pair {
+	if ds.pq.Len() == 0 {
+		return nil
+	}
+	pair := heap.Pop(&ds.pq).(*Pair)
+	delete(ds.keyMap, pair.Key)
+	return pair
+}