fix: track monotonic entropy initialization

CHANGELOG.md

@@ -1,3 +1,12 @@
+## Unreleased
+
+* Fix monotonic generation after an all-zero 80-bit entropy seed. If another
+  call uses the same timestamp, the zero seed previously selected reseeding
+  instead of incrementing. With independent uniform reads, the zero seed occurs
+  with probability 2^-80 (about 1 in 1.21e24); an actual duplicate also
+  requires the following read to be zero, making the two-read event 2^-160.
+  Other entropy readers can have different distributions.
+
 ## 1.3.1 / 2018-10-02
 
 * Use underlying entropy source for random increments in Monotonic (#32)

ulid.go

@@ -613,19 +613,25 @@ type MonotonicEntropy struct {
 	io.Reader
 	ms      uint64
 	inc     uint64
-	entropy uint80
+	entropy monotonicState
 	rand    [8]byte
 	rng     rng
 }
 
 // MonotonicRead implements the MonotonicReader interface.
 func (m *MonotonicEntropy) MonotonicRead(ms uint64, entropy []byte) (err error) {
-	if !m.entropy.IsZero() && m.ms == ms {
+	if m.entropy.initialized && m.ms == ms {
+		previous := m.entropy
 		err = m.increment()
 		m.entropy.AppendTo(entropy)
+		if err != nil {
+			// Keep a failed increment from replacing the last successful state.
+			m.entropy = previous
+		}
 	} else if _, err = io.ReadFull(m.Reader, entropy); err == nil {
 		m.ms = ms
 		m.entropy.SetBytes(entropy)
+		m.entropy.initialized = true
 	}
 	return err
 }
@@ -694,29 +700,27 @@ func (m *MonotonicEntropy) random() (inc uint64, err error) {
 	return 1 + inc, nil
 }
 
-type uint80 struct {
-	Hi uint16
-	Lo uint64
+// monotonicState stores the 80-bit entropy and whether it has been initialized.
+type monotonicState struct {
+	Hi          uint16
+	initialized bool // Fits in the alignment padding before Lo.
+	Lo          uint64
 }
 
-func (u *uint80) SetBytes(bs []byte) {
+func (u *monotonicState) SetBytes(bs []byte) {
 	u.Hi = binary.BigEndian.Uint16(bs[:2])
 	u.Lo = binary.BigEndian.Uint64(bs[2:])
 }
 
-func (u *uint80) AppendTo(bs []byte) {
+func (u *monotonicState) AppendTo(bs []byte) {
 	binary.BigEndian.PutUint16(bs[:2], u.Hi)
 	binary.BigEndian.PutUint64(bs[2:], u.Lo)
 }
 
-func (u *uint80) Add(n uint64) (overflow bool) {
+func (u *monotonicState) Add(n uint64) (overflow bool) {
 	lo, hi := u.Lo, u.Hi
 	if u.Lo += n; u.Lo < lo {
 		u.Hi++
 	}
 	return u.Hi < hi
 }
-
-func (u uint80) IsZero() bool {
-	return u.Hi == 0 && u.Lo == 0
-}

ulid_test.go

@@ -16,6 +16,7 @@ package ulid_test
 import (
 	"bytes"
 	crand "crypto/rand"
+	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -629,6 +630,127 @@ func TestMonotonic(t *testing.T) {
 	}
 }
 
+type errorOnceReader struct {
+	prefix []byte
+	next   io.Reader
+	err    error
+	failed bool
+}
+
+func (r *errorOnceReader) Read(p []byte) (int, error) {
+	if !r.failed {
+		r.failed = true
+		return copy(p, r.prefix), r.err
+	}
+	return r.next.Read(p)
+}
+
+func TestMonotonicZeroEntropy(t *testing.T) {
+	for _, ms := range []uint64{123, 0} {
+		ms := ms
+		t.Run(fmt.Sprintf("timestamp=%d", ms), func(t *testing.T) {
+			entropy := ulid.Monotonic(
+				io.MultiReader(bytes.NewReader(make([]byte, 20)), crand.Reader),
+				1,
+			)
+			first, err := ulid.New(ms, entropy)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if got, want := first.Entropy(), make([]byte, 10); !bytes.Equal(got, want) {
+				t.Fatalf("first entropy: got %x, want %x", got, want)
+			}
+
+			second, err := ulid.New(ms, entropy)
+			if err != nil {
+				t.Fatal(err)
+			}
+			want := make([]byte, 10)
+			want[len(want)-1] = 1
+			if got := second.Entropy(); !bytes.Equal(got, want) {
+				t.Fatalf("second entropy: got %x, want %x", got, want)
+			}
+		})
+	}
+
+	t.Run("nonzero initial entropy", func(t *testing.T) {
+		initial := bytes.Repeat([]byte{0xA5}, 10)
+		entropy := ulid.Monotonic(bytes.NewReader(initial), 1)
+		first, err := ulid.New(123, entropy)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got := first.Entropy(); !bytes.Equal(got, initial) {
+			t.Fatalf("first entropy: got %x, want %x", got, initial)
+		}
+
+		second, err := ulid.New(123, entropy)
+		if err != nil {
+			t.Fatal(err)
+		}
+		want := append([]byte(nil), initial...)
+		want[len(want)-1]++
+		if got := second.Entropy(); !bytes.Equal(got, want) {
+			t.Fatalf("second entropy: got %x, want %x", got, want)
+		}
+	})
+
+	t.Run("changed timestamp", func(t *testing.T) {
+		firstEntropy := bytes.Repeat([]byte{0x12}, 10)
+		nextEntropy := bytes.Repeat([]byte{0x34}, 10)
+		input := append(append([]byte(nil), firstEntropy...), nextEntropy...)
+		entropy := ulid.Monotonic(bytes.NewReader(input), 1)
+
+		first, err := ulid.New(123, entropy)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got := first.Entropy(); !bytes.Equal(got, firstEntropy) {
+			t.Fatalf("first entropy: got %x, want %x", got, firstEntropy)
+		}
+
+		second, err := ulid.New(124, entropy)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got := second.Entropy(); !bytes.Equal(got, nextEntropy) {
+			t.Fatalf("entropy after timestamp change: got %x, want %x", got, nextEntropy)
+		}
+
+		third, err := ulid.New(124, entropy)
+		if err != nil {
+			t.Fatal(err)
+		}
+		want := append([]byte(nil), nextEntropy...)
+		want[len(want)-1]++
+		if got := third.Entropy(); !bytes.Equal(got, want) {
+			t.Fatalf("entropy after repeated timestamp: got %x, want %x", got, want)
+		}
+	})
+
+	t.Run("failed initial read", func(t *testing.T) {
+		readErr := errors.New("initial entropy read failed")
+		initial := bytes.Repeat([]byte{0x5A}, 10)
+		reader := &errorOnceReader{
+			prefix: []byte{0x01, 0x02},
+			next:   bytes.NewReader(initial),
+			err:    readErr,
+		}
+		entropy := ulid.Monotonic(reader, 1)
+
+		if _, err := ulid.New(123, entropy); err != readErr {
+			t.Fatalf("first read error: got %v, want %v", err, readErr)
+		}
+		second, err := ulid.New(123, entropy)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got := second.Entropy(); !bytes.Equal(got, initial) {
+			t.Fatalf("entropy after retry: got %x, want %x", got, initial)
+		}
+	})
+}
+
 func TestMonotonicOverflow(t *testing.T) {
 	t.Parallel()
 
@@ -652,6 +774,46 @@ func TestMonotonicOverflow(t *testing.T) {
 	}
 }
 
+func TestMonotonicOverflowState(t *testing.T) {
+	t.Parallel()
+
+	initial := bytes.Repeat([]byte{0xFF}, 10)
+	input := append(append([]byte(nil), initial...), make([]byte, 20)...)
+	entropy := ulid.Monotonic(bytes.NewReader(input), 1)
+
+	first, err := ulid.New(123, entropy)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := first.Entropy(); !bytes.Equal(got, initial) {
+		t.Fatalf("initial entropy: got %x, want %x", got, initial)
+	}
+
+	for i := 0; i < 2; i++ {
+		if _, err := ulid.New(123, entropy); err != ulid.ErrMonotonicOverflow {
+			t.Fatalf("same-timestamp call %d after maximum entropy: got err %v, want %v", i+2, err, ulid.ErrMonotonicOverflow)
+		}
+	}
+
+	reset, err := ulid.New(124, entropy)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := reset.Entropy(), make([]byte, 10); !bytes.Equal(got, want) {
+		t.Fatalf("entropy after timestamp change: got %x, want %x", got, want)
+	}
+
+	next, err := ulid.New(124, entropy)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := make([]byte, 10)
+	want[len(want)-1] = 1
+	if got := next.Entropy(); !bytes.Equal(got, want) {
+		t.Fatalf("entropy after timestamp reseed: got %x, want %x", got, want)
+	}
+}
+
 func TestMonotonicSafe(t *testing.T) {
 	t.Parallel()