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))
+}

common/ipsetsink/sink_test.go

@@ -0,0 +1,47 @@
+package ipsetsink
+
+import (
+	"fmt"
+	"github.com/clarkduvall/hyperloglog"
+	"testing"
+)
+import . "github.com/smartystreets/goconvey/convey"
+
+func TestSinkInit(t *testing.T) {
+	Convey("Context", t, func() {
+		sink := NewIPSetSink("demo")
+		sink.AddIPToSet("test1")
+		sink.AddIPToSet("test2")
+		data, err := sink.Dump()
+		So(err, ShouldBeNil)
+		structure, err := hyperloglog.NewPlus(18)
+		So(err, ShouldBeNil)
+		err = structure.GobDecode(data)
+		So(err, ShouldBeNil)
+		count := structure.Count()
+		So(count, ShouldBeBetweenOrEqual, 1, 3)
+	})
+}
+
+func TestSinkCounting(t *testing.T) {
+	Convey("Context", t, func() {
+		for itemCount := 300; itemCount <= 10000; itemCount += 200 {
+			sink := NewIPSetSink("demo")
+			for i := 0; i <= itemCount; i++ {
+				sink.AddIPToSet(fmt.Sprintf("demo%v", i))
+			}
+			for i := 0; i <= itemCount; i++ {
+				sink.AddIPToSet(fmt.Sprintf("demo%v", i))
+			}
+			data, err := sink.Dump()
+			So(err, ShouldBeNil)
+			structure, err := hyperloglog.NewPlus(18)
+			So(err, ShouldBeNil)
+			err = structure.GobDecode(data)
+			So(err, ShouldBeNil)
+			count := structure.Count()
+			So((float64(count)/float64(itemCount))-1.0, ShouldAlmostEqual, 0, 0.01)
+		}
+
+	})
+}

common/ipsetsink/sinkcluster/common.go

@@ -0,0 +1,9 @@
+package sinkcluster
+
+import "time"
+
+type SinkEntry struct {
+	RecordingStart time.Time `json:"recordingStart"`
+	RecordingEnd   time.Time `json:"recordingEnd"`
+	Recorded       []byte    `json:"recorded"`
+}

common/ipsetsink/sinkcluster/reader.go

@@ -0,0 +1,60 @@
+package sinkcluster
+
+import (
+	"bufio"
+	"encoding/json"
+	"github.com/clarkduvall/hyperloglog"
+	"io"
+	"time"
+)
+
+func NewClusterCounter(from time.Time, to time.Time) *ClusterCounter {
+	return &ClusterCounter{from: from, to: to}
+}
+
+type ClusterCounter struct {
+	from time.Time
+	to   time.Time
+}
+
+type ClusterCountResult struct {
+	Sum           uint64
+	ChunkIncluded int64
+}
+
+func (c ClusterCounter) Count(reader io.Reader) (*ClusterCountResult, error) {
+	result := ClusterCountResult{}
+	counter, err := hyperloglog.NewPlus(18)
+	if err != nil {
+		return nil, err
+	}
+	inputScanner := bufio.NewScanner(reader)
+	for inputScanner.Scan() {
+		inputLine := inputScanner.Bytes()
+		sinkInfo := SinkEntry{}
+		if err := json.Unmarshal(inputLine, &sinkInfo); err != nil {
+			return nil, err
+		}
+
+		if (sinkInfo.RecordingStart.Before(c.from) && !sinkInfo.RecordingStart.Equal(c.from)) ||
+			sinkInfo.RecordingEnd.After(c.to) {
+			continue
+		}
+
+		restoredCounter, err := hyperloglog.NewPlus(18)
+		if err != nil {
+			return nil, err
+		}
+		err = restoredCounter.GobDecode(sinkInfo.Recorded)
+		if err != nil {
+			return nil, err
+		}
+		result.ChunkIncluded++
+		err = counter.Merge(restoredCounter)
+		if err != nil {
+			return nil, err
+		}
+	}
+	result.Sum = counter.Count()
+	return &result, nil
+}

common/ipsetsink/sinkcluster/writer.go

@@ -0,0 +1,68 @@
+package sinkcluster
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"log"
+	"time"
+
+	"git.torproject.org/pluggable-transports/snowflake.git/v2/common/ipsetsink"
+)
+
+func NewClusterWriter(writer WriteSyncer, writeInterval time.Duration, sink *ipsetsink.IPSetSink) *ClusterWriter {
+	c := &ClusterWriter{
+		writer:        writer,
+		lastWriteTime: time.Now(),
+		writeInterval: writeInterval,
+		current:       sink,
+	}
+	return c
+}
+
+type ClusterWriter struct {
+	writer        WriteSyncer
+	lastWriteTime time.Time
+	writeInterval time.Duration
+	current       *ipsetsink.IPSetSink
+}
+
+type WriteSyncer interface {
+	Sync() error
+	io.Writer
+}
+
+func (c *ClusterWriter) WriteIPSetToDisk() {
+	currentTime := time.Now()
+	data, err := c.current.Dump()
+	if err != nil {
+		log.Println("unable able to write ipset to file:", err)
+		return
+	}
+	entry := &SinkEntry{
+		RecordingStart: c.lastWriteTime,
+		RecordingEnd:   currentTime,
+		Recorded:       data,
+	}
+	jsonData, err := json.Marshal(entry)
+	if err != nil {
+		log.Println("unable able to write ipset to file:", err)
+		return
+	}
+	jsonData = append(jsonData, byte('\n'))
+	_, err = io.Copy(c.writer, bytes.NewReader(jsonData))
+	if err != nil {
+		log.Println("unable able to write ipset to file:", err)
+		return
+	}
+	c.writer.Sync()
+	c.lastWriteTime = currentTime
+	c.current.Reset()
+}
+
+func (c *ClusterWriter) AddIPToSet(ipAddress string) {
+	if c.lastWriteTime.Add(c.writeInterval).Before(time.Now()) {
+		c.WriteIPSetToDisk()
+	}
+	c.current.AddIPToSet(ipAddress)
+}

common/ipsetsink/sinkcluster/writer_test.go

@@ -0,0 +1,33 @@
+package sinkcluster
+
+import (
+	"bytes"
+	"io"
+	"testing"
+	"time"
+
+	"git.torproject.org/pluggable-transports/snowflake.git/v2/common/ipsetsink"
+
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+type writerStub struct {
+	io.Writer
+}
+
+func (w writerStub) Sync() error {
+	return nil
+}
+
+func TestSinkWriter(t *testing.T) {
+
+	Convey("Context", t, func() {
+		buffer := bytes.NewBuffer(nil)
+		writerStubInst := &writerStub{buffer}
+		sink := ipsetsink.NewIPSetSink("demo")
+		clusterWriter := NewClusterWriter(writerStubInst, time.Minute, sink)
+		clusterWriter.AddIPToSet("1")
+		clusterWriter.WriteIPSetToDisk()
+		So(buffer.Bytes(), ShouldNotBeNil)
+	})
+}