From ad4d81fe50f8faaf4481acff4954d001cd7aecca Mon Sep 17 00:00:00 2001
From: lakshit verma <vermalucky2004@gmail.com>
Date: Fri, 21 Aug 2026 21:47:04 +0530
Subject: [PATCH] store: survive corrupt stored values

apply() sliced record values at fixed offsets and getMeta, Count and
Top called Uint64 without checking length, so a truncated or
hand-corrupted value panicked the ingest loop or the stats command.
All decoders now validate lengths; corrupt records are overwritten,
corrupt meta/count values return errors.
---
 internal/store/store.go      | 46 +++++++++++++++++++++++++++++-------
 internal/store/store_test.go | 26 ++++++++++++++++++++
 2 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/internal/store/store.go b/internal/store/store.go
index dd31c16..d6b3186 100644
--- a/internal/store/store.go
+++ b/internal/store/store.go
@@ -182,12 +182,17 @@ func (s *Store) apply(batch *pebble.Batch, r Record) error {
 		if err == pebble.ErrNotFound {
 			st = nil
 		} else {
-			st = &recState{
-				firstSeen: int64(binary.BigEndian.Uint64(existing[:8])),
-				source:    existing[8],
-				seq:       binary.BigEndian.Uint64(existing[9:17]),
+			if len(existing) < 17 {
+				closer.Close()
+				st = nil
+			} else {
+				st = &recState{
+					firstSeen: int64(binary.BigEndian.Uint64(existing[:8])),
+					source:    existing[8],
+					seq:       binary.BigEndian.Uint64(existing[9:17]),
+				}
+				closer.Close()
 			}
-			closer.Close()
 		}
 	}
 	if st == nil {
@@ -228,7 +233,11 @@ func (s *Store) pendingCount(apex string) uint64 {
 		s.pendingCounts[apex] = 0
 		return 0
 	}
-	n := binary.BigEndian.Uint64(v)
+	n, ok := beUint64(v)
+	if !ok {
+		s.pendingCounts[apex] = 0
+		return 0
+	}
 	s.pendingCounts[apex] = n
 	return n
 }
@@ -303,7 +312,11 @@ func (s *Store) Count(apex string) (uint64, error) {
 	if v == nil {
 		return 0, nil
 	}
-	return binary.BigEndian.Uint64(v), nil
+	n, ok := beUint64(v)
+	if !ok {
+		return 0, fmt.Errorf("corrupt count value for %s", apex)
+	}
+	return n, nil
 }
 
 func (s *Store) Total() (uint64, error) {
@@ -320,7 +333,11 @@ func (s *Store) Top(n int) ([]ApexCount, error) {
 	defer it.Close()
 	all := make([]ApexCount, 0, 1024)
 	for ok := it.First(); ok; ok = it.Next() {
-		all = append(all, ApexCount{Apex: string(it.Key()[1:]), Count: binary.BigEndian.Uint64(it.Value())})
+		n, ok := beUint64(it.Value())
+		if !ok {
+			continue
+		}
+		all = append(all, ApexCount{Apex: string(it.Key()[1:]), Count: n})
 	}
 	for i := 0; i < len(all); i++ {
 		for j := i + 1; j < len(all); j++ {
@@ -471,7 +488,11 @@ func (s *Store) getMeta(name string) (uint64, error) {
 	if err != nil || v == nil {
 		return 0, err
 	}
-	return binary.BigEndian.Uint64(v), nil
+	n, ok := beUint64(v)
+	if !ok {
+		return 0, fmt.Errorf("corrupt meta value for %s", name)
+	}
+	return n, nil
 }
 
 func recordKey(apex, sub string) []byte {
@@ -498,6 +519,13 @@ func metaKey(name string) []byte {
 
 func totalKey() []byte { return metaKey("total") }
 
+func beUint64(b []byte) (uint64, bool) {
+	if len(b) != 8 {
+		return 0, false
+	}
+	return binary.BigEndian.Uint64(b), true
+}
+
 func encodeValue(firstSeen int64, source byte, seq uint64) []byte {
 	v := make([]byte, 17)
 	binary.BigEndian.PutUint64(v[:8], uint64(firstSeen))
diff --git a/internal/store/store_test.go b/internal/store/store_test.go
index 168d803..d2ca158 100644
--- a/internal/store/store_test.go
+++ b/internal/store/store_test.go
@@ -1,6 +1,7 @@
 package store
 
 import (
+	"github.com/cockroachdb/pebble/v2"
 	"path/filepath"
 	"testing"
 	"time"
@@ -99,3 +100,28 @@ func TestRecount(t *testing.T) {
 		t.Errorf("total after recount+ingest = %d, want 4", got)
 	}
 }
+
+func TestCorruptValuesDoNotPanic(t *testing.T) {
+	dir := filepath.Join(t.TempDir(), "db")
+	st, err := Open(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := st.Ingest(Record{Apex: "example.com", Sub: "a.example.com", FirstSeen: 1, Source: 1}); err != nil {
+		t.Fatal(err)
+	}
+	time.Sleep(100 * time.Millisecond)
+	st.db.Set([]byte("example.com\x00junk"), []byte{1, 2, 3}, nil)
+	st.db.Set(metaKey("total"), []byte{9}, pebble.Sync)
+	res, err := st.Scan("example.com", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res) != 1 {
+		t.Fatalf("scan = %d results, want 1", len(res))
+	}
+	if _, err := st.Total(); err == nil {
+		t.Error("corrupt meta accepted")
+	}
+	st.Close()
+}