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