Move replay into subpackage

2 files changed, 14 insertions, 7 deletions

diff --git a/replay.go b/replay/replay.go
index 2d9c6e1..993ff58 100644
--- a/replay.go
+++ b/replay/replay.go
@@ -4,9 +4,7 @@
  * Copyright (C) 2017-2018 Mathias N. Hall-Andersen <mathias@hall-andersen.dk>.
  */
 
-package main
-
-/* Copyright (C) 2015-2017 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved. */
+package replay
 
 /* Implementation of RFC6479
  * https://tools.ietf.org/html/rfc6479
@@ -32,6 +30,13 @@ const (
 	BacktrackWords = CounterBitsTotal / _WordSize
 )
 
+func minUint64(a uint64, b uint64) uint64 {
+	if a > b {
+		return b
+	}
+	return a
+}
+
 type ReplayFilter struct {
 	counter   uint64
 	backtrack [BacktrackWords]uintptr
@@ -42,8 +47,8 @@ func (filter *ReplayFilter) Init() {
 	filter.backtrack[0] = 0
 }
 
-func (filter *ReplayFilter) ValidateCounter(counter uint64) bool {
-	if counter >= RejectAfterMessages {
+func (filter *ReplayFilter) ValidateCounter(counter uint64, limit uint64) bool {
+	if counter >= limit {
 		return false
 	}
 
diff --git a/replay_test.go b/replay/replay_test.go
index 8b2e57d..da39498 100644
--- a/replay_test.go
+++ b/replay/replay_test.go
@@ -4,7 +4,7 @@
  * Copyright (C) 2017-2018 Mathias N. Hall-Andersen <mathias@hall-andersen.dk>.
  */
 
-package main
+package replay
 
 import (
 	"testing"
@@ -15,6 +15,8 @@ import (
  *
  */
 
+const RejectAfterMessages = (1 << 64) - (1 << 4) - 1
+
 func TestReplay(t *testing.T) {
 	var filter ReplayFilter
 
@@ -23,7 +25,7 @@ func TestReplay(t *testing.T) {
 	testNumber := 0
 	T := func(n uint64, v bool) {
 		testNumber++
-		if filter.ValidateCounter(n) != v {
+		if filter.ValidateCounter(n, RejectAfterMessages) != v {
 			t.Fatal("Test", testNumber, "failed", n, v)
 		}
 	}