fec.go raw

   1  // Package fec implements Reed Solomon 9/3 forward error correction,
   2  //  intended to be sent as 9 pieces where 3 uncorrupted parts allows assembly of the message
   3  package fec
   4  
   5  import (
   6  	"encoding/binary"
   7  	"fmt"
   8  	"hash/crc32"
   9  	
  10  	"github.com/vivint/infectious"
  11  )
  12  
  13  var (
  14  	rsTotal    = 9
  15  	rsRequired = 3
  16  	rsFEC      = func() *infectious.FEC {
  17  		var e error
  18  		var c *infectious.FEC
  19  		if c, e = infectious.NewFEC(rsRequired, rsTotal); !E.Chk(e) {
  20  			return c
  21  		}
  22  		panic(e)
  23  	}()
  24  )
  25  
  26  // padData appends a 2 byte length prefix, and pads to a multiple of rsTotal.
  27  // Max message size is limited to 1<<32 but in our use will never get near
  28  // this size through higher level protocols breaking packets into sessions
  29  func padData(data []byte) (out []byte) {
  30  	dataLen := len(data)
  31  	prefixBytes := make([]byte, 4)
  32  	binary.LittleEndian.PutUint32(prefixBytes, uint32(dataLen))
  33  	data = append(prefixBytes, data...)
  34  	dataLen = len(data)
  35  	chunkLen := (dataLen) / rsTotal
  36  	chunkMod := (dataLen) % rsTotal
  37  	if chunkMod != 0 {
  38  		chunkLen++
  39  	}
  40  	padLen := rsTotal*chunkLen - dataLen
  41  	out = append(data, make([]byte, padLen)...)
  42  	return
  43  }
  44  
  45  // Encode turns a byte slice into a set of shards with first byte containing the shard number.
  46  func Encode(data []byte) (chunks [][]byte, e error) {
  47  	// First we must pad the data
  48  	data = padData(data)
  49  	shares := make([]infectious.Share, rsTotal)
  50  	output := func(s infectious.Share) {
  51  		shares[s.Number] = s.DeepCopy()
  52  	}
  53  	e = rsFEC.Encode(data, output)
  54  	if e != nil {
  55  		E.Ln(e)
  56  		return
  57  	}
  58  	for i := range shares {
  59  		// Append the chunk number to the front of the chunk
  60  		chunk := append([]byte{byte(shares[i].Number)}, shares[i].Data...)
  61  		// Checksum includes chunk number byte so we know if its checksum is incorrect so could the chunk number be
  62  		checksum := crc32.Checksum(chunk, crc32.MakeTable(crc32.Castagnoli))
  63  		checkBytes := make([]byte, 4)
  64  		binary.LittleEndian.PutUint32(checkBytes, checksum)
  65  		chunk = append(chunk, checkBytes...)
  66  		chunks = append(chunks, chunk)
  67  	}
  68  	return
  69  }
  70  
  71  // Decode takes a set of shards and if there is sufficient to reassemble,
  72  // returns the corrected data
  73  func Decode(chunks [][]byte) (data []byte, e error) {
  74  	var shares []infectious.Share
  75  	for i := range chunks {
  76  		bodyLen := len(chunks[i])
  77  		// log.SPEW(chunks[i])
  78  		body := chunks[i][:bodyLen-4]
  79  		checksum := crc32.Checksum(body, crc32.MakeTable(crc32.Castagnoli))
  80  		var share infectious.Share
  81  		if binary.LittleEndian.Uint32(chunks[i][bodyLen-4:]) == checksum {
  82  			share = infectious.Share{
  83  				Number: int(body[0]),
  84  				Data:   body[1:],
  85  			}
  86  		} else {
  87  			share = infectious.Share{
  88  				Number: i,
  89  				Data:   nil,
  90  			}
  91  		}
  92  		shares = append(shares, share)
  93  	}
  94  	data, e = rsFEC.Decode(nil, shares)
  95  	if len(data) > 4 {
  96  		prefix := data[:4]
  97  		data = data[4:]
  98  		dataLen := int(binary.LittleEndian.Uint32(prefix))
  99  		if len(data) < dataLen {
 100  			I.S(data)
 101  			e = fmt.Errorf("somehow data is corrupted though everything else about it is correct")
 102  		} else {
 103  			data = data[:dataLen]
 104  		}
 105  	}
 106  	return
 107  }
 108