From ce6b08761ed3970ba46e82b632816ba03776314f Mon Sep 17 00:00:00 2001
From: Ethan Frey <ethanfrey@users.noreply.github.com>
Date: Tue, 20 Jun 2017 15:56:12 +0200
Subject: [PATCH] Improve crc32

---
 keys/ecc.go      | 35 ++++++++++++++++++++++++++---------
 keys/ecc_test.go |  5 +++--
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/keys/ecc.go b/keys/ecc.go
index ff16345d3..ba754d525 100644
--- a/keys/ecc.go
+++ b/keys/ecc.go
@@ -26,14 +26,19 @@ func (_ NoECC) AddECC(input []byte) []byte            { return input }
 func (_ NoECC) CheckECC(input []byte) ([]byte, error) { return input, nil }
 
 // CRC32 does the ieee crc32 polynomial check
-type CRC32 struct{}
+type CRC32 struct {
+	Poly  uint32
+	table *crc32.Table
+}
+
+var _ ECC = &CRC32{}
 
-var _ ECC = CRC32{}
+func (c *CRC32) AddECC(input []byte) []byte {
+	table := c.getTable()
 
-func (_ CRC32) AddECC(input []byte) []byte {
 	// get crc and convert to some bytes...
-	crc := crc32.ChecksumIEEE(input)
-	check := make([]byte, 4)
+	crc := crc32.Checksum(input, table)
+	check := make([]byte, crc32.Size)
 	binary.BigEndian.PutUint32(check, crc)
 
 	// append it to the input
@@ -41,16 +46,28 @@ func (_ CRC32) AddECC(input []byte) []byte {
 	return output
 }
 
-func (_ CRC32) CheckECC(input []byte) ([]byte, error) {
-	if len(input) <= 4 {
+func (c *CRC32) CheckECC(input []byte) ([]byte, error) {
+	table := c.getTable()
+
+	if len(input) <= crc32.Size {
 		return nil, errors.New("input too short, no checksum present")
 	}
-	cut := len(input) - 4
+	cut := len(input) - crc32.Size
 	data, check := input[:cut], input[cut:]
 	crc := binary.BigEndian.Uint32(check)
-	calc := crc32.ChecksumIEEE(data)
+	calc := crc32.Checksum(data, table)
 	if crc != calc {
 		return nil, errors.New("Checksum does not match")
 	}
 	return data, nil
 }
+
+func (c *CRC32) getTable() *crc32.Table {
+	if c.table == nil {
+		if c.Poly == 0 {
+			c.Poly = crc32.IEEE
+		}
+		c.table = crc32.MakeTable(c.Poly)
+	}
+	return c.table
+}
diff --git a/keys/ecc_test.go b/keys/ecc_test.go
index f58bc7a9c..a85d4ddd7 100644
--- a/keys/ecc_test.go
+++ b/keys/ecc_test.go
@@ -1,6 +1,7 @@
 package keys
 
 import (
+	"hash/crc32"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -12,7 +13,7 @@ import (
 func TestECCPasses(t *testing.T) {
 	assert := assert.New(t)
 
-	checks := []ECC{NoECC{}, CRC32{}}
+	checks := []ECC{NoECC{}, &CRC32{}, &CRC32{Poly: crc32.Castagnoli}}
 
 	for _, check := range checks {
 		for i := 0; i < 2000; i++ {
@@ -32,7 +33,7 @@ func TestECCPasses(t *testing.T) {
 func TestECCFails(t *testing.T) {
 	assert := assert.New(t)
 
-	checks := []ECC{CRC32{}}
+	checks := []ECC{&CRC32{}, &CRC32{Poly: crc32.Castagnoli}}
 
 	attempts := 2000