Merge pull request #221 from shizunge/prometheus

Rename enable_healthcheck to healthcheck_enabled. Add prometheus_enabled.

.github/workflows/on-pull-request.yml

@@ -10,57 +10,19 @@ on:
       - 'README.md'
       - 'LICENSE'
   workflow_dispatch:
-
+ 
-env:
-  PLATFORMS: "linux/amd64,linux/arm64,linux/arm/v7"
-  
 jobs:
   test:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
         uses: actions/checkout@v7
-
       - name: Set up Go
         uses: actions/setup-go@v7
         with:
           go-version: '1.26.1'
           cache: true
-
       - name: Install dependencies
         run: go mod download
-
       - name: Run Tests
         run: go test -v ./...
-
-  build_container_image:
-    name: Build Docker image
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v7
-      - name: Set up QEMU
-        uses: docker/setup-qemu-action@v4
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v4.2.0
-      - name: Docker meta
-        id: meta
-        uses: docker/metadata-action@v6
-        with:
-          images: |
-            ghcr.io/${{ github.repository }}-development
-          tags: |
-            type=raw,value=dev-{{date 'X'}}
-            type=raw,value=latest
-            type=ref,event=branch
-            type=edge,branch=main
-      - name: Build 
-        uses: docker/build-push-action@v7.3.0
-        with:
-          platforms: ${{ env.PLATFORMS }}
-          push: false
-          tags: ${{ steps.meta.outputs.tags }}
-          labels: ${{ steps.meta.outputs.labels }}
-          provenance: false
-
-

.github/workflows/on-push.yml

@@ -21,21 +21,19 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v7
-
       - name: Set up Go
         uses: actions/setup-go@v7
         with:
           go-version: '1.26.1'
           cache: true
-
       - name: Install dependencies
         run: go mod download
-
       - name: Run Tests
         run: go test -v ./...
   build_and_push:
     name: Build and push Docker image
     runs-on: ubuntu-latest
+    needs: test
     if: ${{ github.actor != 'dependabot[bot]' }}
     steps:
       - name: Checkout code

.github/workflows/on-release.yml

@@ -8,9 +8,25 @@ env:
   PLATFORMS: "linux/amd64,linux/arm64,linux/arm/v7"
   
 jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v7
+      - name: Set up Go
+        uses: actions/setup-go@v7
+        with:
+          go-version: '1.26.1'
+          cache: true
+      - name: Install dependencies
+        run: go mod download
+      - name: Run Tests
+        run: go test -v ./...
+
   build_and_push:
     name: Build and push Docker image
     runs-on: ubuntu-latest
+    needs: test
     steps:
       - name: Checkout code
         uses: actions/checkout@v7

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_enabled"]

README.md

@@ -55,15 +55,17 @@ Usage of ./endlessh-go
   -conn_type string
         Connection type. Possible values are tcp, tcp4, tcp6 (default "tcp")
   -enable_prometheus
-        Enable prometheus
+        Enable prometheus (deprecated, use prometheus_enabled)
   -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_enabled
+        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
@@ -88,6 +90,8 @@ Usage of ./endlessh-go
         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_enabled
+        Enable prometheus
   -prometheus_entry string
         Entry point for prometheus (default "metrics")
   -prometheus_host string
@@ -119,7 +123,7 @@ Endlessh-go exports the following Prometheus metrics.
 | 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`: 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`.
+The metrics is off by default, you can turn it via the CLI argument `-prometheus_enabled`.
 
 It listens to port `2112` and entry point is `/metrics` by default. The port and entry point can be changed via CLI arguments.
 
@@ -129,7 +133,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_enabled` flag is set. By default it listens on `127.0.0.1:51000` at `/health` and returns a JSON like this:
 
 ```json
 {

endlessh_integration_test.go

@@ -263,7 +263,7 @@ func TestEndlesshIntegration_PrometheusMetrics(t *testing.T) {
 	cmd := exec.Command(
 		"go", "run", "main.go",
 		"-port=0",
-		"-enable_prometheus",
+		"-prometheus_enabled",
 		"-prometheus_port=0",
 		"-interval_ms=100",
 		"-logtostderr", "-v=1",
@@ -322,3 +322,61 @@ func TestEndlesshIntegration_PrometheusMetrics(t *testing.T) {
 		t.Errorf("Expected bytes metric not found:\n%s", body)
 	}
 }
+
+func TestEndlesshIntegration_Healthcheck(t *testing.T) {
+	var stderr bytes.Buffer
+
+	cmd := exec.Command(
+		"go", "run", "main.go",
+		"-port=0",
+		"-healthcheck_enabled",
+		"-healthcheck_port=0",
+		"-interval_ms=100",
+		"-logtostderr", "-v=1",
+	)
+	cmd.Stderr = &stderr
+
+	if err := cmd.Start(); err != nil {
+		t.Fatalf("Failed to start server: %v", err)
+	}
+	defer cmd.Process.Kill()
+
+	if !waitForLogMatch(&stderr, "Starting healthcheck on http", waitForListenTimeout) {
+		t.Fatalf("Healthcheck listener did not start: %s", stderr.String())
+	}
+
+	reHealth := regexp.MustCompile(`Starting healthcheck on http://.*:(\d+)/health`)
+	healthMatch := reHealth.FindStringSubmatch(stderr.String())
+	if len(healthMatch) < 2 {
+		t.Fatalf("Could not parse healthcheck port: %s", stderr.String())
+	}
+	healthPort := healthMatch[1]
+
+	// 1. Verify healthcheck HTTP endpoint returns valid JSON
+	resp, err := net.Dial("tcp", "localhost:"+healthPort)
+	if err != nil {
+		t.Fatalf("Failed to connect to healthcheck endpoint: %v", err)
+	}
+	fmt.Fprintf(resp, "GET /health HTTP/1.1\r\nHost: localhost\r\n\r\n")
+
+	buf := make([]byte, 8192)
+	n, _ := resp.Read(buf)
+	body := string(buf[:n])
+	resp.Close()
+
+	if !strings.Contains(body, `"status":"ok"`) {
+		t.Errorf("Expected status:ok in health response: %s", body)
+	}
+
+	// 2. Run the probe mode (-healthcheck) and verify it exits with 0
+	probeCmd := exec.Command(
+		"go", "run", "main.go",
+		"-healthcheck",
+		"-healthcheck_port="+healthPort,
+	)
+	var probeStderr bytes.Buffer
+	probeCmd.Stderr = &probeStderr
+	if err := probeCmd.Run(); err != nil {
+		t.Errorf("Healthcheck probe failed: %v, stderr: %s", err, probeStderr.String())
+	}
+}

examples/docker-simple/docker-compose.yml

@@ -8,7 +8,7 @@ services:
       - -interval_ms=1000
       - -logtostderr
       - -v=1
-      - -enable_prometheus
+      - -prometheus_enabled
       - -geoip_supplier=ip-api
     networks:
       - example_network

examples/kustomize-simple/deployment.yaml

@@ -14,7 +14,7 @@ spec:
             - -interval_ms=1000
             - -logtostderr
             - -v=1
-            - -enable_prometheus
+            - -prometheus_enabled
             - -geoip_supplier=ip-api
             - -host=[::]
             - -prometheus_host=[::]

health/health.go

@@ -21,6 +21,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"strconv"
 	"time"
 
 	"github.com/golang/glog"
@@ -84,3 +85,29 @@ func Probe(host, port string) bool {
 	}
 	return body.Status == "ok"
 }
+
+// SetupHealthcheck probes or starts the healthcheck listener as configured.
+func SetupHealthcheck(healthcheck bool, healthcheckEnabled bool, connType string, healthcheckHost, healthcheckPort *string) {
+	if healthcheck {
+		if !Probe(*healthcheckHost, *healthcheckPort) {
+			os.Exit(1)
+		}
+		os.Exit(0)
+	}
+
+	if healthcheckEnabled {
+		if connType == "tcp6" && *healthcheckHost == "0.0.0.0" {
+			*healthcheckHost = "[::]"
+		}
+		if *healthcheckPort == "0" || *healthcheckPort == "" {
+			l, err := net.Listen("tcp", *healthcheckHost+":0")
+			if err != nil {
+				glog.Fatalf("Failed to pick a free healthcheck port: %v", err)
+			}
+			actualPort := l.Addr().(*net.TCPAddr).Port
+			*healthcheckPort = strconv.Itoa(actualPort)
+			l.Close()
+		}
+		StartListener(*healthcheckHost, *healthcheckPort)
+	}
+}

main.go

@@ -131,25 +131,35 @@ 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")
+	// PROXY protocol flags
-	flag.Var(&connPorts, "port", fmt.Sprintf("SSH listening port. You may provide multiple -port flags to listen to multiple ports. (default %q)", defaultPort))
+	proxyProtocolEnabled := flag.Bool("proxy_protocol_enabled", false, "Enable PROXY protocol support. This causes the server to expect PROXY protocol headers on incoming connections.")
-	prometheusEnabled := flag.Bool("enable_prometheus", false, "Enable prometheus")
+	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.")
-	healthcheckEnabled := flag.Bool("enable_healthcheck", false, "Enable healthcheck")
+
+	// Prometheus metrics flags
+	prometheusEnabledOld := flag.Bool("enable_prometheus", false, "Enable prometheus (deprecated, use prometheus_enabled)")
+	prometheusEnabledNew := flag.Bool("prometheus_enabled", 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)")
+
+	// 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.")
+	// Healthcheck flags
-	healthcheckHost := flag.String("healthcheck_host", health.DefaultHost, "The address for container healthcheck")
+	healthcheckEnabled := flag.Bool("healthcheck_enabled", false, "Enable healthcheck")
-	healthcheckPort := flag.String("healthcheck_port", health.DefaultPort, "HTTP port for container healthcheck; serves JSON with status and uptime at /health")
+	healthcheckHost := flag.String("healthcheck_host", health.DefaultHost, "The address for healthcheck.")
-	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)")
+	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])
@@ -157,14 +167,20 @@ func main() {
 	}
 	flag.Parse()
 
-	if *healthcheck {
+	health.SetupHealthcheck(*healthcheck, *healthcheckEnabled, *connType, healthcheckHost, healthcheckPort)
-		if !health.Probe(*healthcheckHost, *healthcheckPort) {
+
-			os.Exit(1)
+	prometheusEnabled := *prometheusEnabledNew
+	prometheusEnableSet := false
+	flag.Visit(func(f *flag.Flag) {
+		if f.Name == "prometheus_enabled" {
+			prometheusEnableSet = true
 		}
-		os.Exit(0)
+	})
+	if !prometheusEnableSet {
+		prometheusEnabled = *prometheusEnabledOld
 	}
 
-	if *prometheusEnabled {
+	if prometheusEnabled {
 		if *connType == "tcp6" && *prometheusHost == "0.0.0.0" {
 			*prometheusHost = "[::]"
 		}
@@ -180,28 +196,13 @@ func main() {
 		metrics.InitPrometheus(*prometheusHost, *prometheusPort, *prometheusEntry)
 	}
 
-	records := metrics.StartRecording(*maxClients, *prometheusEnabled, *prometheusCleanUnseenSeconds,
+	records := metrics.StartRecording(*maxClients, prometheusEnabled, *prometheusCleanUnseenSeconds,
 		geoip.GeoOption{
 			GeoipSupplier:     *geoipSupplier,
 			MaxMindDbFileName: *maxMindDbFileName,
 		})
 	clients := startSending(*maxClients, *bannerMaxLength, records)
 
-	if *healthcheckEnabled {
-		if *connType == "tcp6" && *healthcheckHost == "0.0.0.0" {
-			*healthcheckHost = "[::]"
-		}
-		if *healthcheckPort == "0" || *healthcheckPort == "" {
-			l, err := net.Listen("tcp", *healthcheckHost+":0")
-			if err != nil {
-				glog.Fatalf("Failed to pick a free healthcheck port: %v", err)
-			}
-			actualPort := l.Addr().(*net.TCPAddr).Port
-			*healthcheckPort = strconv.Itoa(actualPort)
-			l.Close()
-		}
-		health.StartListener(*healthcheckHost, *healthcheckPort)
-	}
 
 	interval := time.Duration(*intervalMs) * time.Millisecond
 	// Listen for incoming connections.

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")
+	}
+}