diff --git a/Dockerfile b/Dockerfile index 6450ecf..6fda926 100644 --- a/Dockerfile +++ b/Dockerfile @@ -19,4 +19,4 @@ HEALTHCHECK --interval=30s --timeout=3s --start-period=10s --retries=3 \ CMD ["/endlessh", "-healthcheck"] USER nobody ENTRYPOINT ["/endlessh"] -CMD ["-logtostderr", "-v=1", "-enable_healthcheck"] +CMD ["-logtostderr", "-v=1", "-healthcheck_enable"] diff --git a/README.md b/README.md index ec9907a..abd50a9 100644 --- a/README.md +++ b/README.md @@ -59,11 +59,13 @@ Usage of ./endlessh-go -geoip_supplier string Supplier to obtain Geohash of IPs. Possible values are "off", "ip-api", "max-mind-db" (default "off") -healthcheck - GET healthcheck_host:healthcheck_port/health and exit 1 if status is not ok or timeout is exceeded (for container healthcheck) + Perform healthcheck and exit. GET healthcheck_host:healthcheck_port/health and exit 1 if status is not ok or timeout is exceeded. + -healthcheck_enable + Enable healthcheck -healthcheck_host string - The address for container healthcheck (default "127.0.0.1") + The address for healthcheck. (default "127.0.0.1") -healthcheck_port string - HTTP port for container healthcheck; serves JSON with status and uptime at /health (default "51000") + HTTP port for healthcheck; Serves JSON with status and uptime at /health. (default "51000") -host string SSH listening address (default "0.0.0.0") -interval_ms int @@ -129,7 +131,7 @@ You could also use an offline GeoIP database from [MaxMind](https://www.maxmind. ## Healthcheck -The endlessh-go server exposes an HTTP health endpoint while the server is running. By default it listens on `127.0.0.1:51000` at `/health` and returns a JSON like this: +The endlessh-go server exposes an HTTP health endpoint when the `-healthcheck_enable` flag is set. By default it listens on `127.0.0.1:51000` at `/health` and returns a JSON like this: ```json { diff --git a/main.go b/main.go index 9392536..a1d51a6 100644 --- a/main.go +++ b/main.go @@ -131,25 +131,34 @@ const defaultPort = "2222" var connPorts arrayStrings func main() { + // Core SSH server flags + connHost := flag.String("host", "0.0.0.0", "SSH listening address") + flag.Var(&connPorts, "port", fmt.Sprintf("SSH listening port. You may provide multiple -port flags to listen to multiple ports. (default %q)", defaultPort)) + connType := flag.String("conn_type", "tcp", "Connection type. Possible values are tcp, tcp4, tcp6") 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") - flag.Var(&connPorts, "port", fmt.Sprintf("SSH listening port. You may provide multiple -port flags to listen to multiple ports. (default %q)", defaultPort)) + + // PROXY protocol flags + proxyProtocolEnabled := flag.Bool("proxy_protocol_enabled", false, "Enable PROXY protocol support. This causes the server to expect PROXY protocol headers on incoming connections.") + proxyProtocolReadHeaderTimeout := flag.Int("proxy_protocol_read_header_timeout_ms", 200, "Timeout for reading the PROXY protocol header in milliseconds. If the connection does not send a valid PROXY protocol header in this time, the header is ignored.") + + // Prometheus metrics flags prometheusEnabled := flag.Bool("enable_prometheus", false, "Enable prometheus") - healthcheckEnabled := flag.Bool("enable_healthcheck", false, "Enable healthcheck") 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)") + + // GeoIP flags 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.") - proxyProtocolEnabled := flag.Bool("proxy_protocol_enabled", false, "Enable PROXY protocol support. This causes the server to expect PROXY protocol headers on incoming connections.") - proxyProtocolReadHeaderTimeout := flag.Int("proxy_protocol_read_header_timeout_ms", 200, "Timeout for reading the PROXY protocol header in milliseconds. If the connection does not send a valid PROXY protocol header in this time, the header is ignored.") - healthcheckHost := flag.String("healthcheck_host", health.DefaultHost, "The address for container healthcheck") - healthcheckPort := flag.String("healthcheck_port", health.DefaultPort, "HTTP port for container healthcheck; serves JSON with status and uptime at /health") - healthcheck := flag.Bool("healthcheck", false, "GET healthcheck_host:healthcheck_port/health and exit 1 if status is not ok or timeout is exceeded (for container healthcheck)") + + // Healthcheck flags + healthcheckEnabled := flag.Bool("healthcheck_enable", false, "Enable healthcheck") + healthcheckHost := flag.String("healthcheck_host", health.DefaultHost, "The address for healthcheck.") + healthcheckPort := flag.String("healthcheck_port", health.DefaultPort, "HTTP port for healthcheck; Serves JSON with status and uptime at /health.") + healthcheck := flag.Bool("healthcheck", false, "Perform healthcheck and exit. GET healthcheck_host:healthcheck_port/health and exit 1 if status is not ok or timeout is exceeded.") flag.Usage = func() { fmt.Fprintf(flag.CommandLine.Output(), "Usage of %v \n", os.Args[0]) diff --git a/metrics/priority_queue_test.go b/metrics/priority_queue_test.go new file mode 100644 index 0000000..74b189d --- /dev/null +++ b/metrics/priority_queue_test.go @@ -0,0 +1,102 @@ +package metrics + +import ( + "testing" + "time" +) + +func TestUpdatablePriorityQueue_PushAndPop(t *testing.T) { + pq := NewUpdatablePriorityQueue() + + now := time.Now() + pq.Update("item1", now.Add(2*time.Second)) + pq.Update("item2", now.Add(1*time.Second)) + pq.Update("item3", now.Add(3*time.Second)) + + if pq.pq.Len() != 3 { + t.Fatalf("Expected length 3, got %d", pq.pq.Len()) + } + + first := pq.Pop() + if first == nil { + t.Fatal("Pop returned nil") + } + if first.Key != "item2" { + t.Errorf("Expected first to be item2, got %s", first.Key) + } + + second := pq.Pop() + if second.Key != "item1" { + t.Errorf("Expected second to be item1, got %s", second.Key) + } + + third := pq.Pop() + if third.Key != "item3" { + t.Errorf("Expected third to be item3, got %s", third.Key) + } + + if pq.Pop() != nil { + t.Error("Expected pop on empty pq to return nil") + } +} + +func TestUpdatablePriorityQueue_UpdateExisting(t *testing.T) { + pq := NewUpdatablePriorityQueue() + + now := time.Now() + pq.Update("item1", now.Add(2*time.Second)) + pq.Update("item2", now.Add(3*time.Second)) + + // Update item2 to have an earlier time + pq.Update("item2", now.Add(1*time.Second)) + + first := pq.Peek() + if first == nil { + t.Fatal("Peek returned nil") + } + if first.Key != "item2" { + t.Errorf("Expected first to be item2 after update, got %s", first.Key) + } + + popFirst := pq.Pop() + if popFirst.Key != "item2" { + t.Errorf("Expected popped first to be item2, got %s", popFirst.Key) + } + + popSecond := pq.Pop() + if popSecond.Key != "item1" { + t.Errorf("Expected popped second to be item1, got %s", popSecond.Key) + } +} + +func TestUpdatablePriorityQueue_UpdateToLater(t *testing.T) { + pq := NewUpdatablePriorityQueue() + + now := time.Now() + pq.Update("item1", now.Add(1*time.Second)) + pq.Update("item2", now.Add(2*time.Second)) + + // Update item1 to have a later time than item2 + pq.Update("item1", now.Add(3*time.Second)) + + first := pq.Pop() + if first == nil || first.Key != "item2" { + t.Errorf("Expected first to be item2, got %v", first) + } + + second := pq.Pop() + if second == nil || second.Key != "item1" { + t.Errorf("Expected second to be item1, got %v", second) + } +} + +func TestUpdatablePriorityQueue_Empty(t *testing.T) { + pq := NewUpdatablePriorityQueue() + + if pq.Peek() != nil { + t.Error("Expected Peek on empty pq to return nil") + } + if pq.Pop() != nil { + t.Error("Expected Pop on empty pq to return nil") + } +}