NextMyDay/internal/core/auth/totp/totp_test.go

package totp_test

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	pqotp "github.com/pquerna/otp/totp"

	"git.lohmar.co.uk/cclohmar/NextMyDay/internal/core/auth/totp"
)

// mapStore is an in-memory Store. Encryption at rest (A9) is the real store's
// job (proven in the SPEC-006 flow tests); here we pin the TOTP mechanics.
type mapStore struct {
	secrets  map[int64][]byte
	verified map[int64]bool
}

func newMapStore() *mapStore {
	return &mapStore{secrets: map[int64][]byte{}, verified: map[int64]bool{}}
}

func (m *mapStore) PutTOTPSecret(_ context.Context, id int64, secret []byte) error {
	m.secrets[id] = append([]byte(nil), secret...)
	delete(m.verified, id)
	return nil
}

func (m *mapStore) TOTPSecret(_ context.Context, id int64) ([]byte, error) {
	s, ok := m.secrets[id]
	if !ok {
		return nil, errors.New("no secret")
	}
	return s, nil
}

func (m *mapStore) MarkTOTPVerified(_ context.Context, id int64) error {
	m.verified[id] = true
	return nil
}

func TestBeginIssuesOtpauthURLAndStoresSecret(t *testing.T) {
	st := newMapStore()
	svc := totp.New("MyDay", st)

	issued, err := svc.Begin(context.Background(), 7, "guid-abc")
	if err != nil {
		t.Fatalf("Begin: %v", err)
	}
	if !strings.HasPrefix(issued.URL, "otpauth://totp/") {
		t.Errorf("URL = %q, want an otpauth:// URL", issued.URL)
	}
	secret, ok := st.secrets[7]
	if !ok || len(secret) == 0 {
		t.Fatal("Begin did not store a secret")
	}
	if !strings.Contains(issued.URL, string(secret)) {
		t.Error("the otpauth URL does not carry the stored secret")
	}
}

func TestVerifyAcceptsAValidCodeAndMarksVerified(t *testing.T) {
	st := newMapStore()
	svc := totp.New("MyDay", st)
	if _, err := svc.Begin(context.Background(), 1, "guid"); err != nil {
		t.Fatalf("Begin: %v", err)
	}
	code, err := pqotp.GenerateCode(string(st.secrets[1]), time.Now())
	if err != nil {
		t.Fatalf("GenerateCode: %v", err)
	}
	if err := svc.Verify(context.Background(), 1, code); err != nil {
		t.Fatalf("Verify(valid) = %v, want nil", err)
	}
	if !st.verified[1] {
		t.Error("a valid code did not mark the secret verified")
	}
}

func TestVerifyRejectsABadOrEmptyCode(t *testing.T) {
	st := newMapStore()
	svc := totp.New("MyDay", st)
	if _, err := svc.Begin(context.Background(), 1, "guid"); err != nil {
		t.Fatalf("Begin: %v", err)
	}
	if err := svc.Verify(context.Background(), 1, ""); !errors.Is(err, totp.ErrInvalidCode) {
		t.Errorf("Verify(\"\") = %v, want ErrInvalidCode", err)
	}
	if err := svc.Verify(context.Background(), 1, "000000"); !errors.Is(err, totp.ErrInvalidCode) {
		t.Errorf("Verify(bad) = %v, want ErrInvalidCode", err)
	}
	if st.verified[1] {
		t.Error("a bad code marked the secret verified")
	}
}

// TestVerifyRejectsAReplay covers L5s: a code that already worked must not work
// again within the skew window.
func TestVerifyRejectsAReplay(t *testing.T) {
	st := newMapStore()
	svc := totp.New("MyDay", st)
	ctx := context.Background()
	if _, err := svc.Begin(ctx, 1, "guid"); err != nil {
		t.Fatalf("Begin: %v", err)
	}
	now := time.Now()
	svc.SetClock(func() time.Time { return now })
	secret := string(st.secrets[1])

	code, err := pqotp.GenerateCode(secret, now)
	if err != nil {
		t.Fatalf("GenerateCode: %v", err)
	}
	if err := svc.Verify(ctx, 1, code); err != nil {
		t.Fatalf("first Verify = %v, want nil", err)
	}
	// The same code, still inside its window, is a replay (L5s).
	if err := svc.Verify(ctx, 1, code); !errors.Is(err, totp.ErrInvalidCode) {
		t.Errorf("replayed Verify = %v, want ErrInvalidCode (L5s)", err)
	}

	// A code from an EARLIER step is refused once a newer step was accepted.
	newer, err := pqotp.GenerateCode(secret, now.Add(30*time.Second))
	if err != nil {
		t.Fatalf("GenerateCode: %v", err)
	}
	svc.SetClock(func() time.Time { return now.Add(30 * time.Second) })
	if err := svc.Verify(ctx, 1, newer); err != nil {
		t.Fatalf("Verify(newer) = %v, want nil", err)
	}
	svc.SetClock(func() time.Time { return now }) // the skew window still covers step N
	if err := svc.Verify(ctx, 1, code); !errors.Is(err, totp.ErrInvalidCode) {
		t.Errorf("an older step after a newer one = %v, want ErrInvalidCode (L5s)", err)
	}

	// A fresh secret resets the counter (re-enrollment, L5s).
	if _, err := svc.Begin(ctx, 1, "guid"); err != nil {
		t.Fatalf("Begin (re-enroll): %v", err)
	}
	fresh, err := pqotp.GenerateCode(string(st.secrets[1]), now)
	if err != nil {
		t.Fatalf("GenerateCode: %v", err)
	}
	if err := svc.Verify(ctx, 1, fresh); err != nil {
		t.Errorf("Verify after re-enrollment = %v, want nil", err)
	}
}