Skip to content

Commit

Permalink
fix crash when prometheus is disable
Browse files Browse the repository at this point in the history
  • Loading branch information
shizunge committed Oct 27, 2021
1 parent de031db commit 31cf50b
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 36 deletions.
10 changes: 8 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
66 changes: 36 additions & 30 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,42 +38,46 @@ 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)
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)
}
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,
}
}

Expand All @@ -84,22 +88,24 @@ 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
}

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()
}
24 changes: 20 additions & 4 deletions main.go
Original file line number Diff line number Diff line change
@@ -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 (
Expand Down Expand Up @@ -95,18 +111,18 @@ 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])
flag.PrintDefaults()
}
flag.Parse()

if *enablePrometheus {
if *prometheusEnabled {
initPrometheus(*connHost, *prometheusPort, *prometheusEntry)
}

Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit 31cf50b

Please sign in to comment.