From 42ed49e0b48c750dbae0df8ddffbab92d16ce543 Mon Sep 17 00:00:00 2001 From: r4ulcl <25440046+r4ulcl@users.noreply.github.com> Date: Sun, 15 Dec 2024 11:41:13 +0100 Subject: [PATCH] Remove duplicated code send message --- manager/api/API.go | 12 +++ manager/api/APItask.go | 10 +-- worker/managerrequest/managerRequest.go | 100 ++++++++---------------- 3 files changed, 49 insertions(+), 73 deletions(-) diff --git a/manager/api/API.go b/manager/api/API.go index 6a50347..a3fa778 100644 --- a/manager/api/API.go +++ b/manager/api/API.go @@ -103,3 +103,15 @@ func handleEntityStatus[T any](w http.ResponseWriter, r *http.Request, db *sql.D http.Error(w, "{ \"error\" : \"Invalid "+entityName+" encode body: "+err.Error()+"\"}", http.StatusBadRequest) } } + +func getUsername(r *http.Request, verbose, debug bool) (bool, string) { + username, ok := r.Context().Value(utils.UsernameKey).(string) + if debug { + log.Println("getUsername", username) + } + if !ok && (debug || verbose) { + log.Println("API { \"error\" : \"Unauthorized\" }") + } + + return ok, username +} diff --git a/manager/api/APItask.go b/manager/api/APItask.go index 88b76f6..a7c9a47 100644 --- a/manager/api/APItask.go +++ b/manager/api/APItask.go @@ -42,9 +42,8 @@ import ( // @security ApiKeyAuth // @router /task [get] func HandleTaskGet(w http.ResponseWriter, r *http.Request, db *sql.DB, verbose, debug bool) { - _, ok := r.Context().Value(utils.UsernameKey).(string) + ok, _ := getUsername(r, verbose, debug) if !ok { - // if not username is a worker http.Error(w, "{ \"error\" : \"Unauthorized\" }", http.StatusUnauthorized) return } @@ -95,11 +94,8 @@ func HandleTaskGet(w http.ResponseWriter, r *http.Request, db *sql.DB, verbose, // @security ApiKeyAuth // @router /task [post] func HandleTaskPost(w http.ResponseWriter, r *http.Request, db *sql.DB, verbose, debug bool, wg *sync.WaitGroup) { - username, okUser := r.Context().Value(utils.UsernameKey).(string) - if !okUser { - if debug { - log.Println("API { \"error\" : \"Unauthorized\" }") - } + ok, username := getUsername(r, verbose, debug) + if !ok { http.Error(w, "{ \"error\" : \"Unauthorized\" }", http.StatusUnauthorized) return } diff --git a/worker/managerrequest/managerRequest.go b/worker/managerrequest/managerRequest.go index fe9b88e..5d78623 100644 --- a/worker/managerrequest/managerRequest.go +++ b/worker/managerrequest/managerRequest.go @@ -72,42 +72,52 @@ func SendMessage(conn *websocket.Conn, message []byte, verbose, debug bool, writ return nil } -// AddWorker sends a POST request to add a worker to the manager -func AddWorker(config *utils.WorkerConfig, verbose, debug bool, writeLock *sync.Mutex) error { - // Create a Worker object with the provided configuration - worker := globalstructs.Worker{ - Name: config.Name, - DefaultThreads: config.DefaultThreads, - IddleThreads: config.DefaultThreads, - UP: true, - DownCount: 0, +// SendWebSocketMessage is a helper function to send a WebSocket message to the manager +func SendWebSocketMessage(config *utils.WorkerConfig, messageType string, payload interface{}, verbose, debug bool, writeLock *sync.Mutex) error { + // Marshal the payload into JSON + payloadData, err := json.Marshal(payload) + if err != nil { + log.Println("Error encoding JSON payload:", err) + return err } - // Marshal the worker object into JSON - payload, _ := json.Marshal(worker) - + // Create the WebSocket message msg := globalstructs.WebsocketMessage{ - Type: "addWorker", - JSON: string(payload), + Type: messageType, + JSON: string(payloadData), } + // Debug logging + if debug { + log.Printf("ManagerRequest msg (%s): %v", messageType, msg) + } + + // Marshal the WebSocket message into JSON jsonData, err := json.Marshal(msg) if err != nil { - log.Println("Error encoding JSON:", err) + log.Println("Error encoding WebSocket message:", err) return err } - err = SendMessage(config.Conn, jsonData, verbose, debug, writeLock) - if err != nil { - return err + // Send the message + return SendMessage(config.Conn, jsonData, verbose, debug, writeLock) +} + +// AddWorker sends a POST request to add a worker to the manager +func AddWorker(config *utils.WorkerConfig, verbose, debug bool, writeLock *sync.Mutex) error { + worker := globalstructs.Worker{ + Name: config.Name, + DefaultThreads: config.DefaultThreads, + IddleThreads: config.DefaultThreads, + UP: true, + DownCount: 0, } - return nil + return SendWebSocketMessage(config, "addWorker", worker, verbose, debug, writeLock) } -// DeleteWorker sends a POST request to delete a worker to the manager +// DeleteWorker sends a POST request to delete a worker from the manager func DeleteWorker(config *utils.WorkerConfig, verbose, debug bool, writeLock *sync.Mutex) error { - // Create a Worker object with the provided configuration worker := globalstructs.Worker{ Name: config.Name, IddleThreads: -1, @@ -115,52 +125,10 @@ func DeleteWorker(config *utils.WorkerConfig, verbose, debug bool, writeLock *sy DownCount: 0, } - // Marshal the worker object into JSON - payload, _ := json.Marshal(worker) - - msg := globalstructs.WebsocketMessage{ - Type: "deleteWorker", - JSON: string(payload), - } - - jsonData, err := json.Marshal(msg) - if err != nil { - log.Println("Error encoding JSON:", err) - return err - } - - err = SendMessage(config.Conn, jsonData, verbose, debug, writeLock) - if err != nil { - return err - } - - return nil + return SendWebSocketMessage(config, "deleteWorker", worker, verbose, debug, writeLock) } -// CallbackTaskMessage sends a POST request to the manager to callback with a task message +// CallbackTaskMessage sends a POST request to the manager with a task message func CallbackTaskMessage(config *utils.WorkerConfig, task *globalstructs.Task, verbose, debug bool, writeLock *sync.Mutex) error { - // Marshal the task object into JSON - payload, _ := json.Marshal(task) - - msg := globalstructs.WebsocketMessage{ - Type: "callbackTask", - JSON: string(payload), - } - - if debug { - log.Println("ManagerRequest msg callback:", msg) - } - - jsonData, err := json.Marshal(msg) - if err != nil { - log.Println("Error encoding JSON:", err) - return err - } - - err = SendMessage(config.Conn, jsonData, verbose, debug, writeLock) - if err != nil { - return err - } - - return nil + return SendWebSocketMessage(config, "callbackTask", task, verbose, debug, writeLock) }