Merge pull request #10 from VikingPingvin/agent_multi_input

Agent multi input
VikingPingvin · Dec 4, 2020 · 4f88ebb · 4f88ebb
2 parents f8bb246 + c87ae22
commit 4f88ebb
Show file tree

Hide file tree

Showing 8 changed files with 289 additions and 207 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,14 +9,17 @@
 
 - Basic Auth (HTPASSWD)
 - SQLITE database with ORM setup
-- Basic frontend with a tree-like view
-- Ability for Agent to send multiple files with a single run
+- Basic web frontend with a tree-like view
 - Ability for Agent to request artifact from the Server
 - Configurable Agent with YML
 - Configurable Server with YML
 - Parts of code is covered by tests
 ---
-# 0.1.0 (2020.12.02)
+# 0.2 (Unreleased)
+## Feature
+- Ability for Agent to send multiple files with a single configuration
+---
+# 0.1 (2020.12.03)
 ## FEATURE 
 
 - Send file over TCP /w Protobuf message to statically configured server

diff --git a/cmd/agent.go b/cmd/agent.go
@@ -43,7 +43,9 @@ func init() {
 	// Cobra supports local flags which will only run when this command
 	// is called directly, e.g.:
 	// agentCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
-	agentCmd.Flags().StringVar(&locker.InputArgPath, "file", "", "Specify path for input artifact file")
+	agentCmd.Flags().StringVar(&locker.InputArgPath, "file", "", "[path,...] Absolute or relative path(s). Multiple paths must be separated with ','")
 
-	agentCmd.Flags().StringVar(&locker.InputArgNamespace, "namespace", "", "Specify namespace in the following format: namespace/project/job-id")
+	agentCmd.Flags().StringVar(&locker.InputArgNamespace, "namespace", "", "[namespace/project/job-id] Separator must be '/'")
+
+	agentCmd.Flags().StringVar(&locker.InputArgConsume, "consume", "", "[namespace/project/job-id] Requests the specified artifact to download from the Locker Server.")
 }
diff --git a/cmd/root.go b/cmd/root.go
@@ -34,7 +34,7 @@ var rootCmd = &cobra.Command{
 	Long:  `A BLOB storage application`,
 	// Uncomment the following line if your bare application
 	// has an action associated with it:
-	Run: func(cmd *cobra.Command, args []string) {},
+	//Run: func(cmd *cobra.Command, args []string) {},
 }
 
 // Execute adds all child commands to the root command and sets flags appropriately.

diff --git a/locker/agent.go b/locker/agent.go
@@ -1,28 +1,19 @@
 package locker
 
 import (
-	"bufio"
-	"crypto/sha256"
-	"encoding/binary"
-	"errors"
-	"fmt"
-	"io"
 	"net"
-	"os"
-	"path/filepath"
-	"strings"
-
-	"vikingPingvin/locker/locker/messaging"
-	"vikingPingvin/locker/locker/messaging/protobuf"
+	"sync"
 
 	"github.com/rs/xid"
 	"github.com/rs/zerolog/log"
-	"google.golang.org/protobuf/proto"
 )
 
 // InputArgPath Relative or absolute path of files for Cobra CLI
-var InputArgPath string
-var InputArgNamespace string
+var (
+	InputArgPath      string
+	InputArgNamespace string
+	InputArgConsume   string
+)
 
 // InputData Populated at the start of the program
 type InputData struct {
@@ -44,204 +35,42 @@ type ArtifactAgent struct {
 	Port string
 }
 
-func (a ArtifactAgent) Start(inputData *InputData) bool {
-	connection, err := net.Dial("tcp", "localhost:27001")
-	if err != nil {
-		panic(err)
+func (a ArtifactAgent) Start(inputDataArray []*InputData) bool {
+	// TODO: move sendConcurrent to config file
+	const sendConcurrent = true
+
+	var wg sync.WaitGroup
+	for _, singleInputData := range inputDataArray {
+		if sendConcurrent {
+			wg.Add(1)
+			go sendArtifactToServer(singleInputData, &wg)
+		} else {
+			sendArtifactToServer(singleInputData, &wg)
+		}
 	}
-	defer connection.Close()
-	log.Info().Msg("Agent connected to Locker Server...")
-
-	// Send Metadata message
-	parseAndSendMetaData(connection, inputData)
-
-	// Send Payload message(s)
-	parseAndSendPayload(connection, inputData)
-
-	// Listen for ACK from server
-	listenForACK(connection, inputData)
 
+	wg.Wait()
 	return true
 }
 
-// If --file cli input is not null, parse file
-func parseAndSendMetaData(connection net.Conn, inputData *InputData) (fileInfo os.FileInfo, err error) {
-	fileInfo, err = os.Stat(inputData.FilePath)
-	if os.IsNotExist(err) {
-		log.Error().Msgf("Parsing file Input error: %v", err)
-		return fileInfo, err
-	}
-	inputData.FileHash = hashFile(inputData.FilePath)
-	inputData.FileName = fileInfo.Name()
-
-	log.Info().
-		Str("file name", fileInfo.Name()).
-		Str("Namespace", fmt.Sprintf("%s/%s/%s", inputData.NameSpace, inputData.Project, inputData.JobID)).
-		Str("size", fmt.Sprintf("%d", fileInfo.Size())).
-		Str("hash", fmt.Sprintf("%v", inputData.FileHash)).
-		Str("id", inputData.ID.String()).
-		Msg("Artifact metadata parsing finished")
-
-	message, err := messaging.CreateMessage_FileMeta(
-		inputData.ID.Bytes(),
-		protobuf.MessageType_META,
-		inputData.NameSpace,
-		inputData.Project,
-		inputData.JobID,
-		inputData.FileName,
-		inputData.FileHash)
+func sendArtifactToServer(artifact *InputData, wg *sync.WaitGroup) {
+	connection, err := net.Dial("tcp", "localhost:27001")
 	if err != nil {
 		panic(err)
 	}
+	defer connection.Close()
+	log.Info().Msg("Agent connected to Locker Server...")
 
 	// Send Metadata message
-	log.Info().Msg("Sending MetaData Packet")
-	messaging.SendProtoBufMessage(connection, message)
-
-	return fileInfo, err
-}
-
-//func parseAndSendPayload(bytes *[]byte, numBytes int) {
-func parseAndSendPayload(connection net.Conn, inputData *InputData) {
-
-	f, err := os.Open(inputData.FilePath)
-	defer f.Close()
-	if err != nil {
-		log.Error().Msgf("Cannot open file %s", inputData.FilePath)
-	}
-
-	log.Info().Msg("Started sending Payload Packets...")
-	reader := bufio.NewReader(f)
-	isPayloadFinal := false
-
-	buffer := make([]byte, 1024)
-	for {
-		n, ioErr := reader.Read(buffer)
-		if ioErr == io.EOF {
-			isPayloadFinal = true
-			// Send terminating payload protobuf message
-			terminalMessage, err := messaging.CreateMessage_FilePackage(
-				inputData.ID.Bytes(),
-				protobuf.MessageType_PACKAGE,
-				make([]byte, 1),
-				isPayloadFinal,
-			)
-			if err != nil {
-				log.Fatal().Msg("Fatal Error during payload protobuf assembly")
-			}
-			messaging.SendProtoBufMessage(connection, terminalMessage)
-			break
-		}
-
-		message, err := messaging.CreateMessage_FilePackage(
-			inputData.ID.Bytes(),
-			protobuf.MessageType_PACKAGE,
-			(buffer)[:n],
-			isPayloadFinal)
-
-		if err != nil {
-			log.Fatal().Msg("Fatal Error during payload protobuf assembly")
-		}
-
-		messaging.SendProtoBufMessage(connection, message)
-	}
-	log.Info().Msg("Finished sending Payload Packets...")
-}
-
-// Given a valid file path, returns a SHA256 hash
-func hashFile(path string) (hash []byte) {
-	f, err := os.Open(path)
-	defer f.Close()
-	if err != nil {
-		log.Err(err).Msgf("Cannot open file %s", path)
-	}
-
-	hasher := sha256.New()
-	if _, err := io.Copy(hasher, f); err != nil {
-		log.Err(err).Msg("Error calculating SHA256 Hash")
-	}
-	return hasher.Sum(nil)
-}
-
-func listenForACK(connection net.Conn, inputData *InputData) {
+	parseAndSendMetaData(connection, artifact)
 
-	// TODO: Make into const in message_handler (also server.go)
-	// sizePrefix is 4 bytes protobug message size
-	sizePrefix := make([]byte, 4)
-
-	_, _ = io.ReadFull(connection, sizePrefix)
-	protoLength := int(binary.BigEndian.Uint32(sizePrefix))
-
-	ackPacketRaw := make([]byte, protoLength)
-	_, _ = io.ReadFull(connection, ackPacketRaw)
-	genericProto := &protobuf.LockerMessage{}
-	if err := proto.Unmarshal(ackPacketRaw, genericProto); err != nil {
-		log.Err(err).Msg("Error during unmarshalling")
-	}
-
-	if genericProto.GetAck().ProtoReflect().IsValid() {
-		ackPacket := genericProto.GetAck()
-
-		serverResult := ackPacket.GetServerSuccess()
-		ackID, _ := xid.FromBytes(ackPacket.GetId())
-		if ackID != inputData.ID {
-			log.Warn().
-				Str("respone_id", ackID.String()).
-				Str("original_id", inputData.ID.String()).
-				Msg("Response ID mismatch.")
-		}
-
-		log.Info().
-			Str("id_back", ackID.String()).
-			Msgf("ACK packet recieved from server with success flag: %v", serverResult)
-	}
-}
-
-// Parse raw CLI input parameters to internal data structures
-func parseInputArguments() *InputData {
-	var err error
-	var inputPath string
-
-	if len(InputArgPath) == 0 {
-		err = errors.New("--file empty")
-		log.Err(err).Str("agent", "parseInputArguments").Msgf("No input file was given.")
-	}
-	if len(InputArgNamespace) == 0 {
-		err = errors.New("--namespace empty")
-		log.Err(err).Str("agent", "parseInputArguments").Msgf("No input namespace was given.")
-	}
-	if err != nil {
-		os.Exit(1)
-	}
-
-	inputPath = InputArgPath
-	if !filepath.IsAbs(InputArgPath) {
-		cwd, err := os.Getwd()
-		if err != nil {
-			log.Err(err).Msg("Error during CWD PATH parsing")
-			os.Exit(1)
-		}
-		log.Debug().Msgf("Relative path of input: %s", InputArgPath)
-		inputPath = filepath.Join(cwd, InputArgPath)
-
-	}
+	// Send Payload message(s)
+	parseAndSendPayload(connection, artifact)
 
-	fullNameSpace := InputArgNamespace
-	namePaths := strings.Split(fullNameSpace, "/")
-	if len(namePaths) != 3 {
-		err = errors.New("Namespace must contain 3 values separated by '/'")
-		log.Err(err).Msg("Namespace values not valid")
-		os.Exit(1)
-	}
+	// Listen for ACK from server
+	listenForACK(connection, artifact)
 
-	data := &InputData{
-		FilePath:  inputPath,
-		NameSpace: namePaths[0],
-		Project:   namePaths[1],
-		JobID:     namePaths[2],
-		ID:        xid.New(),
-	}
-	return data
+	wg.Done()
 }
 
 // ExecuteAgent : Entrypoint for Locker agent start