Harden Go SDK head store persistence

cmd/attesto/main.go

@@ -476,8 +476,9 @@ func (a *app) doctor(ctx context.Context, args []string) error {
 	}
 
 	headStore := attesto.NewFileHeadStore("")
-	headStore.Set("__doctor__", 1, strings.Repeat("0", 64))
-	if seq, hash, ok := headStore.Get("__doctor__"); ok && seq == 1 && hash == strings.Repeat("0", 64) {
+	if err := headStore.SetWithError("__doctor__", 1, strings.Repeat("0", 64)); err != nil {
+		fail("head_store", err)
+	} else if seq, hash, ok := headStore.Get("__doctor__"); ok && seq == 1 && hash == strings.Repeat("0", 64) {
 		pass("head_store", nil)
 	} else {
 		fail("head_store", errors.New("head store readback failed"))

heads.go

@@ -15,6 +15,11 @@ type HeadStore interface {
 	Set(streamID string, seqNo int64, eventHash string)
 }
 
+type errorAwareHeadStore interface {
+	HeadStore
+	SetWithError(streamID string, seqNo int64, eventHash string) error
+}
+
 // ForkDetectedError reports that a receipt did not extend the last accepted head
 // for its stream. The store is NOT advanced when this is returned.
 type ForkDetectedError struct {
@@ -58,6 +63,11 @@ func (s *MemoryHeadStore) Set(streamID string, seqNo int64, eventHash string) {
 	s.heads[streamID] = [2]any{seqNo, eventHash}
 }
 
+func (s *MemoryHeadStore) SetWithError(streamID string, seqNo int64, eventHash string) error {
+	s.Set(streamID, seqNo, eventHash)
+	return nil
+}
+
 // FileHeadStore persists heads to a JSON file (default ~/.attesto/heads.json),
 // giving fork detection across separate process invocations. Writes are atomic.
 type FileHeadStore struct {
@@ -78,15 +88,26 @@ func NewFileHeadStore(path string) *FileHeadStore {
 }
 
 func (s *FileHeadStore) load() map[string][2]json.RawMessage {
-	raw, err := os.ReadFile(s.path)
+	out, err := s.loadWithError()
 	if err != nil {
 		return map[string][2]json.RawMessage{}
 	}
+	return out
+}
+
+func (s *FileHeadStore) loadWithError() (map[string][2]json.RawMessage, error) {
+	raw, err := os.ReadFile(s.path)
+	if os.IsNotExist(err) {
+		return map[string][2]json.RawMessage{}, nil
+	}
+	if err != nil {
+		return nil, err
+	}
 	out := map[string][2]json.RawMessage{}
 	if err := json.Unmarshal(raw, &out); err != nil {
-		return map[string][2]json.RawMessage{}
+		return nil, err
 	}
-	return out
+	return out, nil
 }
 
 func (s *FileHeadStore) Get(streamID string) (int64, string, bool) {
@@ -105,37 +126,59 @@ func (s *FileHeadStore) Get(streamID string) (int64, string, bool) {
 }
 
 func (s *FileHeadStore) Set(streamID string, seqNo int64, eventHash string) {
+	_ = s.SetWithError(streamID, seqNo, eventHash)
+}
+
+func (s *FileHeadStore) SetWithError(streamID string, seqNo int64, eventHash string) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	heads := map[string][2]any{}
-	for key, entry := range s.load() {
+	stored, err := s.loadWithError()
+	if err != nil {
+		return err
+	}
+	for key, entry := range stored {
 		var n int64
 		var h string
-		_ = json.Unmarshal(entry[0], &n)
-		_ = json.Unmarshal(entry[1], &h)
+		if err := json.Unmarshal(entry[0], &n); err != nil {
+			return fmt.Errorf("load stored head seq for %s: %w", key, err)
+		}
+		if err := json.Unmarshal(entry[1], &h); err != nil {
+			return fmt.Errorf("load stored head hash for %s: %w", key, err)
+		}
 		heads[key] = [2]any{n, h}
 	}
 	heads[streamID] = [2]any{seqNo, eventHash}
 	body, err := json.Marshal(heads)
 	if err != nil {
-		return
+		return err
 	}
 	if err := os.MkdirAll(filepath.Dir(s.path), 0o700); err != nil {
-		return
+		return err
 	}
 	tmp, err := os.CreateTemp(filepath.Dir(s.path), ".heads-")
 	if err != nil {
-		return
+		return err
 	}
 	tmpName := tmp.Name()
 	_, writeErr := tmp.Write(body)
 	closeErr := tmp.Close()
 	if writeErr != nil || closeErr != nil {
 		_ = os.Remove(tmpName)
-		return
+		if writeErr != nil {
+			return writeErr
 		}
-	_ = os.Chmod(tmpName, 0o600)
-	_ = os.Rename(tmpName, s.path)
+		return closeErr
+	}
+	if err := os.Chmod(tmpName, 0o600); err != nil {
+		_ = os.Remove(tmpName)
+		return err
+	}
+	if err := os.Rename(tmpName, s.path); err != nil {
+		_ = os.Remove(tmpName)
+		return err
+	}
+	return nil
 }
 
 // checkAndAdvanceHead verifies a receipt extends the stored head, then advances
@@ -160,6 +203,9 @@ func checkAndAdvanceHead(store HeadStore, receipt EventReceipt) error {
 			}
 		}
 	}
+	if errorAware, ok := store.(errorAwareHeadStore); ok {
+		return errorAware.SetWithError(receipt.StreamID, receipt.SeqNo, receipt.EventHash)
+	}
 	store.Set(receipt.StreamID, receipt.SeqNo, receipt.EventHash)
 	return nil
 }

heads_test.go

@@ -81,6 +81,32 @@ func TestFileHeadStorePersistsAndIs0600(t *testing.T) {
 	}
 }
 
+func TestFileHeadStorePersistenceFailureReturned(t *testing.T) {
+	path := t.TempDir()
+	store := NewFileHeadStore(path)
+	if err := checkAndAdvanceHead(store, receipt(1, "h1", "")); err == nil {
+		t.Fatal("expected persistence failure when head-store path is a directory")
+	}
+}
+
+func TestFileHeadStoreCorruptFileFailsClosed(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "heads.json")
+	if err := os.WriteFile(path, []byte("{not-json"), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	store := NewFileHeadStore(path)
+	if err := store.SetWithError("str_demo", 1, "h1"); err == nil {
+		t.Fatal("expected corrupt head store to fail closed")
+	}
+	raw, err := os.ReadFile(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(raw) != "{not-json" {
+		t.Fatal("corrupt head store must not be silently overwritten")
+	}
+}
+
 func TestExactReplayOfStoredHeadIsBenign(t *testing.T) {
 	// [P3.3 regression] A deduplicated resend returns the same receipt; the
 	// head tracker must treat (same seqNo, same eventHash) as a no-op.