From 76349a6b55810a460f6d7c1af54f9ac80315fc53 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 17 Jun 2026 15:14:54 +0200 Subject: [PATCH] Harden Go SDK head store persistence --- cmd/attesto/main.go | 5 ++-- heads.go | 70 +++++++++++++++++++++++++++++++++++++-------- heads_test.go | 26 +++++++++++++++++ 3 files changed, 87 insertions(+), 14 deletions(-) diff --git a/cmd/attesto/main.go b/cmd/attesto/main.go index 4e80213..ae74a0c 100644 --- a/cmd/attesto/main.go +++ b/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")) diff --git a/heads.go b/heads.go index 6464281..acdb4d3 100644 --- a/heads.go +++ b/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 + } + return closeErr } - _ = os.Chmod(tmpName, 0o600) - _ = os.Rename(tmpName, s.path) + 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 } diff --git a/heads_test.go b/heads_test.go index f50e0ca..fdabab5 100644 --- a/heads_test.go +++ b/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.