github.com/simpleiot/simpleiot@v0.18.3/client/file-transfer.go (about)

     1  package client
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"io"
     7  	"log"
     8  	"os"
     9  	"path"
    10  	"time"
    11  
    12  	"github.com/nats-io/nats.go"
    13  	"github.com/simpleiot/simpleiot/internal/pb"
    14  	"google.golang.org/protobuf/proto"
    15  )
    16  
    17  // Companion file in api/send-file.go
    18  
    19  type fileDownload struct {
    20  	name string
    21  	data []byte
    22  	seq  int32
    23  }
    24  
    25  // ListenForFile listens for a file sent from server. dir is the directly to place
    26  // downloaded files.
    27  func ListenForFile(nc *nats.Conn, dir, deviceID string, callback func(path string)) error {
    28  	dl := fileDownload{}
    29  	_, err := nc.Subscribe(fmt.Sprintf("device.%v.file", deviceID), func(m *nats.Msg) {
    30  		chunk := &pb.FileChunk{}
    31  
    32  		err := proto.Unmarshal(m.Data, chunk)
    33  
    34  		if err != nil {
    35  			log.Println("Error decoding file chunk:", err)
    36  			err := nc.Publish(m.Reply, []byte("error decoding"))
    37  			if err != nil {
    38  				log.Println("Error replying to file download:", err)
    39  				return
    40  			}
    41  		}
    42  
    43  		if chunk.Seq == 0 {
    44  			// we are starting a new stream
    45  			dl.name = chunk.FileName
    46  			dl.data = []byte{}
    47  			dl.seq = 0
    48  		} else if chunk.Seq != dl.seq+1 {
    49  			log.Println("Seq # error in file download:", dl.seq, chunk.Seq)
    50  			err := nc.Publish(m.Reply, []byte("seq error"))
    51  			if err != nil {
    52  				log.Println("Error replying to file download:", err)
    53  				return
    54  			}
    55  		}
    56  
    57  		// process data from server
    58  		dl.data = append(dl.data, chunk.Data...)
    59  		dl.seq = chunk.Seq
    60  
    61  		switch chunk.State {
    62  		case pb.FileChunk_ERROR:
    63  			log.Println("Server error getting chunk")
    64  			// reset download
    65  			dl = fileDownload{}
    66  		case pb.FileChunk_DONE:
    67  			filePath := path.Join(dir, dl.name)
    68  			err := os.WriteFile(filePath, dl.data, 0644)
    69  			if err != nil {
    70  				log.Println("Error writing dl file:", err)
    71  				err := nc.Publish(m.Reply, []byte("error writing"))
    72  				if err != nil {
    73  					log.Println("Error replying to file download:", err)
    74  					return
    75  				}
    76  			}
    77  
    78  			callback(filePath)
    79  		}
    80  
    81  		err = nc.Publish(m.Reply, []byte("OK"))
    82  		if err != nil {
    83  			log.Println("Error replying to file download:", err)
    84  		}
    85  	})
    86  
    87  	return err
    88  }
    89  
    90  // SendFile can be used to send a file to a device. Callback provides bytes transferred.
    91  func SendFile(nc *nats.Conn, deviceID string, reader io.Reader, name string, callback func(int)) error {
    92  	done := false
    93  	seq := int32(0)
    94  
    95  	bytesTx := 0
    96  
    97  	// send file in chunks
    98  	for {
    99  		var err error
   100  		data := make([]byte, 50*1024)
   101  		count, err := reader.Read(data)
   102  		data = data[:count]
   103  
   104  		chunk := &pb.FileChunk{Seq: seq, Data: data}
   105  
   106  		if seq == 0 {
   107  			chunk.FileName = name
   108  		}
   109  
   110  		if err != nil {
   111  			if err != io.EOF {
   112  				chunk.State = pb.FileChunk_ERROR
   113  			} else {
   114  				chunk.State = pb.FileChunk_DONE
   115  			}
   116  			done = true
   117  		}
   118  
   119  		out, err := proto.Marshal(chunk)
   120  
   121  		if err != nil {
   122  			return err
   123  		}
   124  
   125  		subject := fmt.Sprintf("device.%v.file", deviceID)
   126  
   127  		retry := 0
   128  		for ; retry < 3; retry++ {
   129  			msg, err := nc.Request(subject, out, time.Minute)
   130  
   131  			if err != nil {
   132  				log.Println("Error sending file, retrying:", retry, err)
   133  				continue
   134  			}
   135  
   136  			msgS := string(msg.Data)
   137  
   138  			if msgS != "OK" {
   139  				log.Println("Error from device when sending file:", retry, msgS)
   140  				continue
   141  			}
   142  
   143  			// we must have sent OK, break out of loop
   144  			break
   145  		}
   146  
   147  		if retry >= 3 {
   148  			return errors.New("Error sending file to device")
   149  		}
   150  
   151  		bytesTx += count
   152  		callback(bytesTx)
   153  
   154  		if done {
   155  			break
   156  		}
   157  
   158  		seq++
   159  	}
   160  
   161  	return nil
   162  }