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