Added handling of —permit-arguments option

app/app.go

@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"io/ioutil"
 	"log"
@@ -26,6 +27,11 @@ import (
 	"github.com/kr/pty"
 )
 
+type InitMessage struct {
+	Arguments string `json:"Arguments,omitempty"`
+	AuthToken string `json:"AuthToken,omitempty"`
+}
+
 type App struct {
 	command []string
 	options *Options
@@ -54,6 +60,7 @@ type Options struct {
 	EnableReconnect     bool                   `hcl:"enable_reconnect"`
 	ReconnectTime       int                    `hcl:"reconnect_time"`
 	Once                bool                   `hcl:"once"`
+	PermitArguments     bool                   `hcl:"permit_arguments"`
 	Preferences         map[string]interface{} `hcl:"preferences"`
 }
 
@@ -272,14 +279,42 @@ func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	_, initMessage, err := conn.ReadMessage()
-	if err != nil || string(initMessage) != app.options.Credential {
+	_, stream, err := conn.ReadMessage()
+	if err != nil {
 		log.Print("Failed to authenticate websocket connection")
 		conn.Close()
 		return
 	}
+	var init InitMessage
 
-	cmd := exec.Command(app.command[0], app.command[1:]...)
+	err = json.Unmarshal(stream, &init)
+	if err != nil {
+		log.Printf("Failed to parse init message %v", err)
+		conn.Close()
+		return
+	}
+	if init.AuthToken != app.options.Credential {
+		log.Print("Failed to authenticate websocket connection")
+		conn.Close()
+		return
+	}
+	argv := app.command[1:]
+	if app.options.PermitArguments {
+		if init.Arguments == "" {
+			init.Arguments = "?"
+		}
+		query, err := url.Parse(init.Arguments)
+		if err != nil {
+			log.Print("Failed to parse arguments")
+			conn.Close()
+			return
+		}
+		params := query.Query()["arg"]
+		if len(params) != 0 {
+			argv = append(argv, params...)
+		}
+	}
+	cmd := exec.Command(app.command[0], argv...)
 	ptyIo, err := pty.Start(cmd)
 	if err != nil {
 		log.Print("Failed to execute command")