[management] Confine the activity writes to the reverse proxy

The user half reused nothing: SaveUserLastLogin already exists and is the
same call the dashboard and device login paths make, so the parallel
RefreshUserLastLogin is gone and the proxy uses the established one.

Reaching it no longer widens shared interfaces. The proxy service already
receives the store, narrowed to ProxyTokenChecker; that interface now
carries the two writes the proxy makes, so users.Manager, peers.Manager and
Peer are untouched and the exclusion predicate moved into the proxy package
next to its only caller.

RefreshPeerLastSeen stays on the store because nothing there fits:
SavePeerStatus rewrites the connected flag and session token from a caller
snapshot, which would race the sync stream that owns them.
This commit is contained in:
mlsmaycon
2026-08-09 08:11:30 +00:00
parent 48d9161056
commit d2f93fcd90
12 changed files with 71 additions and 195 deletions
+33 -9
View File
@@ -61,6 +61,22 @@ type ProxyTokenChecker interface {
IsProxyAccessTokenValid(ctx context.Context, tokenID string) (bool, error)
}
// ProxyStore is the slice of the management store this service reaches
// directly. It is declared here, and satisfied by the store, so the reverse
// proxy owns the surface it needs instead of widening manager interfaces the
// rest of management shares.
type ProxyStore interface {
ProxyTokenChecker
// SaveUserLastLogin is the same write the dashboard and device login paths
// use, reused here so a proxy SSO sign-in lands in the one place activity
// accounting reads.
SaveUserLastLogin(ctx context.Context, accountID, userID string, lastLogin time.Time) error
// RefreshPeerLastSeen stamps only LastSeen. SavePeerStatus is not usable
// here: it rewrites the connected flag and session token from a caller
// snapshot, which would race the sync stream that owns them.
RefreshPeerLastSeen(ctx context.Context, accountID, peerID string) error
}
// ProxyServiceServer implements the ProxyService gRPC server
// AgentNetworkSynthesizer produces in-memory reverse-proxy services from
// Agent Network provider/policy state for the proxy snapshot path; synthesised
@@ -118,7 +134,7 @@ type ProxyServiceServer struct {
tokenStore *OneTimeTokenStore
// Checker for proxy access token validity
tokenChecker ProxyTokenChecker
proxyStore ProxyStore
// OIDC configuration for proxy authentication
oidcConfig ProxyOIDCConfig
@@ -189,7 +205,7 @@ func enforceAccountScope(ctx context.Context, requestAccountID string) error {
}
// NewProxyServiceServer creates a new proxy service server.
func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, pkceStore *PKCEVerifierStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer {
func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, pkceStore *PKCEVerifierStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, proxyStore ProxyStore) *ProxyServiceServer {
ctx, cancel := context.WithCancel(context.Background())
s := &ProxyServiceServer{
accessLogManager: accessLogMgr,
@@ -200,7 +216,7 @@ func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeT
usersManager: usersManager,
idpManager: idpManager,
proxyManager: proxyMgr,
tokenChecker: tokenChecker,
proxyStore: proxyStore,
snapshotBatchSize: snapshotBatchSizeFromEnv(),
cancel: cancel,
}
@@ -695,8 +711,8 @@ func (s *ProxyServiceServer) heartbeat(ctx context.Context, conn *proxyConnectio
log.WithContext(ctx).Debugf("Failed to update proxy %s heartbeat: %v", p.ID, err)
}
if conn.tokenID != "" && s.tokenChecker != nil {
valid, err := s.tokenChecker.IsProxyAccessTokenValid(ctx, conn.tokenID)
if conn.tokenID != "" && s.proxyStore != nil {
valid, err := s.proxyStore.IsProxyAccessTokenValid(ctx, conn.tokenID)
if err != nil {
log.WithContext(ctx).Warnf("failed to check token validity for proxy %s: %v", conn.proxyID, err)
continue
@@ -1697,11 +1713,11 @@ func (s *ProxyServiceServer) GenerateSessionToken(ctx context.Context, domain, u
// have no interactive login to record. A failure is logged and dropped: the
// next login marks it again and no authorization decision reads it.
func (s *ProxyServiceServer) recordUserLogin(ctx context.Context, accountID string, user *types.User) {
if user.IsServiceUser {
if s.proxyStore == nil || user.IsServiceUser {
return
}
if err := s.usersManager.RefreshLastLogin(ctx, accountID, user.Id, time.Now().UTC()); err != nil {
if err := s.proxyStore.SaveUserLastLogin(ctx, accountID, user.Id, time.Now().UTC()); err != nil {
log.WithContext(ctx).Debugf("record proxy login for user %s: %v", user.Id, err)
}
}
@@ -2086,7 +2102,7 @@ const proxyPeerSeenInterval = time.Hour
// costs nothing. A failure is logged and dropped: the next request marks it
// again.
func (s *ProxyServiceServer) recordPeerSeen(ctx context.Context, accountID string, peer *peer.Peer) {
if !peer.CountsTowardActivity() {
if s.proxyStore == nil || !peerCountsTowardActivity(peer) {
return
}
@@ -2094,11 +2110,19 @@ func (s *ProxyServiceServer) recordPeerSeen(ctx context.Context, accountID strin
return
}
if err := s.peersManager.RefreshLastSeen(ctx, accountID, peer.ID); err != nil {
if err := s.proxyStore.RefreshPeerLastSeen(ctx, accountID, peer.ID); err != nil {
log.WithContext(ctx).Debugf("record proxy activity for peer %s: %v", peer.ID, err)
}
}
// peerCountsTowardActivity reports whether the peer represents a device a
// person actually runs. Embedded proxy peers are infrastructure and browser
// (WASM) clients are ephemeral sessions, so activity accounting ignores both
// and a write for them could never count.
func peerCountsTowardActivity(peer *peer.Peer) bool {
return !peer.ProxyMeta.Embedded && peer.Meta.KernelVersion != "wasm"
}
// resolvePeerOwner returns the user a peer is linked to, once per request so
// the status gate and the identity resolution below share a single lookup.
// Unlinked peers (machine agents) have no owner. A lookup that fails returns
@@ -123,20 +123,6 @@ type mockUsersManager struct {
users map[string]*types.User
err error
getUserCalls int
loginMarks []loginMark
}
// loginMark records a RefreshLastLogin call so tests can assert what the proxy
// wrote, rather than that it merely called something.
type loginMark struct {
accountID string
userID string
at time.Time
}
func (m *mockUsersManager) RefreshLastLogin(_ context.Context, accountID, userID string, loginAt time.Time) error {
m.loginMarks = append(m.loginMarks, loginMark{accountID: accountID, userID: userID, at: loginAt})
return nil
}
func (m *mockUsersManager) GetUser(ctx context.Context, userID string) (*types.User, error) {
@@ -168,17 +154,36 @@ type mockTunnelPeersManager struct {
peerErr error
groups []*types.Group
groupsErr error
seenMarks []seenMark
}
// seenMark records a RefreshLastSeen call. The timestamp is the database's, so
// there is nothing from the caller to assert beyond who was marked.
// mockProxyStore records the activity writes the proxy makes so tests can
// assert what was written, not merely that something was called.
type mockProxyStore struct {
loginMarks []loginMark
seenMarks []seenMark
}
type loginMark struct {
accountID string
userID string
at time.Time
}
type seenMark struct {
accountID string
peerID string
}
func (m *mockTunnelPeersManager) RefreshLastSeen(_ context.Context, accountID, peerID string) error {
func (m *mockProxyStore) IsProxyAccessTokenValid(_ context.Context, _ string) (bool, error) {
return true, nil
}
func (m *mockProxyStore) SaveUserLastLogin(_ context.Context, accountID, userID string, lastLogin time.Time) error {
m.loginMarks = append(m.loginMarks, loginMark{accountID: accountID, userID: userID, at: lastLogin})
return nil
}
func (m *mockProxyStore) RefreshPeerLastSeen(_ context.Context, accountID, peerID string) error {
m.seenMarks = append(m.seenMarks, seenMark{accountID: accountID, peerID: peerID})
return nil
}
@@ -820,14 +825,15 @@ func TestValidateTunnelPeerRecordsActivity(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
peersManager := &mockTunnelPeersManager{peer: tt.peer}
proxyStore := &mockProxyStore{}
server := &ProxyServiceServer{
proxyStore: proxyStore,
serviceManager: &mockReverseProxyManager{
proxiesByAccount: map[string][]*service.Service{
accountID: {{Domain: domain, AccountID: accountID}},
},
},
peersManager: peersManager,
peersManager: &mockTunnelPeersManager{peer: tt.peer},
usersManager: &mockUsersManager{users: map[string]*types.User{}},
}
@@ -840,12 +846,12 @@ func TestValidateTunnelPeerRecordsActivity(t *testing.T) {
require.True(t, resp.GetValid(), "peer should be granted access")
if !tt.expectMark {
assert.Empty(t, peersManager.seenMarks, "peer should not have been marked seen")
assert.Empty(t, proxyStore.seenMarks, "peer should not have been marked seen")
return
}
require.Len(t, peersManager.seenMarks, 1, "peer should have been marked seen exactly once")
mark := peersManager.seenMarks[0]
require.Len(t, proxyStore.seenMarks, 1, "peer should have been marked seen exactly once")
mark := proxyStore.seenMarks[0]
assert.Equal(t, accountID, mark.accountID, "activity must be recorded against the service account")
assert.Equal(t, peerID, mark.peerID, "activity must be recorded against the calling peer")
})
@@ -860,16 +866,17 @@ func TestValidateTunnelPeerDeniedRecordsNoActivity(t *testing.T) {
accountID = "account1"
)
peersManager := &mockTunnelPeersManager{
peer: &peer.Peer{ID: "peer1", Name: "agent", UserID: "user1", Status: &peer.PeerStatus{LastSeen: time.Now().Add(-3 * time.Hour)}},
}
proxyStore := &mockProxyStore{}
server := &ProxyServiceServer{
proxyStore: proxyStore,
serviceManager: &mockReverseProxyManager{
proxiesByAccount: map[string][]*service.Service{
accountID: {{Domain: domain, AccountID: accountID}},
},
},
peersManager: peersManager,
peersManager: &mockTunnelPeersManager{
peer: &peer.Peer{ID: "peer1", Name: "agent", UserID: "user1", Status: &peer.PeerStatus{LastSeen: time.Now().Add(-3 * time.Hour)}},
},
// The owner is blocked, so the tunnel gate denies before the mint.
usersManager: &mockUsersManager{users: map[string]*types.User{
"user1": {Id: "user1", AccountID: accountID, Blocked: true},
@@ -883,7 +890,7 @@ func TestValidateTunnelPeerDeniedRecordsNoActivity(t *testing.T) {
require.NoError(t, err)
require.False(t, resp.GetValid(), "blocked owner should be denied")
assert.Empty(t, peersManager.seenMarks, "a denied peer must not be marked seen")
assert.Empty(t, proxyStore.seenMarks, "a denied peer must not be marked seen")
}
func TestGetAccountProxyByDomain(t *testing.T) {
@@ -587,10 +587,6 @@ type testValidateSessionUsersManager struct {
store store.Store
}
func (m *testValidateSessionUsersManager) RefreshLastLogin(ctx context.Context, accountID, userID string, loginAt time.Time) error {
return m.store.RefreshUserLastLogin(ctx, accountID, userID, loginAt)
}
func (m *testValidateSessionUsersManager) GetUser(ctx context.Context, userID string) (*types.User, error) {
return m.store.GetUserByUserID(ctx, store.LockingStrengthNone, userID)
}