From c922eba2d82503e6ef187d3d3f06572d4b0bacec Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Sat, 24 Sep 2022 21:28:09 +0100
Subject: [PATCH] Fix sending arguments to the admin socket in `yggdrasilctl`

---
 cmd/yggdrasilctl/main.go | 14 ++++++++------
 src/admin/admin.go       | 27 ++++++++++++++-------------
 src/core/nodeinfo.go     |  8 ++++++--
 3 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/cmd/yggdrasilctl/main.go b/cmd/yggdrasilctl/main.go
index 1e93b98f..f57ab79a 100644
--- a/cmd/yggdrasilctl/main.go
+++ b/cmd/yggdrasilctl/main.go
@@ -85,7 +85,7 @@ func run() int {
 	encoder := json.NewEncoder(conn)
 	send := &admin.AdminSocketRequest{}
 	recv := &admin.AdminSocketResponse{}
-
+	args := map[string]string{}
 	for c, a := range cmdLineEnv.args {
 		if c == 0 {
 			if strings.HasPrefix(a, "-") {
@@ -96,15 +96,17 @@ func run() int {
 			send.Name = a
 			continue
 		}
-		tokens := strings.SplitN(a, "=", 1)
+		tokens := strings.SplitN(a, "=", 2)
 		switch {
 		case len(tokens) == 1:
-			panic("incomplete argument supplied")
+			logger.Println("Ignoring invalid argument:", a)
 		default:
-			send.Arguments[tokens[0]] = tokens[1]
+			args[tokens[0]] = tokens[1]
 		}
 	}
-
+	if send.Arguments, err = json.Marshal(args); err != nil {
+		panic(err)
+	}
 	if err := encoder.Encode(&send); err != nil {
 		panic(err)
 	}
@@ -139,7 +141,7 @@ func run() int {
 	table.SetNoWhiteSpace(true)
 	table.SetAutoWrapText(false)
 
-	switch strings.ToLower(recv.Request.Name) {
+	switch strings.ToLower(send.Name) {
 	case "list":
 		var resp admin.ListResponse
 		if err := json.Unmarshal(recv.Response, &resp); err != nil {
diff --git a/src/admin/admin.go b/src/admin/admin.go
index 376f79a6..ca1d95ef 100644
--- a/src/admin/admin.go
+++ b/src/admin/admin.go
@@ -30,16 +30,16 @@ type AdminSocket struct {
 }
 
 type AdminSocketRequest struct {
-	Name      string            `json:"request"`
-	Arguments map[string]string `json:"arguments,omitempty"`
-	KeepAlive bool              `json:"keepalive,omitempty"`
+	Name      string          `json:"request"`
+	Arguments json.RawMessage `json:"arguments,omitempty"`
+	KeepAlive bool            `json:"keepalive,omitempty"`
 }
 
 type AdminSocketResponse struct {
-	Status   string             `json:"status"`
-	Error    string             `json:"error,omitempty"`
-	Request  AdminSocketRequest `json:"request"`
-	Response json.RawMessage    `json:"response"`
+	Status   string          `json:"status"`
+	Error    string          `json:"error,omitempty"`
+	Request  json.RawMessage `json:"request"`
+	Response json.RawMessage `json:"response"`
 }
 
 type handler struct {
@@ -298,25 +298,26 @@ func (a *AdminSocket) handleRequest(conn net.Conn) {
 	for {
 		var err error
 		var buf json.RawMessage
+		var req AdminSocketRequest
 		var resp AdminSocketResponse
 		if err := func() error {
 			if err = decoder.Decode(&buf); err != nil {
 				return fmt.Errorf("Failed to find request")
 			}
-			if err = json.Unmarshal(buf, &resp.Request); err != nil {
+			if err = json.Unmarshal(buf, &req); err != nil {
 				return fmt.Errorf("Failed to unmarshal request")
 			}
-			if resp.Request.Name == "" {
+			if req.Name == "" {
 				return fmt.Errorf("No request specified")
 			}
-			reqname := strings.ToLower(resp.Request.Name)
+			reqname := strings.ToLower(req.Name)
 			handler, ok := a.handlers[reqname]
 			if !ok {
 				return fmt.Errorf("Unknown action '%s', try 'list' for help", reqname)
 			}
-			res, err := handler.handler(buf)
+			res, err := handler.handler(req.Arguments)
 			if err != nil {
-				return fmt.Errorf("Handler returned error: %w", err)
+				return err
 			}
 			if resp.Response, err = json.Marshal(res); err != nil {
 				return fmt.Errorf("Failed to marshal response: %w", err)
@@ -330,7 +331,7 @@ func (a *AdminSocket) handleRequest(conn net.Conn) {
 		if err = encoder.Encode(resp); err != nil {
 			a.log.Debugln("Encode error:", err)
 		}
-		if !resp.Request.KeepAlive {
+		if !req.KeepAlive {
 			break
 		} else {
 			continue
diff --git a/src/core/nodeinfo.go b/src/core/nodeinfo.go
index a6132ec2..e02c79c7 100644
--- a/src/core/nodeinfo.go
+++ b/src/core/nodeinfo.go
@@ -4,6 +4,7 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"runtime"
 	"strings"
 	"time"
@@ -160,11 +161,14 @@ func (m *nodeinfo) nodeInfoAdminHandler(in json.RawMessage) (interface{}, error)
 	if err := json.Unmarshal(in, &req); err != nil {
 		return nil, err
 	}
+	if req.Key == "" {
+		return nil, fmt.Errorf("No remote public key supplied")
+	}
 	var key keyArray
 	var kbs []byte
 	var err error
 	if kbs, err = hex.DecodeString(req.Key); err != nil {
-		return nil, err
+		return nil, fmt.Errorf("Failed to decode public key: %w", err)
 	}
 	copy(key[:], kbs)
 	ch := make(chan []byte, 1)
@@ -175,7 +179,7 @@ func (m *nodeinfo) nodeInfoAdminHandler(in json.RawMessage) (interface{}, error)
 	defer timer.Stop()
 	select {
 	case <-timer.C:
-		return nil, errors.New("timeout")
+		return nil, errors.New("Timed out waiting for response")
 	case info := <-ch:
 		var msg json.RawMessage
 		if err := msg.UnmarshalJSON(info); err != nil {