Remove duplicated code send message

r4ulcl · Dec 15, 2024 · 42ed49e · 42ed49e
1 parent 86c6b9f
commit 42ed49e
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 73 deletions.
diff --git 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
@@ -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
@@ -72,95 +72,63 @@ 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,
 		UP:           true,
 		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)
 }