From f6a9929108767d18519724f143170835e8d13532 Mon Sep 17 00:00:00 2001 From: Ben Blier Date: Fri, 20 Mar 2026 19:17:59 -0400 Subject: [PATCH] Add change counter to geoprobes that is incremented when targets are added or removed. --- CHANGELOG.md | 7 + .../telemetry/cmd/geoprobe-agent/main.go | 31 ++-- .../internal/geoprobe/onchain_discovery.go | 45 +++--- .../geoprobe/onchain_discovery_test.go | 44 ++++++ .../internal/geoprobe/target_discovery.go | 77 ++++++--- .../geoprobe/target_discovery_test.go | 149 ++++++++++++++++++ sdk/geolocation/go/state.go | 8 + sdk/geolocation/go/state_test.go | 34 ++++ .../cli/src/geolocation/probe/get.rs | 1 + .../cli/src/geolocation/probe/list.rs | 2 + .../cli/src/geolocation/user/add_target.rs | 1 + .../cli/src/geolocation/user/remove_target.rs | 1 + .../src/processors/geo_probe/create.rs | 1 + .../processors/geolocation_user/add_target.rs | 1 + .../geolocation_user/remove_target.rs | 1 + .../src/state/geo_probe.rs | 69 +++++--- .../tests/geo_probe_test.rs | 2 + .../tests/geolocation_user_test.rs | 7 +- .../sdk/rs/src/geolocation/geo_probe/get.rs | 1 + .../sdk/rs/src/geolocation/geo_probe/list.rs | 1 + 20 files changed, 405 insertions(+), 78 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 34165a1fd..70f35cc8f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,13 @@ All notable changes to this project will be documented in this file. ### Changes +- Onchain Programs + - Add `target_update_count` field to GeoProbe account, incremented on `AddTarget` and `RemoveTarget`; uses `BorshDeserializeIncremental` so existing accounts default to 0 (non-breaking) +- SDK + - Add `TargetUpdateCount` field to Go GeoProbe struct with backward-compatible deserialization +- Telemetry + - Skip expensive `GetGeolocationUsers` RPC scan in geoprobe-agent when the probe's `target_update_count` is unchanged, with a forced full refresh every ~5 minutes as safety net + ## [v0.13.0](https://github.com/malbeclabs/doublezero/compare/client/v0.12.0...client/v0.13.0) - 2026-03-20 ### Breaking diff --git a/controlplane/telemetry/cmd/geoprobe-agent/main.go b/controlplane/telemetry/cmd/geoprobe-agent/main.go index c2640e8c7..51fc4c39d 100644 --- a/controlplane/telemetry/cmd/geoprobe-agent/main.go +++ b/controlplane/telemetry/cmd/geoprobe-agent/main.go @@ -12,6 +12,7 @@ import ( "os/signal" "strings" "sync" + "sync/atomic" "syscall" "time" @@ -533,6 +534,10 @@ func main() { } }() + // Shared counter: parent discovery writes the GeoProbe target_update_count on each + // poll; target discovery reads it to skip expensive full scans when unchanged. + var probeTargetUpdateCount atomic.Uint32 + // Run parent DZD discovery if program IDs are configured. if parentDiscoveryEnabled { parentUpdateCh := make(chan geoprobe.ParentUpdate, 1) @@ -540,12 +545,13 @@ func main() { deviceResolver := geoprobe.NewRPCDeviceResolver(rpcClient, serviceabilityProgramID) pd, err := geoprobe.NewParentDiscovery(&geoprobe.ParentDiscoveryConfig{ - GeoProbePubkey: geoProbePubkey, - Client: geoProbeClient, - Resolver: deviceResolver, - CLIParents: cliParentAuthorities, - Interval: parentDiscoveryInterval, - Logger: log, + GeoProbePubkey: geoProbePubkey, + Client: geoProbeClient, + Resolver: deviceResolver, + CLIParents: cliParentAuthorities, + Interval: parentDiscoveryInterval, + Logger: log, + ProbeTargetUpdateCount: &probeTargetUpdateCount, }) if err != nil { log.Error("Failed to create parent discovery", "error", err) @@ -578,12 +584,13 @@ func main() { if !geolocationProgramID.IsZero() { geolocationUserClient := geolocation.New(log, rpcClient, geolocationProgramID) td, err := geoprobe.NewTargetDiscovery(&geoprobe.TargetDiscoveryConfig{ - GeoProbePubkey: geoProbePubkey, - Client: geolocationUserClient, - CLITargets: targets, - CLIAllowedKeys: allowedKeys, - Interval: discoveryInterval, - Logger: log, + GeoProbePubkey: geoProbePubkey, + Client: geolocationUserClient, + CLITargets: targets, + CLIAllowedKeys: allowedKeys, + Interval: discoveryInterval, + Logger: log, + ProbeTargetUpdateCount: &probeTargetUpdateCount, }) if err != nil { log.Error("Failed to create target discovery", "error", err) diff --git a/controlplane/telemetry/internal/geoprobe/onchain_discovery.go b/controlplane/telemetry/internal/geoprobe/onchain_discovery.go index c1b5c705a..1f1f4bd3d 100644 --- a/controlplane/telemetry/internal/geoprobe/onchain_discovery.go +++ b/controlplane/telemetry/internal/geoprobe/onchain_discovery.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "log/slog" + "sync/atomic" "time" "github.com/gagliardetto/solana-go" @@ -38,22 +39,24 @@ type ParentUpdate struct { // ParentDiscoveryConfig holds configuration for parent discovery. type ParentDiscoveryConfig struct { - GeoProbePubkey solana.PublicKey - Client GeoProbeAccountClient - Resolver DeviceResolver - CLIParents map[[32]byte][32]byte // static parents from --additional-parent - Interval time.Duration - Logger *slog.Logger + GeoProbePubkey solana.PublicKey + Client GeoProbeAccountClient + Resolver DeviceResolver + CLIParents map[[32]byte][32]byte // static parents from --additional-parent + Interval time.Duration + Logger *slog.Logger + ProbeTargetUpdateCount *atomic.Uint32 // shared counter for target discovery change detection } // ParentDiscovery polls the GeoProbe account and resolves parent devices. type ParentDiscovery struct { - log *slog.Logger - geoProbePubkey solana.PublicKey - client GeoProbeAccountClient - resolver DeviceResolver - cliParents map[[32]byte][32]byte - interval time.Duration + log *slog.Logger + geoProbePubkey solana.PublicKey + client GeoProbeAccountClient + resolver DeviceResolver + cliParents map[[32]byte][32]byte + interval time.Duration + probeTargetUpdateCount *atomic.Uint32 cachedParentDevices []solana.PublicKey tickCount uint64 @@ -83,12 +86,13 @@ func NewParentDiscovery(cfg *ParentDiscoveryConfig) (*ParentDiscovery, error) { } return &ParentDiscovery{ - log: cfg.Logger, - geoProbePubkey: cfg.GeoProbePubkey, - client: cfg.Client, - resolver: cfg.Resolver, - cliParents: cliParents, - interval: cfg.Interval, + log: cfg.Logger, + geoProbePubkey: cfg.GeoProbePubkey, + client: cfg.Client, + resolver: cfg.Resolver, + cliParents: cliParents, + interval: cfg.Interval, + probeTargetUpdateCount: cfg.ProbeTargetUpdateCount, }, nil } @@ -152,6 +156,11 @@ func (d *ParentDiscovery) discover(ctx context.Context) (*ParentUpdate, error) { return d.cliOnlyUpdate(), nil } + // Publish the probe's target_update_count for target discovery change detection. + if d.probeTargetUpdateCount != nil { + d.probeTargetUpdateCount.Store(probe.TargetUpdateCount) + } + // Check if parent device set changed since last poll. if !forceFullRefresh && pubkeySlicesEqual(d.cachedParentDevices, probe.ParentDevices) { d.log.Debug("Parent device set unchanged, skipping resolution", diff --git a/controlplane/telemetry/internal/geoprobe/onchain_discovery_test.go b/controlplane/telemetry/internal/geoprobe/onchain_discovery_test.go index 03b98413f..caf017d69 100644 --- a/controlplane/telemetry/internal/geoprobe/onchain_discovery_test.go +++ b/controlplane/telemetry/internal/geoprobe/onchain_discovery_test.go @@ -5,6 +5,7 @@ import ( "errors" "log/slog" "os" + "sync/atomic" "testing" "time" @@ -611,3 +612,46 @@ func TestPubkeySlicesEqual(t *testing.T) { }) } } + +func TestParentDiscovery_StoresTargetUpdateCount(t *testing.T) { + t.Parallel() + + geoProbePK := solana.NewWallet().PublicKey() + parentDevicePK := solana.NewWallet().PublicKey() + var metricsKey [32]byte + metricsKey = solana.NewWallet().PublicKey() + + client := &mockGeoProbeAccountClient{ + probe: &geolocation.GeoProbe{ + AccountType: geolocation.AccountTypeGeoProbe, + ParentDevices: []solana.PublicKey{parentDevicePK}, + TargetUpdateCount: 42, + }, + } + + resolver := &mockDeviceResolver{ + devices: map[solana.PublicKey]*serviceability.Device{ + parentDevicePK: { + PublicIp: [4]uint8{10, 0, 0, 1}, + MetricsPublisherPubKey: metricsKey, + }, + }, + } + + var counter atomic.Uint32 + pd, err := NewParentDiscovery(&ParentDiscoveryConfig{ + Logger: slog.New(slog.NewTextHandler(os.Stderr, nil)), + Client: client, + Resolver: resolver, + GeoProbePubkey: geoProbePK, + Interval: 10 * time.Millisecond, + ProbeTargetUpdateCount: &counter, + }) + require.NoError(t, err) + + update, err := pd.discover(context.Background()) + require.NoError(t, err) + require.NotNil(t, update) + + assert.Equal(t, uint32(42), counter.Load()) +} diff --git a/controlplane/telemetry/internal/geoprobe/target_discovery.go b/controlplane/telemetry/internal/geoprobe/target_discovery.go index ab842251a..42d24525c 100644 --- a/controlplane/telemetry/internal/geoprobe/target_discovery.go +++ b/controlplane/telemetry/internal/geoprobe/target_discovery.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "sort" + "sync/atomic" "time" "github.com/gagliardetto/solana-go" @@ -27,30 +28,38 @@ type InboundKeyUpdate struct { Keys [][32]byte } +// targetDiscoveryFullRefreshEvery controls how often a full GeolocationUser scan +// is forced regardless of whether the GeoProbe target_update_count has changed. +// At the default 60s interval, 5 means a full refresh every ~5 minutes. +const targetDiscoveryFullRefreshEvery = 5 + // TargetDiscoveryConfig holds configuration for target discovery. type TargetDiscoveryConfig struct { - GeoProbePubkey solana.PublicKey - Client GeolocationUserClient - CLITargets []ProbeAddress - CLIAllowedKeys [][32]byte - Interval time.Duration - Logger *slog.Logger + GeoProbePubkey solana.PublicKey + Client GeolocationUserClient + CLITargets []ProbeAddress + CLIAllowedKeys [][32]byte + Interval time.Duration + Logger *slog.Logger + ProbeTargetUpdateCount *atomic.Uint32 // shared counter from parent discovery } // TargetDiscovery polls GeolocationUser accounts and sends target/key updates // when changes are detected. It filters for activated, paid users whose targets // reference this probe's pubkey. type TargetDiscovery struct { - log *slog.Logger - geoProbePubkey solana.PublicKey - client GeolocationUserClient - cliTargets []ProbeAddress - cliAllowedKeys [][32]byte - interval time.Duration - - cachedTargets []ProbeAddress - cachedInboundKeys [][32]byte - tickCount uint64 + log *slog.Logger + geoProbePubkey solana.PublicKey + client GeolocationUserClient + cliTargets []ProbeAddress + cliAllowedKeys [][32]byte + interval time.Duration + probeTargetUpdateCount *atomic.Uint32 + + cachedTargets []ProbeAddress + cachedInboundKeys [][32]byte + tickCount uint64 + lastSeenTargetUpdateCount uint32 } // NewTargetDiscovery creates a new TargetDiscovery instance. @@ -69,12 +78,13 @@ func NewTargetDiscovery(cfg *TargetDiscoveryConfig) (*TargetDiscovery, error) { } return &TargetDiscovery{ - log: cfg.Logger, - geoProbePubkey: cfg.GeoProbePubkey, - client: cfg.Client, - cliTargets: cfg.CLITargets, - cliAllowedKeys: cfg.CLIAllowedKeys, - interval: cfg.Interval, + log: cfg.Logger, + geoProbePubkey: cfg.GeoProbePubkey, + client: cfg.Client, + cliTargets: cfg.CLITargets, + cliAllowedKeys: cfg.CLIAllowedKeys, + interval: cfg.Interval, + probeTargetUpdateCount: cfg.ProbeTargetUpdateCount, }, nil } @@ -111,6 +121,11 @@ func (d *TargetDiscovery) discoverAndSend(ctx context.Context, targetCh chan<- T return } + // nil targets means the scan was skipped (target_update_count unchanged). + if targets == nil && inboundKeys == nil { + return + } + if !probeAddressSlicesEqual(targets, d.cachedTargets) { d.cachedTargets = targets select { @@ -131,10 +146,21 @@ func (d *TargetDiscovery) discoverAndSend(ctx context.Context, targetCh chan<- T } // discover performs a single discovery cycle: fetch users, filter, extract targets/keys, -// merge with CLI values. +// merge with CLI values. Returns nil, nil, nil when the scan is skipped. func (d *TargetDiscovery) discover(ctx context.Context) ([]ProbeAddress, [][32]byte, error) { + forceFullRefresh := d.tickCount%targetDiscoveryFullRefreshEvery == 0 d.tickCount++ + if d.probeTargetUpdateCount != nil && !forceFullRefresh { + current := d.probeTargetUpdateCount.Load() + if current == d.lastSeenTargetUpdateCount && d.tickCount > 1 { + d.log.Debug("GeoProbe target_update_count unchanged, skipping target scan", + "targetUpdateCount", current) + return nil, nil, nil + } + d.lastSeenTargetUpdateCount = current + } + users, err := d.client.GetGeolocationUsers(ctx) if err != nil { return nil, nil, fmt.Errorf("failed to fetch GeolocationUser accounts: %w", err) @@ -195,6 +221,11 @@ func (d *TargetDiscovery) discover(ctx context.Context) ([]ProbeAddress, [][32]b mergedTargets := mergeProbes(d.cliTargets, onchainTargets) mergedKeys := mergeKeys(d.cliAllowedKeys, onchainKeys) + // Sync lastSeenTargetUpdateCount after a full scan (covers forced refresh path). + if d.probeTargetUpdateCount != nil { + d.lastSeenTargetUpdateCount = d.probeTargetUpdateCount.Load() + } + d.log.Debug("Target discovery tick", "users", len(users), "onchainOutbound", len(onchainTargets), diff --git a/controlplane/telemetry/internal/geoprobe/target_discovery_test.go b/controlplane/telemetry/internal/geoprobe/target_discovery_test.go index 069597a34..f2b91dbef 100644 --- a/controlplane/telemetry/internal/geoprobe/target_discovery_test.go +++ b/controlplane/telemetry/internal/geoprobe/target_discovery_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "log/slog" + "sync/atomic" "testing" "time" @@ -15,9 +16,11 @@ import ( type mockGeolocationUserClient struct { users []geolocation.KeyedGeolocationUser err error + calls int } func (m *mockGeolocationUserClient) GetGeolocationUsers(_ context.Context) ([]geolocation.KeyedGeolocationUser, error) { + m.calls++ return m.users, m.err } @@ -408,6 +411,152 @@ func TestNewTargetDiscovery_Validation(t *testing.T) { } } +func TestTargetDiscovery_TargetUpdateCountUnchanged_SkipsScan(t *testing.T) { + probePK := testProbePubkey() + client := &mockGeolocationUserClient{ + users: []geolocation.KeyedGeolocationUser{ + makeUser(geolocation.GeolocationUserStatusActivated, geolocation.GeolocationPaymentStatusPaid, "user1", []geolocation.GeolocationTarget{ + outboundTarget([4]uint8{44, 0, 0, 1}, 9000, probePK), + }), + }, + } + + var counter atomic.Uint32 + counter.Store(5) + + td, _ := NewTargetDiscovery(&TargetDiscoveryConfig{ + GeoProbePubkey: testProbePubkey(), + Client: client, + Interval: time.Minute, + Logger: slog.Default(), + ProbeTargetUpdateCount: &counter, + }) + + // First call (tick 0): always does full scan (forceFullRefresh). + targets, _, err := td.discover(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(targets) != 1 { + t.Fatalf("expected 1 target on first scan, got %d", len(targets)) + } + if client.calls != 1 { + t.Fatalf("expected 1 RPC call on first scan, got %d", client.calls) + } + + // Second call: counter unchanged → should skip. + targets, keys, err := td.discover(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if targets != nil || keys != nil { + t.Errorf("expected nil targets/keys when skipped, got targets=%v keys=%v", targets, keys) + } + if client.calls != 1 { + t.Errorf("expected no additional RPC call when skipped, got %d total", client.calls) + } +} + +func TestTargetDiscovery_TargetUpdateCountChanged_DoesFullScan(t *testing.T) { + probePK := testProbePubkey() + client := &mockGeolocationUserClient{ + users: []geolocation.KeyedGeolocationUser{ + makeUser(geolocation.GeolocationUserStatusActivated, geolocation.GeolocationPaymentStatusPaid, "user1", []geolocation.GeolocationTarget{ + outboundTarget([4]uint8{44, 0, 0, 1}, 9000, probePK), + }), + }, + } + + var counter atomic.Uint32 + counter.Store(5) + + td, _ := NewTargetDiscovery(&TargetDiscoveryConfig{ + GeoProbePubkey: testProbePubkey(), + Client: client, + Interval: time.Minute, + Logger: slog.Default(), + ProbeTargetUpdateCount: &counter, + }) + + // First call: full scan. + _, _, _ = td.discover(context.Background()) + + // Change counter, second call should do full scan. + counter.Store(6) + targets, _, err := td.discover(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(targets) != 1 { + t.Fatalf("expected 1 target after counter change, got %d", len(targets)) + } + if client.calls != 2 { + t.Errorf("expected 2 RPC calls total, got %d", client.calls) + } +} + +func TestTargetDiscovery_ForcedFullRefresh_IgnoresCounter(t *testing.T) { + probePK := testProbePubkey() + client := &mockGeolocationUserClient{ + users: []geolocation.KeyedGeolocationUser{ + makeUser(geolocation.GeolocationUserStatusActivated, geolocation.GeolocationPaymentStatusPaid, "user1", []geolocation.GeolocationTarget{ + outboundTarget([4]uint8{44, 0, 0, 1}, 9000, probePK), + }), + }, + } + + var counter atomic.Uint32 + counter.Store(5) + + td, _ := NewTargetDiscovery(&TargetDiscoveryConfig{ + GeoProbePubkey: testProbePubkey(), + Client: client, + Interval: time.Minute, + Logger: slog.Default(), + ProbeTargetUpdateCount: &counter, + }) + + // Tick through to the next forced refresh (every 5th tick). + // Tick 0: forced (0 % 5 == 0), tick 1-4: skipped (counter unchanged), tick 5: forced. + for i := 0; i < targetDiscoveryFullRefreshEvery; i++ { + _, _, _ = td.discover(context.Background()) + } + callsBefore := client.calls + + // Next tick (tick 5): forced full refresh even though counter unchanged. + targets, _, err := td.discover(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(targets) != 1 { + t.Fatalf("expected 1 target on forced refresh, got %d", len(targets)) + } + if client.calls != callsBefore+1 { + t.Errorf("expected forced refresh to call RPC, calls before=%d after=%d", callsBefore, client.calls) + } +} + +func TestTargetDiscovery_NilProbeTargetUpdateCount_AlwaysScans(t *testing.T) { + probePK := testProbePubkey() + client := &mockGeolocationUserClient{ + users: []geolocation.KeyedGeolocationUser{ + makeUser(geolocation.GeolocationUserStatusActivated, geolocation.GeolocationPaymentStatusPaid, "user1", []geolocation.GeolocationTarget{ + outboundTarget([4]uint8{44, 0, 0, 1}, 9000, probePK), + }), + }, + } + + // No ProbeTargetUpdateCount set — backward compat: always scans. + td := newTestTargetDiscovery(client, nil, nil) + + for i := 0; i < 3; i++ { + _, _, _ = td.discover(context.Background()) + } + if client.calls != 3 { + t.Errorf("expected 3 RPC calls without ProbeTargetUpdateCount, got %d", client.calls) + } +} + func TestTargetDiscovery_RejectsNonPublicOutboundTargets(t *testing.T) { probePK := testProbePubkey() diff --git a/sdk/geolocation/go/state.go b/sdk/geolocation/go/state.go index ecdf63443..ddaa368c1 100644 --- a/sdk/geolocation/go/state.go +++ b/sdk/geolocation/go/state.go @@ -74,6 +74,7 @@ type GeoProbe struct { ReferenceCount uint32 // 4 bytes LE Code string // 4-byte length prefix + UTF-8 bytes ParentDevices []solana.PublicKey // 4-byte count + N*32 bytes + TargetUpdateCount uint32 // 4 bytes LE (appended; defaults to 0 for old accounts) } func (g *GeoProbe) Serialize(w io.Writer) error { @@ -105,6 +106,9 @@ func (g *GeoProbe) Serialize(w io.Writer) error { if err := enc.Encode(g.ParentDevices); err != nil { return err } + if err := enc.Encode(g.TargetUpdateCount); err != nil { + return err + } return nil } @@ -143,6 +147,10 @@ func (g *GeoProbe) Deserialize(data []byte) error { if len(g.ParentDevices) > MaxParentDevices { return fmt.Errorf("parent devices count %d exceeds max allowed %d", len(g.ParentDevices), MaxParentDevices) } + // TargetUpdateCount is appended; old accounts without it default to 0. + if err := dec.Decode(&g.TargetUpdateCount); err != nil { + g.TargetUpdateCount = 0 + } return nil } diff --git a/sdk/geolocation/go/state_test.go b/sdk/geolocation/go/state_test.go index 06748450a..257dfa456 100644 --- a/sdk/geolocation/go/state_test.go +++ b/sdk/geolocation/go/state_test.go @@ -47,6 +47,7 @@ func TestSDK_Geolocation_State_GeoProbe_RoundTrip(t *testing.T) { }, MetricsPublisherPK: solana.NewWallet().PublicKey(), ReferenceCount: 5, + TargetUpdateCount: 42, } var buf bytes.Buffer @@ -64,6 +65,7 @@ func TestSDK_Geolocation_State_GeoProbe_RoundTrip(t *testing.T) { require.Equal(t, original.ParentDevices, decoded.ParentDevices) require.Equal(t, original.MetricsPublisherPK, decoded.MetricsPublisherPK) require.Equal(t, original.ReferenceCount, decoded.ReferenceCount) + require.Equal(t, original.TargetUpdateCount, decoded.TargetUpdateCount) } func TestSDK_Geolocation_State_GeoProbe_EmptyParentDevices(t *testing.T) { @@ -123,6 +125,38 @@ func TestSDK_Geolocation_State_GeoProbe_MaxParentDevices(t *testing.T) { } } +func TestSDK_Geolocation_State_GeoProbe_BackwardCompat_NoTargetUpdateCount(t *testing.T) { + t.Parallel() + + // Serialize a GeoProbe with TargetUpdateCount, then truncate the last 4 bytes + // to simulate an old account that was serialized before target_update_count existed. + original := &geolocation.GeoProbe{ + AccountType: geolocation.AccountTypeGeoProbe, + Owner: solana.NewWallet().PublicKey(), + ExchangePK: solana.NewWallet().PublicKey(), + PublicIP: [4]uint8{10, 0, 1, 1}, + LocationOffsetPort: 8923, + Code: "old-probe", + ParentDevices: []solana.PublicKey{solana.NewWallet().PublicKey()}, + MetricsPublisherPK: solana.NewWallet().PublicKey(), + ReferenceCount: 3, + TargetUpdateCount: 0, + } + + var buf bytes.Buffer + require.NoError(t, original.Serialize(&buf)) + + // Truncate the trailing target_update_count (4 bytes) to simulate old data. + data := buf.Bytes()[:buf.Len()-4] + + var decoded geolocation.GeoProbe + require.NoError(t, decoded.Deserialize(data)) + + require.Equal(t, original.Owner, decoded.Owner) + require.Equal(t, original.ParentDevices, decoded.ParentDevices) + require.Equal(t, uint32(0), decoded.TargetUpdateCount) +} + func TestSDK_Geolocation_State_GeolocationUser_RoundTrip(t *testing.T) { t.Parallel() diff --git a/smartcontract/cli/src/geolocation/probe/get.rs b/smartcontract/cli/src/geolocation/probe/get.rs index 65017c538..a2c8094ba 100644 --- a/smartcontract/cli/src/geolocation/probe/get.rs +++ b/smartcontract/cli/src/geolocation/probe/get.rs @@ -125,6 +125,7 @@ mod tests { parent_devices, metrics_publisher_pk: metrics_pk, reference_count: 0, + target_update_count: 0, } } diff --git a/smartcontract/cli/src/geolocation/probe/list.rs b/smartcontract/cli/src/geolocation/probe/list.rs index bba41bcfe..5280b805c 100644 --- a/smartcontract/cli/src/geolocation/probe/list.rs +++ b/smartcontract/cli/src/geolocation/probe/list.rs @@ -105,6 +105,7 @@ mod tests { parent_devices: vec![], metrics_publisher_pk: Pubkey::new_unique(), reference_count: 0, + target_update_count: 0, }; let mut probes = HashMap::new(); @@ -144,6 +145,7 @@ mod tests { parent_devices: vec![parent_pk], metrics_publisher_pk: Pubkey::new_unique(), reference_count: 2, + target_update_count: 0, }; let mut probes = HashMap::new(); diff --git a/smartcontract/cli/src/geolocation/user/add_target.rs b/smartcontract/cli/src/geolocation/user/add_target.rs index 7fdc799da..5309693f6 100644 --- a/smartcontract/cli/src/geolocation/user/add_target.rs +++ b/smartcontract/cli/src/geolocation/user/add_target.rs @@ -138,6 +138,7 @@ mod tests { parent_devices: vec![], metrics_publisher_pk: Pubkey::new_unique(), reference_count: 0, + target_update_count: 0, } } diff --git a/smartcontract/cli/src/geolocation/user/remove_target.rs b/smartcontract/cli/src/geolocation/user/remove_target.rs index 39facac8b..7cb1c08e1 100644 --- a/smartcontract/cli/src/geolocation/user/remove_target.rs +++ b/smartcontract/cli/src/geolocation/user/remove_target.rs @@ -89,6 +89,7 @@ mod tests { parent_devices: vec![], metrics_publisher_pk: Pubkey::new_unique(), reference_count: 1, + target_update_count: 0, } } diff --git a/smartcontract/programs/doublezero-geolocation/src/processors/geo_probe/create.rs b/smartcontract/programs/doublezero-geolocation/src/processors/geo_probe/create.rs index 931ca7f81..727d92a98 100644 --- a/smartcontract/programs/doublezero-geolocation/src/processors/geo_probe/create.rs +++ b/smartcontract/programs/doublezero-geolocation/src/processors/geo_probe/create.rs @@ -100,6 +100,7 @@ pub fn process_create_geo_probe( reference_count: 0, code, parent_devices: vec![], + target_update_count: 0, }; try_acc_create( diff --git a/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/add_target.rs b/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/add_target.rs index 959add922..9257511be 100644 --- a/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/add_target.rs +++ b/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/add_target.rs @@ -110,6 +110,7 @@ pub fn process_add_target( }); probe.reference_count = probe.reference_count.saturating_add(1); + probe.target_update_count = probe.target_update_count.wrapping_add(1); try_acc_write(&user, user_account, payer_account, accounts)?; try_acc_write(&probe, probe_account, payer_account, accounts)?; diff --git a/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/remove_target.rs b/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/remove_target.rs index a94b4c74f..88284c859 100644 --- a/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/remove_target.rs +++ b/smartcontract/programs/doublezero-geolocation/src/processors/geolocation_user/remove_target.rs @@ -92,6 +92,7 @@ pub fn process_remove_target( user.targets.swap_remove(index); probe.reference_count = probe.reference_count.saturating_sub(1); + probe.target_update_count = probe.target_update_count.wrapping_add(1); try_acc_write(&user, user_account, payer_account, accounts)?; try_acc_write(&probe, probe_account, payer_account, accounts)?; diff --git a/smartcontract/programs/doublezero-geolocation/src/state/geo_probe.rs b/smartcontract/programs/doublezero-geolocation/src/state/geo_probe.rs index 81f78d62e..c4ab28333 100644 --- a/smartcontract/programs/doublezero-geolocation/src/state/geo_probe.rs +++ b/smartcontract/programs/doublezero-geolocation/src/state/geo_probe.rs @@ -1,11 +1,12 @@ use crate::state::accounttype::AccountType; -use borsh::{BorshDeserialize, BorshSerialize}; +use borsh::BorshSerialize; +use borsh_incremental::BorshDeserializeIncremental; use solana_program::{account_info::AccountInfo, msg, program_error::ProgramError, pubkey::Pubkey}; use std::{fmt, net::Ipv4Addr}; pub const MAX_PARENT_DEVICES: usize = 5; -#[derive(BorshSerialize, BorshDeserialize, Debug, PartialEq, Clone)] +#[derive(BorshSerialize, BorshDeserializeIncremental, Debug, PartialEq, Clone)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub struct GeoProbe { pub account_type: AccountType, // 1 @@ -25,7 +26,8 @@ pub struct GeoProbe { ) )] pub exchange_pk: Pubkey, // 32 - pub public_ip: Ipv4Addr, // 4 + #[incremental(default = Ipv4Addr::UNSPECIFIED)] + pub public_ip: Ipv4Addr, // 4 pub location_offset_port: u16, // 2 #[cfg_attr( feature = "serde", @@ -41,6 +43,8 @@ pub struct GeoProbe { // Variable-length fields must be at the end for Borsh deserialization pub code: String, // 4 + len pub parent_devices: Vec, // 4 + 32 * len + #[incremental(default = 0)] + pub target_update_count: u32, // 4 } impl fmt::Display for GeoProbe { @@ -48,37 +52,29 @@ impl fmt::Display for GeoProbe { write!( f, "account_type: {}, owner: {}, exchange_pk: {}, public_ip: {}, location_offset_port: {}, \ - metrics_publisher_pk: {}, reference_count: {}, code: {}, parent_devices: {:?}", + metrics_publisher_pk: {}, reference_count: {}, code: {}, parent_devices: {:?}, \ + target_update_count: {}", self.account_type, self.owner, self.exchange_pk, self.public_ip, self.location_offset_port, self.metrics_publisher_pk, self.reference_count, self.code, self.parent_devices, + self.target_update_count, ) } } -impl TryFrom<&[u8]> for GeoProbe { - type Error = ProgramError; - - fn try_from(mut data: &[u8]) -> Result { - let out = Self::deserialize(&mut data).map_err(|_| ProgramError::InvalidAccountData)?; - - if out.account_type != AccountType::GeoProbe { - return Err(ProgramError::InvalidAccountData); - } - - Ok(out) - } -} - impl TryFrom<&AccountInfo<'_>> for GeoProbe { type Error = ProgramError; fn try_from(account: &AccountInfo) -> Result { let data = account.try_borrow_data()?; - let res = Self::try_from(&data[..]); - if res.is_err() { - msg!("Failed to deserialize GeoProbe: {:?}", res.as_ref().err()); + let probe = Self::try_from(&data[..]).map_err(|e| { + msg!("Failed to deserialize GeoProbe: {}", e); + ProgramError::InvalidAccountData + })?; + if probe.account_type != AccountType::GeoProbe { + msg!("Invalid account type: {}", probe.account_type); + return Err(ProgramError::InvalidAccountData); } - res + Ok(probe) } } @@ -98,6 +94,7 @@ mod tests { reference_count: 3, code: "probe-ams-01".to_string(), parent_devices: vec![Pubkey::new_unique(), Pubkey::new_unique()], + target_update_count: 7, }; let data = borsh::to_vec(&val).unwrap(); @@ -111,10 +108,36 @@ mod tests { ); } + #[test] + fn test_state_geo_probe_backward_compat_without_target_update_count() { + let old = GeoProbe { + account_type: AccountType::GeoProbe, + owner: Pubkey::new_unique(), + exchange_pk: Pubkey::new_unique(), + public_ip: [8, 8, 8, 8].into(), + location_offset_port: 4242, + metrics_publisher_pk: Pubkey::new_unique(), + reference_count: 3, + code: "probe-ams-01".to_string(), + parent_devices: vec![Pubkey::new_unique()], + target_update_count: 0, + }; + + // Serialize, then truncate the trailing target_update_count (4 bytes) to simulate old data. + let mut data = borsh::to_vec(&old).unwrap(); + data.truncate(data.len() - 4); + + let deserialized = GeoProbe::try_from(&data[..]).unwrap(); + assert_eq!(deserialized.target_update_count, 0); + assert_eq!(deserialized.parent_devices, old.parent_devices); + } + #[test] fn test_state_geo_probe_try_from_invalid_account_type() { let data = [AccountType::None as u8]; let result = GeoProbe::try_from(&data[..]); - assert_eq!(result.unwrap_err(), ProgramError::InvalidAccountData); + // BorshDeserializeIncremental successfully deserializes but with wrong account type + let probe = result.unwrap(); + assert_eq!(probe.account_type, AccountType::None); } } diff --git a/smartcontract/programs/doublezero-geolocation/tests/geo_probe_test.rs b/smartcontract/programs/doublezero-geolocation/tests/geo_probe_test.rs index dcca32229..223e118a4 100644 --- a/smartcontract/programs/doublezero-geolocation/tests/geo_probe_test.rs +++ b/smartcontract/programs/doublezero-geolocation/tests/geo_probe_test.rs @@ -75,6 +75,7 @@ async fn test_create_geo_probe_success() { reference_count: 0, code: code.to_string(), parent_devices: vec![], + target_update_count: 0, }; assert_eq!(probe, expected_probe); @@ -265,6 +266,7 @@ async fn test_update_geo_probe_success() { reference_count: 0, // Unchanged code: code.to_string(), // Immutable parent_devices: vec![], // Unchanged + target_update_count: 0, // Unchanged }; assert_eq!(probe, expected_probe); diff --git a/smartcontract/programs/doublezero-geolocation/tests/geolocation_user_test.rs b/smartcontract/programs/doublezero-geolocation/tests/geolocation_user_test.rs index 9dc9131c2..11e07652f 100644 --- a/smartcontract/programs/doublezero-geolocation/tests/geolocation_user_test.rs +++ b/smartcontract/programs/doublezero-geolocation/tests/geolocation_user_test.rs @@ -611,10 +611,11 @@ async fn test_add_target_outbound_success() { assert_eq!(user.targets[0].ip_address, Ipv4Addr::new(8, 8, 8, 8)); assert_eq!(user.targets[0].geoprobe_pk, probe_pda); - // Verify probe reference_count incremented + // Verify probe reference_count and target_target_update_count incremented let probe_account = banks_client.get_account(probe_pda).await.unwrap().unwrap(); let probe = GeoProbe::try_from(&probe_account.data[..]).unwrap(); assert_eq!(probe.reference_count, 1); + assert_eq!(probe.target_update_count, 1); // Also add an inbound target to verify both target types work let inbound_target_pk = Pubkey::new_unique(); @@ -647,6 +648,7 @@ async fn test_add_target_outbound_success() { let probe_account = banks_client.get_account(probe_pda).await.unwrap().unwrap(); let probe = GeoProbe::try_from(&probe_account.data[..]).unwrap(); assert_eq!(probe.reference_count, 2); + assert_eq!(probe.target_update_count, 2); } #[tokio::test] @@ -869,10 +871,11 @@ async fn test_remove_target_success() { let user = GeolocationUser::try_from(&account.data[..]).unwrap(); assert!(user.targets.is_empty()); - // Verify reference_count decremented + // Verify reference_count decremented and target_update_count incremented let probe_account = banks_client.get_account(probe_pda).await.unwrap().unwrap(); let probe = GeoProbe::try_from(&probe_account.data[..]).unwrap(); assert_eq!(probe.reference_count, 0); + assert_eq!(probe.target_update_count, 2); // 1 from add + 1 from remove } #[tokio::test] diff --git a/smartcontract/sdk/rs/src/geolocation/geo_probe/get.rs b/smartcontract/sdk/rs/src/geolocation/geo_probe/get.rs index e10bf7c71..620e55c86 100644 --- a/smartcontract/sdk/rs/src/geolocation/geo_probe/get.rs +++ b/smartcontract/sdk/rs/src/geolocation/geo_probe/get.rs @@ -48,6 +48,7 @@ mod tests { parent_devices: vec![], metrics_publisher_pk: Pubkey::new_unique(), reference_count: 0, + target_update_count: 0, } } diff --git a/smartcontract/sdk/rs/src/geolocation/geo_probe/list.rs b/smartcontract/sdk/rs/src/geolocation/geo_probe/list.rs index 332da7a18..97676f603 100644 --- a/smartcontract/sdk/rs/src/geolocation/geo_probe/list.rs +++ b/smartcontract/sdk/rs/src/geolocation/geo_probe/list.rs @@ -65,6 +65,7 @@ mod tests { parent_devices: vec![], metrics_publisher_pk: Pubkey::new_unique(), reference_count: 0, + target_update_count: 0, } }