large packet based data transfer issue fixed

Author: mrasif

daemon/daemon.go

@@ -1,8 +1,11 @@
 package daemon
 
 import (
+	"bufio"
 	"context"
+	"encoding/binary"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"os"
@@ -164,7 +167,6 @@ func (ps *daemonMgr) StatusDaemon(ctx context.Context) error {
 	return nil
 }
 
-// Send commands to the daemon
 func (ps *daemonMgr) SendCommandToDaemon(ctx context.Context, packet models.Packet) error {
 	// Ensure the daemon is running
 	if !ps.IsDaemonRunning(ctx) {
@@ -178,48 +180,118 @@ func (ps *daemonMgr) SendCommandToDaemon(ctx context.Context, packet models.Pack
 	defer conn.Close()
 
 	data := packet.ToByte()
+	packetSize := uint32(len(data))
 
-	// Send request to daemon
-	_, err = conn.Write(data)
+	// Send the size header (4 bytes)
+	sizeHeader := make([]byte, 4)
+	binary.BigEndian.PutUint32(sizeHeader, packetSize)
+	_, err = conn.Write(sizeHeader)
 	if err != nil {
-		return fmt.Errorf("failed to send command: %v", err)
+		return fmt.Errorf("failed to send size header: %v", err)
 	}
 
-	// Read the response from the daemon
-	buf := make([]byte, 1024)
-	n, err := conn.Read(buf)
+	// Send the data in chunks
+	const chunkSize = 4096
+	for start := 0; start < len(data); start += chunkSize {
+		end := start + chunkSize
+		if end > len(data) {
+			end = len(data)
+		}
+		_, err = conn.Write(data[start:end])
+		if err != nil {
+			return fmt.Errorf("failed to send data chunk: %v", err)
+		}
+	}
+
+	// Read the response size header (4 bytes)
+	responseSizeHeader := make([]byte, 4)
+	_, err = io.ReadFull(conn, responseSizeHeader)
 	if err != nil {
-		return fmt.Errorf("failed to read response: %v", err)
+		return fmt.Errorf("failed to read response size header: %v", err)
 	}
-	result, err := utils.ToPacket(buf[:n])
+	responseSize := binary.BigEndian.Uint32(responseSizeHeader)
+
+	// Read the full response data in chunks
+	fullResponse := make([]byte, responseSize)
+	_, err = io.ReadFull(conn, fullResponse)
+	if err != nil {
+		return fmt.Errorf("failed to read full response: %v", err)
+	}
+
+	// Decode the response
+	result, err := utils.ToPacket(fullResponse)
 	if err != nil {
 		return fmt.Errorf("failed to decode response: %v", err)
 	}
-	fmt.Printf(string(result.Data))
+	fmt.Println(string(result.Data))
 	return nil
 }
 
 func (ps *daemonMgr) handleClient(ctx context.Context, conn net.Conn) {
 	defer conn.Close()
 
-	buf := make([]byte, 1024)
-	n, err := conn.Read(buf)
+	// Use a buffered reader for handling fragmented data
+	reader := bufio.NewReader(conn)
+
+	// Read the size header (4 bytes)
+	header := make([]byte, 4)
+	_, err := io.ReadFull(reader, header)
+	if err != nil {
+		log.Printf("Error reading size header: %v", err)
+		return
+	}
+
+	packetSize := binary.BigEndian.Uint32(header)
+	log.Printf("Expected packet size: %d", packetSize)
+
+	// Read the full packet based on the size
+	packetData := make([]byte, packetSize)
+	_, err = io.ReadFull(reader, packetData)
 	if err != nil {
-		log.Printf("Error reading command: %v", err)
+		log.Printf("Error reading packet data: %v", err)
 		return
 	}
 
-	packet, err := utils.ToPacket(buf[:n])
+	// Decode the packet
+	packet, err := utils.ToPacket(packetData)
 	if err != nil {
 		log.Printf("Error decoding packet: %v", err)
 		return
 	}
 
+	// Process the packet
 	result, err := ps.handler.Handle(ctx, packet)
 	if err != nil {
-		log.Printf("Error executing command: %v\n", err)
+		log.Printf("Error executing command: %v", err)
+		return
+	}
+
+	// Prepare the response
+	responseData := result.ToByte()
+	responseSize := uint32(len(responseData))
+
+	// Send the size header (4 bytes)
+	sizeHeader := make([]byte, 4)
+	binary.BigEndian.PutUint32(sizeHeader, responseSize)
+	_, err = conn.Write(sizeHeader)
+	if err != nil {
+		log.Printf("Error sending response size header: %v", err)
+		return
+	}
+
+	// Send the response in chunks
+	const chunkSize = 4096
+	for start := 0; start < len(responseData); start += chunkSize {
+		end := start + chunkSize
+		if end > len(responseData) {
+			end = len(responseData)
+		}
+		_, err = conn.Write(responseData[start:end])
+		if err != nil {
+			log.Printf("Error sending response data: %v", err)
+			return
+		}
 	}
-	conn.Write(result.ToByte())
 }
 
 func NewDaemonMgr(

handler/credientialHandler.go

@@ -44,12 +44,17 @@ func (h *handler) removeCrediential(packet models.Packet) (models.Packet, error)
 		result.Data = []byte("Uanble to decode crediential\n")
 		return result, err
 	}
+
+	if len(doc.Name) == 0 {
+		result.Data = []byte("Crediential name is required\n")
+		return result, nil
+	}
+
 	err = h.configRepo.RemoveCrediential(doc.Name)
 	if err != nil {
 		log.Println("failed to remove crediential.", err)
 	}
 	result.Data = []byte("crediential removed\n")
-	log.Println("Crediential removed.")
 	return result, nil
 }
 
@@ -78,7 +83,6 @@ func (h *handler) listCrediential(_ models.Packet) (models.Packet, error) {
 	table.Render()
 
 	result.Data = buf.Bytes()
-	log.Println("Crediential removed.")
 	return result, nil
 }
 
@@ -95,6 +99,11 @@ func (h *handler) detailCrediential(packet models.Packet) (models.Packet, error)
 		return result, err
 	}
 
+	if len(cred.Name) == 0 {
+		result.Data = []byte("Crediential name is required\n")
+		return result, nil
+	}
+
 	cred, err = h.configRepo.DetailCrediential(cred.Name)
 	if err != nil {
 		log.Println("failed to read crediential.", err)
@@ -106,8 +115,11 @@ func (h *handler) detailCrediential(packet models.Packet) (models.Packet, error)
 
 	table.Append([]string{"Name:", cred.Name})
 	table.Append([]string{"Type:", cred.Type})
-	table.Append([]string{"Password:", cred.Password})
-	table.Append([]string{"Key File:", cred.KeyFile})
+	if cred.Type == constants.CREDIENTIAL_PASS {
+		table.Append([]string{"Password:", cred.Password})
+	} else {
+		table.Append([]string{"Key File:", cred.KeyFile})
+	}
 
 	table.SetBorder(true) // Enable/Disable borders
 	table.SetAlignment(tablewriter.ALIGN_LEFT)

handler/hostHanler.go

@@ -45,6 +45,12 @@ func (h *handler) removeHost(packet models.Packet) (models.Packet, error) {
 		result.Data = []byte("Uanble to decode host\n")
 		return result, err
 	}
+
+	if len(doc.Name) == 0 {
+		result.Data = []byte("Host name is required\n")
+		return result, nil
+	}
+
 	err = h.configRepo.RemoveHost(doc.Name)
 	if err != nil {
 		log.Println("failed to remove host.", err)
@@ -79,7 +85,6 @@ func (h *handler) listHost(_ models.Packet) (models.Packet, error) {
 	table.Render()
 
 	result.Data = buf.Bytes()
-	log.Println("Crediential removed.")
 	return result, nil
 }
 
@@ -96,6 +101,11 @@ func (h *handler) detailHost(packet models.Packet) (models.Packet, error) {
 		return result, err
 	}
 
+	if len(doc.Name) == 0 {
+		result.Data = []byte("Host name is required\n")
+		return result, nil
+	}
+
 	doc, err = h.configRepo.DetailHost(doc.Name)
 	if err != nil {
 		log.Println("failed to read host.", err)

handler/tunnelHandler.go

@@ -46,6 +46,12 @@ func (h *handler) removeTunnel(packet models.Packet) (models.Packet, error) {
 		result.Data = []byte("Uanble to decode tunnel\n")
 		return result, err
 	}
+
+	if len(doc.Name) == 0 {
+		result.Data = []byte("Tunnel name is required\n")
+		return result, nil
+	}
+
 	err = h.configRepo.RemoveTunnel(doc.Name)
 	if err != nil {
 		log.Println("failed to remove tunnel.", err)
@@ -106,6 +112,11 @@ func (h *handler) detailTunnel(packet models.Packet) (models.Packet, error) {
 		return result, err
 	}
 
+	if len(doc.Name) == 0 {
+		result.Data = []byte("Tunnel name is required\n")
+		return result, nil
+	}
+
 	doc, err = h.configRepo.DetailTunnel(doc.Name)
 	if err != nil {
 		log.Println("failed to read tunnel.", err)