Allow the server to execute commands

5 files changed, 163 insertions, 6 deletions

diff --git a/cmd/server/main.go b/cmd/server/main.go
index e326fa65c..18aca2f71 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -5,5 +5,5 @@ import (
 )
 
 func main() {
-	server.NewService().Serve("0.0.0.0:6666")
+	server.NewService().Serve("0.0.0.0:6666", server.CommandExecutorCallback)
 }
diff --git a/server/command_executor.go b/server/command_executor.go
new file mode 100644
index 000000000..46659fb43
--- /dev/null
+++ b/server/command_executor.go
@@ -0,0 +1,83 @@
+package server
+
+import (
+	"bytes"
+	"encoding/json"
+	"log"
+	"os/exec"
+	"syscall"
+)
+
+type cmdRequest struct {
+	Cmd []string `json:"cmd"`
+}
+
+type cmdResponse struct {
+	Status string `json:"status"`
+	Message string `json:"message"`
+}
+
+func CommandExecutorCallback(input []byte) []byte {
+	req := cmdRequest{}
+
+	err := json.Unmarshal(input, &req)
+	if err != nil {
+		return errorResponse("Error parsing JSON request")
+	}
+
+	output, err := runCommand(req.Cmd[0], req.Cmd[1:]...)
+
+	if err != nil {
+		return errorResponse(string(output.Bytes()))
+	}
+
+	return successResponse(string(output.Bytes()))
+}
+
+func runCommand(name string, args ...string) (bytes.Buffer, error) {
+	var stdoutBuf bytes.Buffer
+	var stderrBuf bytes.Buffer
+
+	cmd := makeCommand(name, args...)
+	cmd.Stdout = &stdoutBuf
+	cmd.Stderr = &stderrBuf
+
+	err := cmd.Run()
+	if err != nil {
+		return stderrBuf, err
+	}
+
+	return stdoutBuf, nil
+}
+
+// Based on git.gitCommand from gitlab-workhorse
+func makeCommand(name string, args ...string) *exec.Cmd {
+	cmd := exec.Command(name, args...)
+
+	// Start the command in its own process group (nice for signalling)
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+
+	return cmd
+}
+
+func errorResponse(message string) []byte {
+	return makeResponse("error", message)
+}
+
+func successResponse(message string) []byte {
+	return makeResponse("success", message)
+}
+
+func makeResponse(status string, message string) []byte {
+	res          := cmdResponse{status, message}
+	tempBuf, err := json.Marshal(res)
+
+	if err != nil {
+		log.Fatalln("Failed marshalling a JSON response")
+	}
+
+	buf := bytes.NewBuffer(tempBuf)
+	buf.WriteString("\n")
+
+	return buf.Bytes()
+}
diff --git a/server/command_executor_test.go b/server/command_executor_test.go
new file mode 100644
index 000000000..eeb37ce31
--- /dev/null
+++ b/server/command_executor_test.go
@@ -0,0 +1,67 @@
+package server
+
+import (
+	"bufio"
+	"bytes"
+	"os"
+	"net"
+	"testing"
+	"time"
+)
+
+const serviceAddress = "127.0.0.1:6667"
+
+func TestMain(m *testing.M) {
+	service := NewService()
+
+	go service.Serve(serviceAddress, CommandExecutorCallback)
+	defer service.Stop()
+
+	time.Sleep(10 * time.Millisecond)
+	os.Exit(m.Run())
+}
+
+func TestRunningCommandSuccessfully(t *testing.T) {
+	res := responseForCommand(`{"cmd":["ls", "-hal"]}`, t)
+
+	if !bytes.Contains(res, []byte(`{"status":"success","message":"total`)) {
+		t.Fatalf("Expected a successful response, got this response: %s", res)
+	}
+}
+
+func TestRunningCommandUnsuccessfully(t *testing.T) {
+	res := responseForCommand(`{"cmd":["ls", "/file-that-does-not-exist"]}`, t)
+
+	if !bytes.Contains(res, []byte(`{"status":"error","message":"ls: cannot access`)) {
+		t.Fatalf("Expected a failure response, got this response: %s", res)
+	}
+}
+
+func TestMalformedCommand(t *testing.T) {
+	res := responseForCommand(`{"cmd":["ls", "/file-that-does-not-exist"}`, t)
+
+	if !bytes.Equal(res, []byte(`{"status":"error","message":"Error parsing JSON request"}`)) {
+		t.Fatalf("Expected a failure response, got this response: %s", res)
+	}
+}
+
+func responseForCommand(cmd string, t *testing.T) []byte {
+	conn, err := net.Dial("tcp", serviceAddress)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer conn.Close()
+
+	if _, err := conn.Write([]byte(cmd + "\n")); err != nil {
+		t.Error(err)
+	}
+
+	reader := bufio.NewReader(conn)
+	buffer, err := reader.ReadBytes('\n')
+	if err != nil {
+		t.Error(err)
+	}
+
+	return bytes.TrimSpace(buffer)
+}
diff --git a/server/server.go b/server/server.go
index da864109f..72a06c3e5 100644
--- a/server/server.go
+++ b/server/server.go
@@ -14,6 +14,8 @@ type Service struct {
 	waitGroup *sync.WaitGroup
 }
 
+type Callback func([]byte) []byte
+
 func NewService() *Service {
 	service := &Service{
 		ch:        make(chan bool),
@@ -23,7 +25,7 @@ func NewService() *Service {
 	return service
 }
 
-func (s *Service) Serve(address string) {
+func (s *Service) Serve(address string, cb Callback) {
 	listener, err := newListener(address)
 	if err != nil {
 		log.Fatalln(err)
@@ -48,7 +50,7 @@ func (s *Service) Serve(address string) {
 		}
 		log.Println("Client connected from ", conn.RemoteAddr())
 		s.waitGroup.Add(1)
-		go s.serve(conn)
+		go s.serve(conn, cb)
 	}
 }
 
@@ -65,9 +67,10 @@ func (s *Service) Stop() {
 	s.waitGroup.Wait()
 }
 
-func (s *Service) serve(conn *net.TCPConn) {
+func (s *Service) serve(conn *net.TCPConn, cb Callback) {
 	defer conn.Close()
 	defer s.waitGroup.Done()
+
 	for {
 		select {
 		case <-s.ch:
@@ -75,7 +78,9 @@ func (s *Service) serve(conn *net.TCPConn) {
 			return
 		default:
 		}
+
 		conn.SetDeadline(time.Now().Add(1e9))
+
 		reader := bufio.NewReader(conn)
 		buffer, err := reader.ReadBytes('\n')
 		if err != nil {
@@ -88,7 +93,9 @@ func (s *Service) serve(conn *net.TCPConn) {
 			}
 			log.Println(err)
 		}
-		if _, err := conn.Write(buffer); nil != err {
+
+		ret := cb(buffer)
+		if _, err := conn.Write(ret); nil != err {
 			log.Println(err)
 			return
 		}
diff --git a/server/server_test.go b/server/server_test.go
index 73efb1a79..83f2f5453 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -12,7 +12,7 @@ func TestServerStandingUp(t *testing.T) {
 	service := NewService()
 	address := "127.0.0.1:6666"
 
-	go service.Serve(address)
+	go service.Serve(address, func(input []byte) []byte { return input })
 	defer service.Stop()
 
 	// Give service a little time to start listening for connections