Merge pull request #148 from davrodpin/misc/rpc

Add rpc command

Author: davrodpin

cmd/misc.go

@@ -0,0 +1,16 @@
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+)
+
+var miscCmd = &cobra.Command{
+	Use:   "misc",
+	Short: "A set of miscelaneous commands",
+	Args:  cobra.MinimumNArgs(1),
+	Run:   func(cmd *cobra.Command, arg []string) {},
+}
+
+func init() {
+	rootCmd.AddCommand(miscCmd)
+}

cmd/misc_rpc.go

@@ -0,0 +1,50 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/davrodpin/mole/mole"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+)
+
+var (
+	method, params string
+
+	miscRpcCmd = &cobra.Command{
+		Use:   "rpc [alias or id] [method] [params]",
+		Short: "Executes a remote procedure call on a given mole instance",
+		Long:  "Executes a remote procedure call on a given mole instance",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) < 2 {
+				return fmt.Errorf("not enough arguments.")
+			}
+
+			id = args[0]
+			method = args[1]
+
+			if len(args) == 3 {
+				params = args[2]
+			}
+
+			return nil
+		},
+		Run: func(cmd *cobra.Command, arg []string) {
+			resp, err := mole.Rpc(id, method, params)
+			if err != nil {
+				log.WithError(err).WithFields(log.Fields{
+					"id": id,
+				}).Error("error executing remote procedure.")
+				os.Exit(1)
+			}
+
+			fmt.Printf(resp)
+		},
+	}
+)
+
+func init() {
+	miscCmd.AddCommand(miscRpcCmd)
+}

count.out

@@ -56,3 +56,14 @@ func CreateInstanceDir(appId string) (string, error) {
 
 	return d, nil
 }
+
+// InstanceDir returns the location where all files related to a specific mole
+// instance are persisted.
+func InstanceDir(id string) (string, error) {
+	home, err := Dir()
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(home, id), nil
+}

fsutils/fsutils.go

@@ -1,13 +1,17 @@
 package mole
 
 import (
+	"context"
+	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strconv"
 
 	"github.com/davrodpin/mole/fsutils"
+	"github.com/davrodpin/mole/rpc"
+
 	"github.com/hpcloud/tail"
 )
 
@@ -129,3 +133,30 @@ func ShowLogs(id string, follow bool) error {
 
 	return nil
 }
+
+// Rpc calls a remote procedure on another mole instance given its id or alias.
+func Rpc(id, method string, params interface{}) (string, error) {
+	d, err := fsutils.InstanceDir(id)
+	if err != nil {
+		return "", err
+	}
+
+	rf := filepath.Join(d, "rpc")
+
+	addr, err := ioutil.ReadFile(rf)
+	if err != nil {
+		return "", err
+	}
+
+	resp, err := rpc.Call(context.Background(), string(addr), method, params)
+	if err != nil {
+		return "", err
+	}
+
+	r, err := json.MarshalIndent(resp, "", "  ")
+	if err != nil {
+		return "", err
+	}
+
+	return string(r), nil
+}

mole/app.go

@@ -2,18 +2,16 @@ package rpc
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"net"
 	"sync"
 
-	"encoding/json"
-
 	log "github.com/sirupsen/logrus"
 	"github.com/sourcegraph/jsonrpc2"
 )
 
 var registeredMethods = sync.Map{}
-var listener net.Listener
 
 const (
 	// DefaultAddress is the network address used by the rpc server if none is given.
@@ -29,7 +27,7 @@ func Start(address string) (net.Addr, error) {
 		address = DefaultAddress
 	}
 
-	listener, err = net.Listen("tcp", address)
+	lis, err := net.Listen("tcp", address)
 	if err != nil {
 		return nil, err
 	}
@@ -39,7 +37,7 @@ func Start(address string) (net.Addr, error) {
 
 	go func() {
 		for {
-			conn, err := listener.Accept()
+			conn, err := lis.Accept()
 			if err != nil {
 				log.WithError(err).Warnf("error establishing connection with rpc client.")
 			}
@@ -48,7 +46,7 @@ func Start(address string) (net.Addr, error) {
 		}
 	}()
 
-	return listener.Addr(), nil
+	return lis.Addr(), nil
 }
 
 // Handler handles JSON-RPC requests and notifications.
@@ -79,8 +77,7 @@ func (h *Handler) Handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2
 		return
 	}
 
-	m, _ := registeredMethods.Load(req.Method)
-	rm, err := m.(Method)(req.Params)
+	params, err := req.Params.MarshalJSON()
 	if err != nil {
 		log.WithFields(log.Fields{
 			"notification": req.Notif,
@@ -96,9 +93,31 @@ func (h *Handler) Handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2
 			})
 
 			sendResponse(ctx, conn, req, resp)
+		}
+
+		return
+	}
+
+	m, _ := registeredMethods.Load(req.Method)
+	rm, err := m.(Method)(params)
+	if err != nil {
+		log.WithFields(log.Fields{
+			"notification": req.Notif,
+			"method":       req.Method,
+			"id":           req.ID,
+		}).WithError(err).Warn("error executing rpc method.")
 
-			return
+		if !req.Notif {
+			resp := &jsonrpc2.Response{}
+			resp.SetResult(jsonrpc2.Error{
+				Code:    jsonrpc2.CodeInternalError,
+				Message: fmt.Sprintf("error executing rpc method %s", req.Method),
+			})
+
+			sendResponse(ctx, conn, req, resp)
 		}
+
+		return
 	}
 
 	if !req.Notif {
@@ -116,10 +135,10 @@ func Register(name string, method Method) {
 // Method represents a procedure that can be called remotely.
 type Method func(params interface{}) (json.RawMessage, error)
 
-// Call initiates a JSON-RPC call using the specified method and waits for the
-// response.
-func Call(ctx context.Context, method string, params interface{}) (map[string]interface{}, error) {
-	tc, err := net.Dial("tcp", listener.Addr().String())
+// Call initiates a JSON-RPC call to a given rpc server address, using the
+// specified method and waits for the response.
+func Call(ctx context.Context, addr, method string, params interface{}) (map[string]interface{}, error) {
+	tc, err := net.Dial("tcp", addr)
 	if err != nil {
 		return nil, err
 	}

rpc/rpc.go

@@ -4,21 +4,37 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"net"
 	"os"
 	"testing"
 
 	"github.com/davrodpin/mole/rpc"
 )
 
+var (
+	addr net.Addr
+)
+
 func TestHandler(t *testing.T) {
+	responseTemplate := `{"message":"%s"}`
 	method := "test"
-	expectedResponse := `{"message":"test"}`
+	paramValue := "param"
+	expectedResponse := fmt.Sprintf(responseTemplate, paramValue)
 
 	rpc.Register(method, func(params interface{}) (json.RawMessage, error) {
-		return json.RawMessage(expectedResponse), nil
+		var r []uint8
+		var ok bool
+
+		if r, ok = params.([]uint8); !ok {
+			return nil, fmt.Errorf("invalid parameter")
+		}
+
+		m := fmt.Sprintf(responseTemplate, string(r[1:6]))
+
+		return json.RawMessage(m), nil
 	})
 
-	response, err := rpc.Call(context.Background(), method, "param")
+	response, err := rpc.Call(context.Background(), addr.String(), method, paramValue)
 	if err != nil {
 		t.Errorf("error while calling remote procedure: %v", err)
 	}
@@ -37,7 +53,7 @@ func TestMethodNotRegistered(t *testing.T) {
 	method := "methodnotregistered"
 	expectedResponse := fmt.Sprintf(`{"code":-32601,"data":null,"message":"method %s not found"}`, method)
 
-	response, err := rpc.Call(context.Background(), method, "param")
+	response, err := rpc.Call(context.Background(), addr.String(), method, "param")
 	if err != nil {
 		t.Errorf("error while calling remote procedure: %v", err)
 	}
@@ -60,7 +76,7 @@ func TestMethodWithError(t *testing.T) {
 		return nil, fmt.Errorf("error")
 	})
 
-	response, err := rpc.Call(context.Background(), method, "param")
+	response, err := rpc.Call(context.Background(), addr.String(), method, "param")
 	if err != nil {
 		t.Errorf("error while calling remote procedure: %v", err)
 	}
@@ -76,7 +92,9 @@ func TestMethodWithError(t *testing.T) {
 }
 
 func TestMain(m *testing.M) {
-	_, err := rpc.Start(rpc.DefaultAddress)
+	var err error
+
+	addr, err = rpc.Start(rpc.DefaultAddress)
 	if err != nil {
 		fmt.Printf("error initializing rpc server: %v", err)
 		os.Exit(1)