fix crash when prometheus is disable

README.md

@@ -14,17 +14,23 @@ Clone the repo then build from source:
 
 ```
 go build .
-./endlessh-go
+./endlessh-go &
 ```
 
 Alternatively, you can use the [docker image](https://hub.docker.com/r/shizunge/endlessh-go):
 
 ```
-sudo docker run -d shizunge/endlessh-go
+sudo docker run -d -p 2222:2222 shizunge/endlessh-go
 ```
 
 It listens to port `2222` by default.
 
+Then you can try to connect to the endlessh server. Your SSH client should hang there.
+
+```
+ssh -p 2222 localhost
+```
+
 If you want log like the [C implementation](https://github.com/skeeto/endlessh), you need to set both CLI arguments `-logtostderr` and `-v=1`, then the log will go to the stderr. You can set different log destinations via CLI arguments.
 
 ## Usage

client.go

@@ -38,19 +38,20 @@ func randStringBytes(n int64) []byte {
 }
 
 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
-	bytes_sent    int
+	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) *client {
+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)
@@ -58,22 +59,25 @@ func NewClient(conn net.Conn, interval time.Duration, maxClient int64, geoipSupp
 	if err != nil {
 		glog.Warningf("Failed to obatin the geohash of %v: %v.", addr.IP, err)
 	}
-	clientIP.With(prometheus.Labels{
-		"ip":       addr.IP.String(),
-		"geohash":  geohash,
-		"country":  country,
-		"location": location}).Inc()
+	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,
-		bytes_sent: 0,
+		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,
 	}
 }
 
@@ -84,15 +88,17 @@ func (c *client) Send(bannerMaxLength int64) error {
 		c.last = time.Now()
 		c.next = time.Now().Add(c.interval)
 		atomic.AddInt64(&numTotalMilliseconds, millisecondsSpent)
-		clientSeconds.With(prometheus.Labels{"ip": addr.IP.String()}).Add(float64(millisecondsSpent) / 1000)
+		if c.prometheusEnabled {
+			clientSeconds.With(prometheus.Labels{"ip": addr.IP.String()}).Add(float64(millisecondsSpent) / 1000)
+		}
 	}(c)
 	length := rand.Int63n(bannerMaxLength)
-	bytes_sent, err := c.conn.Write(randStringBytes(length))
+	bytesSent, err := c.conn.Write(randStringBytes(length))
 	if err != nil {
 		return err
 	}
-	c.bytes_sent += bytes_sent
-	atomic.AddInt64(&numTotalBytes, int64(bytes_sent))
+	c.bytesSent += bytesSent
+	atomic.AddInt64(&numTotalBytes, int64(bytesSent))
 	return nil
 }
 
@@ -100,6 +106,6 @@ 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.bytes_sent)
+	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()
 }

main.go

@@ -1,3 +1,19 @@
+// 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 (
@@ -95,10 +111,10 @@ func main() {
 	connType := flag.String("conn_type", "tcp", "Connection type. Possible values are tcp, tcp4, tcp6")
 	connHost := flag.String("host", "0.0.0.0", "Listening address")
 	connPort := flag.String("port", "2222", "Listening port")
-	enablePrometheus := flag.Bool("enable_prometheus", false, "Enable prometheus")
+	prometheusEnabled := flag.Bool("enable_prometheus", false, "Enable prometheus")
 	prometheusPort := flag.String("prometheus_port", "2112", "The port for prometheus")
 	prometheusEntry := flag.String("prometheus_entry", "metrics", "Entry point for prometheus")
-	geoipSupplier := flag.String("geoip_supplier", "ip-api", "Supplier to obtain Geohash of IPs. Possible values are \"ip-api\", \"freegeoip\"")
+	geoipSupplier := flag.String("geoip_supplier", "off", "Supplier to obtain Geohash of IPs. Possible values are \"off\", \"ip-api\", \"freegeoip\"")
 
 	flag.Usage = func() {
 		fmt.Fprintf(flag.CommandLine.Output(), "Usage of %v \n", os.Args[0])
@@ -106,7 +122,7 @@ func main() {
 	}
 	flag.Parse()
 
-	if *enablePrometheus {
+	if *prometheusEnabled {
 		initPrometheus(*connHost, *prometheusPort, *prometheusEntry)
 	}
 
@@ -155,7 +171,7 @@ func main() {
 			for numCurrentClients >= *maxClients {
 				time.Sleep(interval)
 			}
-			clients <- NewClient(conn, interval, *maxClients, *geoipSupplier)
+			clients <- NewClient(conn, interval, *maxClients, *geoipSupplier, *prometheusEnabled)
 		}
 	}
 	listener()