diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d8fe2c3 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +endlessh-go diff --git a/README.md b/README.md index 2f2f89a..a6c82fd 100644 --- a/README.md +++ b/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 - SSH listening port (default "2222") + -port value + 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.
Labels:
| -| endlessh_client_trapped_time_seconds | count | Seconds a client spends on endlessh.
Labels:
| +| endlessh_client_open_count | count | Number of connections of clients.
Labels:
| +| endlessh_client_trapped_time_seconds | count | Seconds a client spends on endlessh.
Labels:
| The metrics is off by default, you can turn it via the CLI argument `-enable_prometheus`. diff --git a/client.go b/client.go deleted file mode 100644 index b81ad99..0000000 --- a/client.go +++ /dev/null @@ -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 . -// - -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() -} diff --git a/client/client.go b/client/client.go new file mode 100644 index 0000000..17f4c86 --- /dev/null +++ b/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 . +// + +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) +} diff --git a/dashboard/endlessh.json b/dashboard/endlessh.json index 4cb94ea..096cd88 100755 --- a/dashboard/endlessh.json +++ b/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": "" } diff --git a/coordinates/country.go b/geoip/country.go similarity index 98% rename from coordinates/country.go rename to geoip/country.go index e09ca2a..ebe6ad0 100644 --- a/coordinates/country.go +++ b/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 . // -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}, diff --git a/geoip.go b/geoip/geoip.go similarity index 80% rename from geoip.go rename to geoip/geoip.go index a42cd83..6701197 100644 --- a/geoip.go +++ b/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 . // -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 ( - maxMindDbFileName *string -) +type GeoOption struct { + 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) { - db, err := geoip2.Open(*maxMindDbFileName) +func geohashAndLocationFromMaxMindDb(ipAddr, maxMindDbFileName string) (string, string, string, error) { + 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) { - switch geoipSupplier { +func GeohashAndLocation(ipAddr string, option GeoOption) (string, string, string, error) { + 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) } } diff --git a/go.mod b/go.mod index 31796f2..08dc9d3 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 614f397..ffad79c 100644 --- a/go.sum +++ b/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= diff --git a/main.go b/main.go index f2d4c92..64b996e 100644 --- a/main.go +++ b/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 ( - numCurrentClients int64 - 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) +func startSending(maxClients int64, bannerMaxLength int64, records chan<- metrics.RecordEntry) chan *client.Client { + clients := make(chan *client.Client, maxClients) go func() { - glog.Infof("Starting Prometheus on %v:%v, entry point is /%v", prometheusHost, prometheusPort, prometheusEntry) - http.ListenAndServe(prometheusHost+":"+prometheusPort, nil) + for { + 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 err != nil { - glog.Errorf("Error listening: %v", err) - os.Exit(1) + if len(connPorts) == 0 { + connPorts = append(connPorts, defaultPort) } - // Close the listener when the application closes. - defer l.Close() - glog.Infof("Listening on %v:%v", *connHost, *connPort) - - clients := make(chan *client, *maxClients) - go func() { - for { - c, more := <-clients - if !more { - return + for _, connPort := range connPorts { + startAccepting(*maxClients, *connType, *connHost, connPort, interval, clients, records) + } + for { + if *prometheusCleanUnseenSeconds <= 0 { + time.Sleep(time.Duration(1<<63 - 1)) + } else { + time.Sleep(time.Second * time.Duration(60)) + records <- metrics.RecordEntry{ + 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() } diff --git a/metrics/metrics.go b/metrics/metrics.go new file mode 100644 index 0000000..fc6773b --- /dev/null +++ b/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 . +// + +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 +} diff --git a/metrics/priority_queue.go b/metrics/priority_queue.go new file mode 100644 index 0000000..a7559fe --- /dev/null +++ b/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 +}