Refactor ota package (#97)

Return a specific 'Econder' struct instead of a 'WriteCloser' interface. The 'Encoder' wraps a 'io.Writer' that's where the encoded file will be written to. In this way the consumer can decide when and how to open/close the writer.
Add tests.

Author: polldo

command/ota/generate.go

@@ -18,8 +18,8 @@
 package ota
 
 import (
-	"bytes"
 	"errors"
+	"fmt"
 	"io/ioutil"
 	"os"
 
@@ -51,17 +51,16 @@ func Generate(binFile string, outFile string, fqbn string) error {
 		return err
 	}
 
-	var w bytes.Buffer
-	otaWriter := inota.NewWriter(&w, arduinoVendorID, productID)
-	_, err = otaWriter.Write(data)
+	out, err := os.Create(outFile)
 	if err != nil {
 		return err
 	}
-	otaWriter.Close()
+	defer out.Close()
 
-	err = ioutil.WriteFile(outFile, w.Bytes(), os.FileMode(0644))
+	enc := inota.NewEncoder(out, arduinoVendorID, productID)
+	err = enc.Encode(data)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to encode binary file: %w", err)
 	}
 
 	return nil

command/ota/massupload.go

@@ -171,6 +171,7 @@ func run(uploader otaUploader, ids []string, otaFile string, expiration int) []R
 			results = append(results, r)
 			continue
 		}
+		defer file.Close()
 		jobs <- job{id: id, file: file}
 	}
 	close(jobs)

command/ota/upload.go

@@ -71,6 +71,7 @@ func Upload(params *UploadParams, cred *config.Credentials) error {
 	if err != nil {
 		return fmt.Errorf("%s: %w", "cannot open ota file", err)
 	}
+	defer file.Close()
 
 	expiration := otaExpirationMins
 	if params.Deferred {

internal/ota/encoder.go

@@ -18,60 +18,48 @@
 package ota
 
 import (
-	"bufio"
 	"encoding/binary"
+	"fmt"
 	"hash/crc32"
 	"io"
 	"strconv"
 
 	"github.com/arduino/arduino-cloud-cli/internal/lzss"
-	"github.com/juju/errors"
 )
 
-// A writer is a buffered, flushable writer.
-type writer interface {
-	io.Writer
-	Flush() error
-}
-
-// encoder encodes a binary into an .ota file.
-type encoder struct {
-	// w is the writer that compressed bytes are written to.
-	w writer
+// Encoder writes a binary to an output stream in the ota format.
+type Encoder struct {
+	// w is the stream where encoded bytes are written.
+	w io.Writer
 
-	// vendorID is the ID of the board vendor
+	// vendorID is the ID of the board vendor.
 	vendorID string
 
-	// is the ID of the board vendor is the ID of the board model
+	// productID is the ID of the board model.
 	productID string
 }
 
-// NewWriter creates a new `WriteCloser` for the the given VID/PID.
-func NewWriter(w io.Writer, vendorID, productID string) io.WriteCloser {
-	bw, ok := w.(writer)
-	if !ok {
-		bw = bufio.NewWriter(w)
-	}
-	return &encoder{
-		w:         bw,
+// NewEncoder creates a new ota encoder.
+func NewEncoder(w io.Writer, vendorID, productID string) *Encoder {
+	return &Encoder{
+		w:         w,
 		vendorID:  vendorID,
 		productID: productID,
 	}
 }
 
-// Write writes a compressed representation of p to e's underlying writer.
-func (e *encoder) Write(binaryData []byte) (int, error) {
-	//log.Println("original binaryData is", len(binaryData), "bytes length")
-
-	// Magic number (VID/PID)
+// Encode compresses data using a lzss algorithm, encodes the result
+// in ota format and writes it to e's underlying writer.
+func (e *Encoder) Encode(data []byte) error {
+	// Compute the magic number (VID/PID)
 	magicNumber := make([]byte, 4)
 	vid, err := strconv.ParseUint(e.vendorID, 16, 16)
 	if err != nil {
-		return 0, errors.Annotate(err, "OTA encoder: failed to parse vendorID")
+		return fmt.Errorf("cannot parse vendorID: %w", err)
 	}
 	pid, err := strconv.ParseUint(e.productID, 16, 16)
 	if err != nil {
-		return 0, errors.Annotate(err, "OTA encoder: failed to parse productID")
+		return fmt.Errorf("cannot parse productID: %w", err)
 	}
 
 	binary.LittleEndian.PutUint16(magicNumber[0:2], uint16(pid))
@@ -82,61 +70,43 @@ func (e *encoder) Write(binaryData []byte) (int, error) {
 		Compression: true,
 	}
 
-	// Compress the compiled binary
-	compressed := lzss.Encode(binaryData)
-
+	compressed := lzss.Encode(data)
 	// Prepend magic number and version field to payload
-	var binDataComplete []byte
-	binDataComplete = append(binDataComplete, magicNumber...)
-	binDataComplete = append(binDataComplete, version.AsBytes()...)
-	binDataComplete = append(binDataComplete, compressed...)
-	//log.Println("binDataComplete is", len(binDataComplete), "bytes length")
+	var outData []byte
+	outData = append(outData, magicNumber...)
+	outData = append(outData, version.Bytes()...)
+	outData = append(outData, compressed...)
 
-	headerSize, err := e.writeHeader(binDataComplete)
+	err = e.writeHeader(outData)
 	if err != nil {
-		return headerSize, err
+		return fmt.Errorf("cannot write data header to output stream: %w", err)
 	}
 
-	payloadSize, err := e.writePayload(binDataComplete)
+	_, err = e.w.Write(outData)
 	if err != nil {
-		return payloadSize, err
+		return fmt.Errorf("cannot write encoded data to output stream: %w", err)
 	}
 
-	return headerSize + payloadSize, nil
-}
-
-// Close closes the encoder, flushing any pending output. It does not close or
-// flush e's underlying writer.
-func (e *encoder) Close() error {
-	return e.w.Flush()
+	return nil
 }
 
-func (e *encoder) writeHeader(binDataComplete []byte) (int, error) {
-
+func (e *Encoder) writeHeader(data []byte) error {
 	// Write the length of the content
 	lengthAsBytes := make([]byte, 4)
-	binary.LittleEndian.PutUint32(lengthAsBytes, uint32(len(binDataComplete)))
-
-	n, err := e.w.Write(lengthAsBytes)
+	binary.LittleEndian.PutUint32(lengthAsBytes, uint32(len(data)))
+	_, err := e.w.Write(lengthAsBytes)
 	if err != nil {
-		return n, err
+		return err
 	}
 
-	// Calculate the checksum for binDataComplete
-	crc := crc32.ChecksumIEEE(binDataComplete)
-
-	// encode the checksum uint32 value as 4 bytes
+	// Write the checksum uint32 value as 4 bytes
+	crc := crc32.ChecksumIEEE(data)
 	crcAsBytes := make([]byte, 4)
 	binary.LittleEndian.PutUint32(crcAsBytes, crc)
-
-	n, err = e.w.Write(crcAsBytes)
+	_, err = e.w.Write(crcAsBytes)
 	if err != nil {
-		return n, err
+		return err
 	}
 
-	return len(lengthAsBytes) + len(crcAsBytes), nil
-}
-
-func (e *encoder) writePayload(data []byte) (int, error) {
-	return e.w.Write(data)
+	return nil
 }

internal/ota/encoder_test.go

@@ -20,7 +20,7 @@ package ota
 import (
 	"bytes"
 	"encoding/hex"
-	"log"
+	"io/ioutil"
 
 	"fmt"
 	"hash/crc32"
@@ -37,30 +37,25 @@ func TestComputeCrc32Checksum(t *testing.T) {
 	assert.Equal(t, crc, uint32(2090640218))
 }
 
-func TestEncoderWrite(t *testing.T) {
-
+func TestEncode(t *testing.T) {
 	// Setup test data
 	data, _ := hex.DecodeString("DEADBEEF") // uncompressed, or 'ef 6b 77 de f0' (compressed w/ LZSS)
 
 	var w bytes.Buffer
 	vendorID := "2341"  // Arduino
 	productID := "8054" // MRK Wifi 1010
 
-	otaWriter := NewWriter(&w, vendorID, productID)
-	defer otaWriter.Close()
+	enc := NewEncoder(&w, vendorID, productID)
 
-	n, err := otaWriter.Write(data)
+	err := enc.Encode(data)
 	if err != nil {
 		t.Error(err)
-		t.Fail()
 	}
-	log.Println("written ota of", n, "bytes length")
 
-	otaWriter.Close()
 	actual := w.Bytes()
 
-	// You can get the expected result creating an `.ota` file using Alex's tools:
-	// https://github.com/arduino-libraries/ArduinoIoTCloud/tree/master/extras/tools
+	// Expected result has been computed with the following tool:
+	// https://github.com/arduino-libraries/ArduinoIoTCloud/tree/master/extras/tools .
 	expected, _ := hex.DecodeString("11000000a1744bd4548041230000000000000040ef6b77def0")
 
 	res := bytes.Compare(expected, actual)
@@ -72,3 +67,51 @@ func TestEncoderWrite(t *testing.T) {
 
 	assert.Assert(t, res == 0) // 0 means equal
 }
+
+// Expected '.ota' files contained in testdata have been computed with the following tool:
+// https://github.com/arduino-libraries/ArduinoIoTCloud/tree/master/extras/tools .
+func TestEncodeFiles(t *testing.T) {
+	tests := []struct {
+		name    string
+		infile  string
+		outfile string
+	}{
+		{
+			name:    "blink",
+			infile:  "testdata/blink.bin",
+			outfile: "testdata/blink.ota",
+		},
+		{
+			name:    "cloud sketch",
+			infile:  "testdata/cloud.bin",
+			outfile: "testdata/cloud.ota",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			input, err := ioutil.ReadFile(tt.infile)
+			if err != nil {
+				t.Fatal("couldn't open test file")
+			}
+
+			want, err := ioutil.ReadFile(tt.outfile)
+			if err != nil {
+				t.Fatal("couldn't open test file")
+			}
+
+			var got bytes.Buffer
+			vendorID := "2341"  // Arduino
+			productID := "8057" // Nano 33 IoT
+			otaenc := NewEncoder(&got, vendorID, productID)
+			err = otaenc.Encode(input)
+			if err != nil {
+				t.Error(err)
+			}
+
+			if !bytes.Equal(want, got.Bytes()) {
+				t.Error("encoding failed")
+			}
+		})
+	}
+}

internal/ota/testdata/blink.bin

@@ -32,8 +32,8 @@ type Version struct {
 	PayloadBuildNum uint32
 }
 
-// AsBytes builds a 8 byte length representation of the Version Struct for the OTA update.
-func (v *Version) AsBytes() []byte {
+// Bytes builds a 8 byte length representation of the Version Struct for the OTA update.
+func (v *Version) Bytes() []byte {
 	version := []byte{0, 0, 0, 0, 0, 0, 0, 0}
 
 	// Set compression

internal/ota/testdata/blink.ota

@@ -34,7 +34,7 @@ func TestVersionWithCompressionEnabled(t *testing.T) {
 	}
 
 	expected := []byte{0, 0, 0, 0, 0, 0, 0, 0x40}
-	actual := version.AsBytes()
+	actual := version.Bytes()
 
 	// create a tabwriter for formatting the output
 	w := new(tabwriter.Writer)