From 8f445949c356807d38bb3d67e932a6e16d794570 Mon Sep 17 00:00:00 2001 From: Andrea Waltlova Date: Mon, 17 Jul 2023 17:35:24 +0200 Subject: [PATCH] Refactor server.go file and add tests Signed-off-by: Andrea Waltlova --- src/runner.go | 9 +++--- src/server.go | 67 ++++++++++++++++++++++++++------------------- src/server_test.go | 68 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 112 insertions(+), 32 deletions(-) diff --git a/src/runner.go b/src/runner.go index 11c3bbd..e1553c7 100644 --- a/src/runner.go +++ b/src/runner.go @@ -71,8 +71,8 @@ func verifyYamlFile(yamlData []byte) bool { // If signature is valid then extracts the bash script to temporary file, // sets env variables if present and then runs the script. // Return stdout of executed script or error message if the signature wasn't valid. -func processSignedScript(yamlFileContet []byte) string { - signatureIsValid := verifyYamlFile(yamlFileContet) +func processSignedScript(incomingContent []byte) string { + signatureIsValid := verifyYamlFile(incomingContent) if !signatureIsValid { errorMsg := "Signature of yaml file is invalid" log.Errorln(errorMsg) @@ -82,9 +82,10 @@ func processSignedScript(yamlFileContet []byte) string { // Parse the YAML data into the yamlConfig struct var yamlContent signedYamlFile - err := yaml.Unmarshal(yamlFileContet, &yamlContent) + err := yaml.Unmarshal(incomingContent, &yamlContent) if err != nil { log.Errorln(err) + return "Yaml couldn't be unmarshaled" } // Set env variables @@ -123,8 +124,8 @@ func processSignedScript(yamlFileContet []byte) string { // Execute the script log.Infoln("Executing bash script") - out, err := exec.Command("/bin/sh", scriptFileName).Output() + if err != nil { log.Errorln("Failed to execute script: ", err) return "" diff --git a/src/server.go b/src/server.go index 017b1d9..4a26e43 100644 --- a/src/server.go +++ b/src/server.go @@ -40,6 +40,40 @@ func createDataMessage(commandOutput string, metadata map[string]string, directi return data } +// Processes signed script and sends message back to dispatcher +func processData(d *pb.Data) *pb.Data { + log.Infoln("Processing received yaml data") + commandOutput := processSignedScript(d.GetContent()) + + // Create a data message to send back to the dispatcher. + log.Infof("Creating payload for message %s", d.GetMessageId()) + data := createDataMessage(commandOutput, d.GetMetadata(), d.GetDirective(), d.GetMessageId()) + return data +} + +// Sends data back to dispatcher +func sendDataToDispatcher(data *pb.Data) *pb.Data { + // Dial the Dispatcher and call "Finish" + conn, err := grpc.Dial(yggdDispatchSocketAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + log.Error(err) + } + defer conn.Close() + + // Create a client of the Dispatch service + client := pb.NewDispatcherClient(conn) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + // Call "Send" + log.Infof("Sending response message to %s", data.GetResponseTo()) + log.Infoln("pb.Data message: ", data) + if _, err := client.Send(ctx, data); err != nil { + log.Error(err) + } + return data +} + // jobServer implements the Worker gRPC service as defined by the yggdrasil // gRPC protocol. It accepts Assignment messages, unmarshals the data into a // string, and echoes the content back to the Dispatch service by calling the @@ -59,35 +93,12 @@ type jobServer struct { // 4. Creates a client of the Dispatcher service. // 5. Constructs a data message to send back to the dispatcher. // 6. Sends the data message using the "Send" method of the Dispatcher service. -func (s *jobServer) Send(ctx context.Context, d *pb.Data) (*pb.Receipt, error) { - go func() { - log.Infoln("Processing received yaml data") - commandOutput := processSignedScript(d.GetContent()) - - // Dial the Dispatcher and call "Finish" - conn, err := grpc.Dial( - yggdDispatchSocketAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - log.Error(err) - } - defer conn.Close() +func (s *jobServer) Send(_ context.Context, d *pb.Data) (*pb.Receipt, error) { - // Create a client of the Dispatch service - c := pb.NewDispatcherClient(conn) - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - // Create a data message to send back to the dispatcher. - log.Infof("Creating payload for message %s", d.GetMessageId()) - data := createDataMessage( - commandOutput, d.GetMetadata(), d.GetDirective(), d.GetMessageId()) - - // Call "Send" - log.Infof("Sending message to %s", d.GetMessageId()) - log.Infoln("pb.Data message: ", data) - if _, err := c.Send(ctx, data); err != nil { - log.Error(err) - } + // Goroutine processing the data, cancels the context when processing is done + go func() { + data := processData(d) + sendDataToDispatcher(data) }() // Respond to the start request that the work was accepted. diff --git a/src/server_test.go b/src/server_test.go index 18fb51b..a21a8a1 100644 --- a/src/server_test.go +++ b/src/server_test.go @@ -3,6 +3,9 @@ package main import ( "strings" "testing" + + "github.com/google/uuid" + pb "github.com/redhatinsights/yggdrasil/protocol" ) func TestCreateDataMessage(t *testing.T) { @@ -47,3 +50,68 @@ func TestCreateDataMessage(t *testing.T) { t.Errorf("Expected Content to be empty, but got %s", data.Content) } } + +func TestProcessData(t *testing.T) { + // FIXME: this should ideally test that all correct functions are called + // Probably easiest would be to move them to interface and then make the function argument take mock interface + yggdDispatchSocketAddr = "mock-target" + + shouldVerifyYaml := false + shouldDoInsightsCoreGPGCheck := false + temporaryWorkerDirectory := "test-dir" + config = &Config{ + VerifyYAML: &shouldVerifyYaml, + TemporaryWorkerDirectory: &temporaryWorkerDirectory, + InsightsCoreGPGCheck: &shouldDoInsightsCoreGPGCheck, + } + + yamlData := []byte(` +vars: + _insights_signature: "invalid-signature" + _insights_signature_exclude: "/vars/insights_signature,/vars/content_vars" + content: | + #!/bin/sh + echo "$RHC_WORKER_FOO $RHC_WORKER_BAR!" + content_vars: + FOO: Hello + BAR: World`) + + returnURL := "bar" + testData := &pb.Data{ + Content: yamlData, + Metadata: map[string]string{ + "return_content_type": "foo", + "return_url": returnURL, + "correlation_id": "000", + }, + Directive: "Your directive", + MessageId: "Your message ID", + } + + data := processData(testData) + expectedOutput := "Hello World!" + + if !strings.Contains(string(data.GetContent()), expectedOutput) { + t.Errorf("Expected content to contain '%s', but it didn't", expectedOutput) + } + + if data.GetDirective() != returnURL { + t.Errorf("Expected directive to contain '%s', but it didn't", returnURL) + } +} + +func TestSendDataToDispatcher(t *testing.T) { + // Tests only that the function doesn't modify data sent to dispatcher + + yggdDispatchSocketAddr = "mock-target" + + testData := &pb.Data{ + MessageId: uuid.New().String(), + ResponseTo: "mock-id", + } + + data := sendDataToDispatcher(testData) + if data != testData { + t.Errorf("Function should NOT change data before sent, but it did: %s", data) + } +}