proxy: limit concurrent traceroute requests

Author: xddxdd

proxy/main.go

@@ -61,19 +61,21 @@ func accessHandler(next http.Handler) http.Handler {
 }
 
 type settingType struct {
-	birdSocket  string
-	listen      []string
-	allowedNets []*net.IPNet
-	tr_bin      string
-	tr_flags    []string
-	tr_raw      bool
+	birdSocket        string
+	listen            []string
+	allowedNets       []*net.IPNet
+	tr_bin            string
+	tr_flags          []string
+	tr_raw            bool
+	tr_max_concurrent int
 }
 
 var setting settingType
 
 // Wrapper of tracer
 func main() {
 	parseSettings()
+	initTracerouteSemaphore(setting.tr_max_concurrent)
 	tracerouteAutodetect()
 
 	mux := http.NewServeMux()

proxy/settings.go

@@ -11,12 +11,13 @@ import (
 )
 
 type viperSettingType struct {
-	BirdSocket      string   `mapstructure:"bird_socket"`
-	Listen          []string `mapstructure:"listen"`
-	AllowedNets     string   `mapstructure:"allowed_ips"`
-	TracerouteBin   string   `mapstructure:"traceroute_bin"`
-	TracerouteFlags string   `mapstructure:"traceroute_flags"`
-	TracerouteRaw   bool     `mapstructure:"traceroute_raw"`
+	BirdSocket              string   `mapstructure:"bird_socket"`
+	Listen                  []string `mapstructure:"listen"`
+	AllowedNets             string   `mapstructure:"allowed_ips"`
+	TracerouteBin           string   `mapstructure:"traceroute_bin"`
+	TracerouteFlags         string   `mapstructure:"traceroute_flags"`
+	TracerouteRaw           bool     `mapstructure:"traceroute_raw"`
+	TracerouteMaxConcurrent int      `mapstructure:"traceroute_max_concurrent"`
 }
 
 // Parse settings with viper, and convert to legacy setting format
@@ -52,6 +53,9 @@ func parseSettings() {
 	pflag.Bool("traceroute_raw", false, "whether to display traceroute outputs raw; set via parameter or environment variable BIRDLG_TRACEROUTE_RAW")
 	viper.BindPFlag("traceroute_raw", pflag.Lookup("traceroute_raw"))
 
+	pflag.Int("traceroute_max_concurrent", 10, "max concurrent traceroute requests allowed")
+	viper.BindPFlag("traceroute_max_concurrent", pflag.Lookup("traceroute_max_concurrent"))
+
 	pflag.Parse()
 
 	if err := viper.ReadInConfig(); err != nil {
@@ -101,6 +105,7 @@ func parseSettings() {
 	}
 
 	setting.tr_raw = viperSettings.TracerouteRaw
+	setting.tr_max_concurrent = viperSettings.TracerouteMaxConcurrent
 
 	fmt.Printf("%#v\n", setting)
 }

proxy/traceroute.go

@@ -11,6 +11,14 @@ import (
 	"github.com/google/shlex"
 )
 
+var tracerouteSemaphore chan struct{}
+
+func initTracerouteSemaphore(maxConcurrent int) {
+	if maxConcurrent > 0 {
+		tracerouteSemaphore = make(chan struct{}, maxConcurrent)
+	}
+}
+
 func tracerouteArgsToString(cmd string, args []string, target []string) string {
 	var cmdCombined = append([]string{cmd}, args...)
 	cmdCombined = append(cmdCombined, target...)
@@ -83,6 +91,20 @@ func tracerouteAutodetect() {
 }
 
 func tracerouteHandler(httpW http.ResponseWriter, httpR *http.Request) {
+	// Check concurrency limit
+	if setting.tr_max_concurrent > 0 {
+		select {
+		case tracerouteSemaphore <- struct{}{}:
+			// Successfully acquired semaphore slot
+			defer func() { <-tracerouteSemaphore }()
+		default:
+			// Semaphore is full, reject request
+			httpW.WriteHeader(http.StatusServiceUnavailable)
+			httpW.Write([]byte("Too many concurrent traceroute requests. Please try again later.\n"))
+			return
+		}
+	}
+
 	query := string(httpR.URL.Query().Get("q"))
 	query = strings.TrimSpace(query)
 

proxy/traceroute_test.go

@@ -97,6 +97,7 @@ func TestTracerouteAutodetectFlagsOnly(t *testing.T) {
 }
 
 func TestTracerouteHandlerWithoutQuery(t *testing.T) {
+	initTracerouteSemaphore(setting.tr_max_concurrent)
 	r := httptest.NewRequest(http.MethodGet, "/traceroute", nil)
 	w := httptest.NewRecorder()
 	tracerouteHandler(w, r)
@@ -107,6 +108,7 @@ func TestTracerouteHandlerWithoutQuery(t *testing.T) {
 }
 
 func TestTracerouteHandlerShlexError(t *testing.T) {
+	initTracerouteSemaphore(setting.tr_max_concurrent)
 	r := httptest.NewRequest(http.MethodGet, "/traceroute?q="+url.QueryEscape("\"1.1.1.1"), nil)
 	w := httptest.NewRecorder()
 	tracerouteHandler(w, r)
@@ -117,6 +119,7 @@ func TestTracerouteHandlerShlexError(t *testing.T) {
 }
 
 func TestTracerouteHandlerNoTracerouteFound(t *testing.T) {
+	initTracerouteSemaphore(setting.tr_max_concurrent)
 	setting.tr_bin = ""
 	setting.tr_flags = nil
 
@@ -130,6 +133,7 @@ func TestTracerouteHandlerNoTracerouteFound(t *testing.T) {
 }
 
 func TestTracerouteHandlerExecuteError(t *testing.T) {
+	initTracerouteSemaphore(setting.tr_max_concurrent)
 	setting.tr_bin = "sh"
 	setting.tr_flags = []string{"-c", "false"}
 	setting.tr_raw = true
@@ -144,6 +148,7 @@ func TestTracerouteHandlerExecuteError(t *testing.T) {
 }
 
 func TestTracerouteHandlerRaw(t *testing.T) {
+	initTracerouteSemaphore(setting.tr_max_concurrent)
 	setting.tr_bin = "sh"
 	setting.tr_flags = []string{"-c", "echo Mock"}
 	setting.tr_raw = true
@@ -156,6 +161,7 @@ func TestTracerouteHandlerRaw(t *testing.T) {
 }
 
 func TestTracerouteHandlerPostprocess(t *testing.T) {
+	initTracerouteSemaphore(setting.tr_max_concurrent)
 	setting.tr_bin = "sh"
 	setting.tr_flags = []string{"-c", "echo \"first line\n 2 *\nthird line\""}
 	setting.tr_raw = false
@@ -166,3 +172,49 @@ func TestTracerouteHandlerPostprocess(t *testing.T) {
 	assert.Equal(t, w.Code, http.StatusOK)
 	assert.Equal(t, w.Body.String(), "first line\nthird line\n\n1 hops not responding.")
 }
+
+func TestTracerouteHandlerConcurrencyLimit(t *testing.T) {
+	// Set a low limit for testing
+	maxConcurrent := 2
+	initTracerouteSemaphore(maxConcurrent)
+	setting.tr_max_concurrent = maxConcurrent
+
+	// Use a slow command to keep requests running
+	setting.tr_bin = "sh"
+	setting.tr_flags = []string{"-c", "sleep 1; echo Done"}
+	setting.tr_raw = true
+
+	// Launch more concurrent requests than the limit
+	numRequests := 5
+	responses := make(chan int, numRequests)
+
+	for i := 0; i < numRequests; i++ {
+		go func() {
+			r := httptest.NewRequest(http.MethodGet, "/traceroute?q="+url.QueryEscape("1.1.1.1"), nil)
+			w := httptest.NewRecorder()
+			tracerouteHandler(w, r)
+			responses <- w.Code
+		}()
+	}
+
+	// Collect all responses
+	statusCodes := make(map[int]int)
+	for i := 0; i < numRequests; i++ {
+		code := <-responses
+		statusCodes[code]++
+	}
+
+	// Verify that some requests succeeded (200) and some were rejected (503)
+	if statusCodes[http.StatusOK] == 0 {
+		t.Error("Expected at least one request to succeed with 200")
+	}
+	if statusCodes[http.StatusServiceUnavailable] == 0 {
+		t.Error("Expected at least one request to be rejected with 503")
+	}
+
+	// Verify we didn't get any unexpected status codes
+	totalRequests := statusCodes[http.StatusOK] + statusCodes[http.StatusServiceUnavailable]
+	if totalRequests != numRequests {
+		t.Errorf("Expected %d total requests, got %d", numRequests, totalRequests)
+	}
+}