diff --git a/tms/heartbeat/go.mod b/tms/heartbeat/go.mod index 4c930170d41808c54873baf641d2f5625710979c..b1cde07d17ec8075814c1f8aa99b04bd7883f942 100644 --- a/tms/heartbeat/go.mod +++ b/tms/heartbeat/go.mod @@ -9,6 +9,7 @@ replace gitee.com/OpenCloudOS/ocmanager/tms/proto => ../proto require ( gitee.com/OpenCloudOS/ocmanager/tms/pkg v0.0.0-00010101000000-000000000000 gitee.com/OpenCloudOS/ocmanager/tms/proto v0.0.0-00010101000000-000000000000 + github.com/alicebob/miniredis/v2 v2.37.0 github.com/redis/go-redis/v9 v9.19.0 gopkg.in/yaml.v3 v3.0.1 trpc.group/trpc-go/trpc-go v1.0.3 @@ -35,6 +36,7 @@ require ( github.com/spf13/cast v1.3.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.43.0 // indirect + github.com/yuin/gopher-lua v1.1.1 // indirect go.uber.org/atomic v1.11.0 // indirect go.uber.org/automaxprocs v1.3.0 // indirect go.uber.org/multierr v1.10.0 // indirect diff --git a/tms/heartbeat/go.sum b/tms/heartbeat/go.sum index 9ca4d03c848d36927645f52263388df9c38b1091..62f32d9b5455328e53e635b50cdb305751716c84 100644 --- a/tms/heartbeat/go.sum +++ b/tms/heartbeat/go.sum @@ -1,5 +1,7 @@ github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= +github.com/alicebob/miniredis/v2 v2.37.0 h1:RheObYW32G1aiJIj81XVt78ZHJpHonHLHW7OLIshq68= +github.com/alicebob/miniredis/v2 v2.37.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= github.com/andybalholm/brotli v1.0.4 h1:V7DdXeJtZscaqfNuAdSRuRFzuiKlHSC/Zh3zl9qY3JY= github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= @@ -79,6 +81,8 @@ github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyC github.com/valyala/fasthttp v1.43.0 h1:Gy4sb32C98fbzVWZlTM1oTMdLWGyvxR03VhM6cBIU4g= github.com/valyala/fasthttp v1.43.0/go.mod h1:f6VbjjoI3z1NDOZOv17o6RvtRSWxC77seBFc2uWtgiY= github.com/valyala/tcplisten v1.0.0/go.mod h1:T0xQ8SeCZGxckz9qRXTfG43PvQ/mcWh7FwZEA7Ioqkc= +github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= +github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= diff --git a/tms/heartbeat/internal/service/heartbeat_svc.go b/tms/heartbeat/internal/service/heartbeat_svc.go index 0b74c62be9d779d3663106307a6db1ffc477edd0..4e2ec23077da919b3ae744d3f96f6b741af87c85 100644 --- a/tms/heartbeat/internal/service/heartbeat_svc.go +++ b/tms/heartbeat/internal/service/heartbeat_svc.go @@ -82,6 +82,7 @@ func (s *HeartbeatService) QueryHeartbeat(ctx context.Context, req *pb.QueryHear } hbState := state.Evaluate(time.Unix(data.Timestamp, 0)) + hbState = s.applyAssetStatus(ctx, data.AgentID, hbState) return &pb.QueryHeartbeatResponse{ AgentId: data.AgentID, State: toProtoState(hbState), @@ -108,6 +109,7 @@ func (s *HeartbeatService) BatchQueryHeartbeat(ctx context.Context, req *pb.Batc continue } hbState := state.Evaluate(time.Unix(data.Timestamp, 0)) + hbState = s.applyAssetStatus(ctx, data.AgentID, hbState) results = append(results, &pb.QueryHeartbeatResponse{ AgentId: data.AgentID, State: toProtoState(hbState), @@ -119,6 +121,26 @@ func (s *HeartbeatService) BatchQueryHeartbeat(ctx context.Context, req *pb.Batc return &pb.BatchQueryHeartbeatResponse{Results: results}, nil } +func (s *HeartbeatService) applyAssetStatus(ctx context.Context, agentID uint64, hbState state.State) state.State { + proxy := newAssetSvcClientProxy( + client.WithServiceName("trpc.ocmanager.asset.AssetSvc"), + ) + resp, err := proxy.QueryAsset(ctx, &assetpb.QueryAssetRequest{AgentId: agentID}) + if err != nil { + log.ErrorContextf(ctx, "query asset status failed, agent_id: %d, err: %v", agentID, err) + return state.Offline + } + for _, asset := range resp.GetAssets() { + if asset.GetAgentId() == agentID { + if asset.GetStatus() != "allowed" { + return state.Offline + } + return hbState + } + } + return state.Offline +} + func (s *HeartbeatService) renewCert(ctx context.Context, agentID uint64) ([]byte, []byte) { // 通过 tRPC ClientProxy 调用 AssetSvc。 proxy := newAssetSvcClientProxy( diff --git a/tms/heartbeat/internal/service/heartbeat_svc_test.go b/tms/heartbeat/internal/service/heartbeat_svc_test.go index ddfd350f494bc4025899f9eb045602b720cc7528..5a43b7a32373f8779188659271d305fcde08853a 100644 --- a/tms/heartbeat/internal/service/heartbeat_svc_test.go +++ b/tms/heartbeat/internal/service/heartbeat_svc_test.go @@ -5,15 +5,23 @@ package service import ( "context" + "errors" "testing" + "time" + "github.com/alicebob/miniredis/v2" "trpc.group/trpc-go/trpc-go/client" + "gitee.com/OpenCloudOS/ocmanager/tms/heartbeat/internal/state" + "gitee.com/OpenCloudOS/ocmanager/tms/heartbeat/internal/store" assetpb "gitee.com/OpenCloudOS/ocmanager/tms/proto/asset" + commonpb "gitee.com/OpenCloudOS/ocmanager/tms/proto/common" + hbpb "gitee.com/OpenCloudOS/ocmanager/tms/proto/heartbeat" ) type fakeAssetSvcClient struct { renew func(context.Context, *assetpb.RenewCertRequest, ...client.Option) (*assetpb.RenewCertResponse, error) + query func(context.Context, *assetpb.QueryAssetRequest, ...client.Option) (*assetpb.QueryAssetResponse, error) } func (f *fakeAssetSvcClient) Register(ctx context.Context, req *assetpb.RegisterRequest, opts ...client.Option) (*assetpb.RegisterResponse, error) { @@ -32,6 +40,9 @@ func (f *fakeAssetSvcClient) RevokeCert(ctx context.Context, req *assetpb.Revoke } func (f *fakeAssetSvcClient) QueryAsset(ctx context.Context, req *assetpb.QueryAssetRequest, opts ...client.Option) (*assetpb.QueryAssetResponse, error) { + if f.query != nil { + return f.query(ctx, req, opts...) + } return &assetpb.QueryAssetResponse{}, nil } @@ -60,3 +71,62 @@ func TestRenewCertNilResponseDoesNotPanic(t *testing.T) { t.Fatalf("renewCert returned cert=%q key=%q, want nil/nil", cert, key) } } + +func TestBlockedAssetIsOfflineInHeartbeatQueries(t *testing.T) { + mr := miniredis.RunT(t) + hbStore := store.New(mr.Addr(), "", 0) + t.Cleanup(func() { _ = hbStore.Close() }) + ctx := context.Background() + const agentID uint64 = 1001 + if err := hbStore.Set(ctx, &store.HeartbeatData{ + AgentID: agentID, Timestamp: time.Now().Unix(), + }); err != nil { + t.Fatalf("store heartbeat: %v", err) + } + + oldNewClient := newAssetSvcClientProxy + newAssetSvcClientProxy = func(opts ...client.Option) assetpb.AssetSvcClientProxy { + return &fakeAssetSvcClient{ + query: func(_ context.Context, req *assetpb.QueryAssetRequest, _ ...client.Option) (*assetpb.QueryAssetResponse, error) { + return &assetpb.QueryAssetResponse{Assets: []*assetpb.AssetRecord{{ + AgentId: req.AgentId, Status: "blocked", + }}}, nil + }, + } + } + t.Cleanup(func() { newAssetSvcClientProxy = oldNewClient }) + + svc := New(hbStore) + single, err := svc.QueryHeartbeat(ctx, &hbpb.QueryHeartbeatRequest{AgentId: agentID}) + if err != nil { + t.Fatalf("QueryHeartbeat: %v", err) + } + if single.State != commonpb.HeartbeatState_STATE_OFFLINE { + t.Fatalf("QueryHeartbeat state = %v, want OFFLINE for blocked asset", single.State) + } + + batch, err := svc.BatchQueryHeartbeat(ctx, &hbpb.BatchQueryHeartbeatRequest{AgentIds: []uint64{agentID}}) + if err != nil { + t.Fatalf("BatchQueryHeartbeat: %v", err) + } + if len(batch.Results) != 1 || batch.Results[0].State != commonpb.HeartbeatState_STATE_OFFLINE { + t.Fatalf("BatchQueryHeartbeat results = %+v, want blocked asset OFFLINE", batch.Results) + } +} + +func TestApplyAssetStatusIsOfflineWhenAssetQueryFails(t *testing.T) { + oldNewClient := newAssetSvcClientProxy + newAssetSvcClientProxy = func(opts ...client.Option) assetpb.AssetSvcClientProxy { + return &fakeAssetSvcClient{ + query: func(context.Context, *assetpb.QueryAssetRequest, ...client.Option) (*assetpb.QueryAssetResponse, error) { + return nil, errors.New("asset service unavailable") + }, + } + } + t.Cleanup(func() { newAssetSvcClientProxy = oldNewClient }) + + svc := &HeartbeatService{} + if got := svc.applyAssetStatus(context.Background(), 1001, state.Online); got != state.Offline { + t.Fatalf("applyAssetStatus() = %v, want %v when asset query fails", got, state.Offline) + } +}