feat: Add connection timeout flag to swupdate command

cmd/ovp8xx/cmd/swupdate.go

@@ -34,20 +34,42 @@ func swupdateCommand(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("cannot get timeout: %w", err)
 	}
 
+	connectionTimeout, err := cmd.Flags().GetDuration("online")
+	if err != nil {
+		return fmt.Errorf("cannot get timeout: %w", err)
+	}
+
 	fmt.Printf("Updating firmware on %s:%d with file %s (%v)\n",
 		host,
 		port,
 		filepath.Base(filename),
 		timeout,
 	)
 
-	swu := swupdater.NewSWUpdater(host, port)
+	// notifications is a channel used to receive SWUpdaterNotification events.
+	// It has a buffer size of 10 to allow for asynchronous processing.
+	notifications := make(chan swupdater.SWUpdaterNotification, 10)
 
-	err = swu.Update(filename, timeout)
-	if err != nil {
+	// Print the messages as they come
+	go func() {
+		for n := range notifications {
+			if value, ok := n["swupdater"]; ok {
+				fmt.Println(value)
+			}
+			if value, ok := n["text"]; ok && n["type"] == "message" {
+				fmt.Println(value)
+			}
+		}
+	}()
+
+	// Create a new SWUpdater instance with the specified host, port, and notifications.
+	swu := swupdater.NewSWUpdater(host, port, notifications)
+	if err = swu.Update(filename,
+		connectionTimeout,
+		timeout,
+	); err != nil {
 		return fmt.Errorf("software update failed: %w", err)
 	}
-
 	return nil
 }
 
@@ -63,4 +85,5 @@ func init() {
 	swupdateCmd.Flags().String("file", "", "A file conatining the firmware image")
 	swupdateCmd.Flags().Uint16("port", 8080, "Port number for SWUpdate")
 	swupdateCmd.Flags().Duration("timeout", 5*time.Minute, "The timeout for the upload")
+	swupdateCmd.Flags().Duration("online", 2*time.Minute, "The time to wait for the device to become available")
 }

pkg/swupdater/swupdater.go

@@ -15,27 +15,32 @@ import (
 
 // SWUpdater represents a software updater.
 type SWUpdater struct {
-	hostName  string // The hostname of the updater.
-	port      uint16 // The port number of the updater.
-	urlUpload string // The URL for uploading software updates.
-	urlStatus string // The URL for checking the status of software updates.
+	hostName      string                     // The hostname of the updater.
+	port          uint16                     // The port number of the updater.
+	urlUpload     string                     // The URL for uploading software updates.
+	urlStatus     string                     // The URL for checking the status of software updates.
+	notifications chan SWUpdaterNotification // A channel for receiving notifications.
+	ws            *websocket.Conn
 }
 
+type SWUpdaterNotification map[string]string
+
 // NewSWUpdater creates a new instance of SWUpdater with the specified host name and port.
-func NewSWUpdater(hostName string, port uint16) *SWUpdater {
+func NewSWUpdater(hostName string, port uint16, notifications chan SWUpdaterNotification) *SWUpdater {
 	return &SWUpdater{
-		hostName:  hostName,
-		port:      port,
-		urlUpload: fmt.Sprintf("http://%s:%d/upload", hostName, port),
-		urlStatus: fmt.Sprintf("ws://%s:%d/ws", hostName, port),
+		hostName:      hostName,
+		port:          port,
+		urlUpload:     fmt.Sprintf("http://%s:%d/upload", hostName, port),
+		urlStatus:     fmt.Sprintf("ws://%s:%d/ws", hostName, port),
+		notifications: notifications,
 	}
 }
 
 // Upload performs the upload of the specified file.
 // The filename parameter specifies the name of the file to be uploaded.
 // Returns an error if the upload fails.
 func (s *SWUpdater) upload(filename string) error {
-	fmt.Printf("Uploading software image to %s\n", s.urlUpload)
+	s.statusUpdate(fmt.Sprintf("Uploading software image to %s\n", s.urlUpload))
 	const fieldname string = "file"
 
 	file, err := os.Open(filename)
@@ -86,31 +91,27 @@ func (s *SWUpdater) upload(filename string) error {
 //	  // SWUpdater process completed successfully
 //	}
 func (s *SWUpdater) waitForFinished(done chan error) {
-	c, _, err := websocket.DefaultDialer.Dial(s.urlStatus, nil)
-	if err != nil {
-		done <- fmt.Errorf("cannot connect to websocket: %w", err)
-		return
-	}
-	defer c.Close()
 
 	for {
-		_, message, err := c.ReadMessage()
+		_, message, err := s.ws.ReadMessage()
 		if err != nil {
 			done <- fmt.Errorf("cannot read message from websocket: %w", err)
 			return
 		}
 
-		data := make(map[string]string)
+		data := make(SWUpdaterNotification)
 		err = json.Unmarshal(message, &data)
 		if err != nil {
 			done <- fmt.Errorf("cannot unmarshal message: %w", err)
 			return
 		}
-		fmt.Println("Raw JSON: ", data)
+		// Send notification to channel
+		if s.notifications != nil {
+			s.notifications <- data
+		}
 		if data["type"] != "message" {
 			continue
 		}
-
 		if strings.Contains(data["text"], "SWUPDATE successful") {
 			done <- nil
 			return
@@ -122,11 +123,52 @@ func (s *SWUpdater) waitForFinished(done chan error) {
 	}
 }
 
+func (s *SWUpdater) connect() error {
+	var err error
+	s.ws, _, err = websocket.DefaultDialer.Dial(s.urlStatus, nil)
+	if err != nil {
+		return fmt.Errorf("unable to connect to the status socket: %w", err)
+	}
+	return err
+}
+
+func (s *SWUpdater) disconnect() {
+	s.ws.Close()
+}
+
+// statusUpdate updates the status of the SWUpdater.
+// It sends a notification to the channel with the provided status.
+func (s *SWUpdater) statusUpdate(status string) {
+	notification := make(SWUpdaterNotification)
+	notification["swupdater"] = status
+	// Send notification to channel
+	if s.notifications != nil {
+		s.notifications <- notification
+	}
+}
+
 // Update uploads a software image and waits for the update process to finish.
 // It takes a filename string and a timeout duration as parameters.
 // It returns an error if the upload fails, or if the operation times out.
-func (s *SWUpdater) Update(filename string, timeout time.Duration) error {
+func (s *SWUpdater) Update(filename string, connectionTimeout, timeout time.Duration) error {
 	done := make(chan error)
+	start := time.Now()
+	s.statusUpdate("Waiting for the Device to become ready...")
+	// Retry connection until successful or connectionTimeout occurs
+	for {
+		err := s.connect()
+		if err == nil {
+			s.statusUpdate("Device is ready now")
+			break
+		}
+		if time.Since(start) > connectionTimeout {
+			return fmt.Errorf("connection timeout: %w", err)
+		}
+		time.Sleep(3 * time.Second) // wait for a second before retrying
+	}
+	defer s.disconnect()
+
+	s.statusUpdate("Starting the Software Update process...")
 	go s.waitForFinished(done)
 	err := s.upload(filename)
 	if err != nil {
@@ -140,6 +182,6 @@ func (s *SWUpdater) Update(filename string, timeout time.Duration) error {
 		}
 		return nil
 	case <-time.After(timeout):
-		return errors.New("timeout")
+		return errors.New("a timeout occurred while waiting for the update to finish")
 	}
 }