Add distinct IP counter

common/ipsetsink/sink.go

@@ -0,0 +1,52 @@
+package ipsetsink
+
+import (
+	"crypto/hmac"
+	"hash"
+	"hash/crc64"
+
+	"github.com/clarkduvall/hyperloglog"
+	"golang.org/x/crypto/sha3"
+)
+
+func NewIPSetSink(maskingKey string) *IPSetSink {
+	countDistinct, _ := hyperloglog.NewPlus(18)
+	return &IPSetSink{
+		ipMaskingKey:  maskingKey,
+		countDistinct: countDistinct,
+	}
+}
+
+type IPSetSink struct {
+	ipMaskingKey  string
+	countDistinct *hyperloglog.HyperLogLogPlus
+}
+
+func (s *IPSetSink) maskIPAddress(ipAddress string) []byte {
+	hmacIPMasker := hmac.New(func() hash.Hash {
+		return sha3.New256()
+	}, []byte(s.ipMaskingKey))
+	hmacIPMasker.Write([]byte(ipAddress))
+	return hmacIPMasker.Sum(nil)
+}
+
+func (s *IPSetSink) AddIPToSet(ipAddress string) {
+	s.countDistinct.Add(crc64FromBytes{hashValue(s.maskIPAddress(ipAddress))})
+}
+
+func (s *IPSetSink) Dump() ([]byte, error) {
+	return s.countDistinct.GobEncode()
+}
+
+func (s *IPSetSink) Reset() {
+	s.countDistinct.Clear()
+}
+
+type hashValue []byte
+type crc64FromBytes struct {
+	hashValue
+}
+
+func (c crc64FromBytes) Sum64() uint64 {
+	return crc64.Checksum(c.hashValue, crc64.MakeTable(crc64.ECMA))
+}