diff --git a/backend/services/verification/store.go b/backend/services/verification/store.go index 301bd84..50600a0 100644 --- a/backend/services/verification/store.go +++ b/backend/services/verification/store.go @@ -359,33 +359,48 @@ func (s *Store) cleanupLocked(now time.Time) { } func randomCode(length int) (string, error) { + return randomCodeFromReader(length, rand.Reader) +} + +func randomCodeFromReader(length int, reader io.Reader) (string, error) { if length <= 0 || length > 12 { return "", ErrInvalidChallenge } const digits = "0123456789" code := make([]byte, length) random := make([]byte, length*2) - for written := 0; written < length; { - if _, err := io.ReadFull(rand.Reader, random); err != nil { - clearBytes(random) - clearBytes(code) - return "", err - } - for _, value := range random { - if value >= 250 { - continue + defer clearBytes(random) + defer clearBytes(code) + for { + var counts [10]uint8 + for written := 0; written < length; { + if _, err := io.ReadFull(reader, random); err != nil { + return "", err } - code[written] = digits[int(value)%len(digits)] - written++ - if written == length { + for _, value := range random { + if value >= 250 { + continue + } + digit := int(value) % len(digits) + code[written] = digits[digit] + counts[digit]++ + written++ + if written == length { + break + } + } + } + valid := true + for _, count := range counts { + if count > 2 { + valid = false break } } + if valid { + return string(code), nil + } } - clearBytes(random) - result := string(code) - clearBytes(code) - return result, nil } func randomID() (string, error) { diff --git a/backend/services/verification/store_test.go b/backend/services/verification/store_test.go index aa24130..853e5cc 100644 --- a/backend/services/verification/store_test.go +++ b/backend/services/verification/store_test.go @@ -1,11 +1,87 @@ package verification import ( + "bytes" "errors" + "io" "testing" "time" ) +func TestRandomCodeRejectsThreeOccurrences(t *testing.T) { + cases := []struct { + name string + digits []byte + }{ + {name: "33626000", digits: []byte{3, 3, 6, 2, 6, 0, 0, 0}}, + {name: "separated", digits: []byte{1, 2, 1, 3, 1, 4, 5, 6}}, + {name: "identical", digits: []byte{7, 7, 7, 7, 7, 7, 7, 7}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + raw := make([]byte, DefaultCodeLength*4) + copy(raw, tc.digits) + copy(raw[DefaultCodeLength*2:], []byte{0, 0, 1, 1, 2, 2, 3, 3}) + code, err := randomCodeFromReader(DefaultCodeLength, bytes.NewReader(raw)) + if err != nil { + t.Fatal(err) + } + if code != "00112233" { + t.Fatalf("code = %q, want 00112233", code) + } + }) + } +} + +func TestRandomCodeSupportedLengths(t *testing.T) { + for length := 1; length <= 12; length++ { + raw := bytes.Repeat([]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, 3) + code, err := randomCodeFromReader(length, bytes.NewReader(raw)) + if err != nil { + t.Fatal(err) + } + if len(code) != length { + t.Fatalf("code length = %d, want %d", len(code), length) + } + var counts [10]int + for _, digit := range []byte(code) { + if digit < '0' || digit > '9' { + t.Fatal("non-digit code") + } + counts[digit-'0']++ + if counts[digit-'0'] > 2 { + t.Fatal("digit occurred more than twice") + } + } + } +} + +func TestRandomCodeRejectsBiasedBytes(t *testing.T) { + raw := []byte{250, 251, 252, 253, 254, 255, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4} + code, err := randomCodeFromReader(DefaultCodeLength, bytes.NewReader(raw)) + if err != nil || code != "00112233" { + t.Fatalf("code = %q, error = %v", code, err) + } +} + +func TestRandomCodeReaderFailure(t *testing.T) { + for _, raw := range [][]byte{nil, {1}, make([]byte, DefaultCodeLength*2)} { + code, err := randomCodeFromReader(DefaultCodeLength, bytes.NewReader(raw)) + if code != "" || (!errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF)) { + t.Fatalf("code = %q, error = %v", code, err) + } + } +} + +func TestRandomCodeRejectsInvalidLength(t *testing.T) { + for _, length := range []int{-1, 0, 13} { + code, err := randomCodeFromReader(length, bytes.NewReader(nil)) + if code != "" || !errors.Is(err, ErrInvalidChallenge) { + t.Fatalf("length = %d, code = %q, error = %v", length, code, err) + } + } +} + func TestStoreIssuesAndConsumesCodeOnlyOnce(t *testing.T) { store, err := NewStore(4) if err != nil {