From f49f8bac5e98d063743bd4923daafe2a5be03886 Mon Sep 17 00:00:00 2001
From: Lan Tian <xuyh0120@outlook.com>
Date: Sat, 27 Feb 2021 16:42:42 +0800
Subject: [PATCH] proxy: support arbitraty traceroute arguments

---
 proxy/bird.go       |  4 +++-
 proxy/go.mod        |  5 ++++-
 proxy/go.sum        |  2 ++
 proxy/traceroute.go | 20 +++++++++++++++-----
 4 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/proxy/bird.go b/proxy/bird.go
index f660dad..e74a00b 100644
--- a/proxy/bird.go
+++ b/proxy/bird.go
@@ -59,7 +59,9 @@ func birdHandler(httpW http.ResponseWriter, httpR *http.Request) {
 		// Initialize BIRDv4 socket
 		bird, err := net.Dial("unix", setting.birdSocket)
 		if err != nil {
-			panic(err)
+			httpW.WriteHeader(http.StatusInternalServerError)
+			httpW.Write([]byte(err.Error()))
+			return
 		}
 		defer bird.Close()
 
diff --git a/proxy/go.mod b/proxy/go.mod
index cd5f70f..8195ce6 100644
--- a/proxy/go.mod
+++ b/proxy/go.mod
@@ -2,4 +2,7 @@ module github.com/xddxdd/bird-lg-go/proxy
 
 go 1.15
 
-require github.com/gorilla/handlers v1.5.1
+require (
+	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
+	github.com/gorilla/handlers v1.5.1
+)
diff --git a/proxy/go.sum b/proxy/go.sum
index f27a04f..4517091 100644
--- a/proxy/go.sum
+++ b/proxy/go.sum
@@ -1,4 +1,6 @@
 github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
 github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/gorilla/handlers v1.5.1 h1:9lRY6j8DEeeBT10CvO9hGW0gmky0BprnvDI5vfhUHH4=
 github.com/gorilla/handlers v1.5.1/go.mod h1:t8XrUpc4KVXb7HGyJ4/cEnwQiaxrX/hz1Zv/4g96P1Q=
diff --git a/proxy/traceroute.go b/proxy/traceroute.go
index d14a3ab..7294abd 100644
--- a/proxy/traceroute.go
+++ b/proxy/traceroute.go
@@ -2,13 +2,14 @@ package main
 
 import (
 	"fmt"
-	"html"
 	"net/http"
 	"os/exec"
 	"regexp"
 	"runtime"
 	"strconv"
 	"strings"
+
+	"github.com/google/shlex"
 )
 
 func tracerouteTryExecute(cmd []string, args [][]string) ([]byte, string) {
@@ -31,10 +32,17 @@ func tracerouteTryExecute(cmd []string, args [][]string) ([]byte, string) {
 func tracerouteHandler(httpW http.ResponseWriter, httpR *http.Request) {
 	query := string(httpR.URL.Query().Get("q"))
 	query = strings.TrimSpace(query)
-	query = html.EscapeString(query)
+
 	if query == "" {
 		invalidHandler(httpW, httpR)
 	} else {
+		args, err := shlex.Split(query)
+		if err != nil {
+			httpW.WriteHeader(http.StatusInternalServerError)
+			httpW.Write([]byte(fmt.Sprintf("failed to parse args: %s\n", err.Error())))
+			return
+		}
+
 		var result []byte
 		var errString string
 		skippedCounter := 0
@@ -45,7 +53,8 @@ func tracerouteHandler(httpW http.ResponseWriter, httpR *http.Request) {
 					"traceroute",
 				},
 				[][]string{
-					{"-q1", "-w1", query},
+					append([]string{"-q1", "-w1"}, args...),
+					args,
 				},
 			)
 		} else if runtime.GOOS == "linux" {
@@ -55,8 +64,9 @@ func tracerouteHandler(httpW http.ResponseWriter, httpR *http.Request) {
 					"traceroute",
 				},
 				[][]string{
-					{"-q1", "-N32", "-w1", query},
-					{"-q1", "-w1", query},
+					append([]string{"-q1", "-N32", "-w1"}, args...),
+					append([]string{"-q1", "-w1"}, args...),
+					args,
 				},
 			)
 		} else {