diff --git a/fleetspeak/src/client/daemonservice/daemonservice_test.go b/fleetspeak/src/client/daemonservice/daemonservice_test.go index 9845cb4a..d20791bf 100644 --- a/fleetspeak/src/client/daemonservice/daemonservice_test.go +++ b/fleetspeak/src/client/daemonservice/daemonservice_test.go @@ -28,6 +28,7 @@ import ( "github.com/google/fleetspeak/fleetspeak/src/client/clitesting" "github.com/google/fleetspeak/fleetspeak/src/client/service" "github.com/google/fleetspeak/fleetspeak/src/client/stats" + "github.com/google/fleetspeak/fleetspeak/src/common" "github.com/google/fleetspeak/fleetspeak/src/common/anypbtest" dspb "github.com/google/fleetspeak/fleetspeak/src/client/daemonservice/proto/fleetspeak_daemonservice" @@ -588,6 +589,59 @@ waitLoop: } } +func TestClientIDEnvVar(t *testing.T) { + cid, err := common.StringToClientID("0102030405060708") + if err != nil { + t.Fatalf("Failed to create ClientID: %v", err) + } + sc := clitesting.FakeServiceContext{ + OutChan: make(chan *fspb.Message, 100), + LocalInfo: &service.LocalInfo{ + ClientID: cid, + }, + } + + bashCmd := "echo $FLEETSPEAK_CLIENT_ID" + + cfg := &fspb.ClientServiceConfig{ + Name: "TestDaemonService", + Config: anypbtest.New(t, &dspb.Config{ + Argv: []string{"bash", "-c", bashCmd}, + StdParams: &dspb.Config_StdParams{ + ServiceName: "test_service", + FlushBytes: 1024, + }, + }), + } + + s, err := Factory(cfg) + if err != nil { + t.Fatalf("Factory(...): %v", err) + } + if err := s.Start(&sc); err != nil { + t.Fatalf("Service.Start(...): %v", err) + } + defer s.Stop() + + for { + select { + case <-time.After(time.Second): + t.Fatal("Timeout waiting for StdOutput") + case m := <-sc.OutChan: + if m.MessageType == "StdOutput" { + var out dspb.StdOutputData + if err := m.Data.UnmarshalTo(&out); err != nil { + t.Fatalf("UnmarshalTo failed: %v", err) + } + if got, want := string(out.Stdout), "0102030405060708\n"; got != want { + t.Errorf("Unexpected output: got %q, want %q", got, want) + } + return + } + } + } +} + func init() { // Hack so that daemonservice logs get written to the right place. if d := os.Getenv("TEST_UNDECLARED_OUTPUTS_DIR"); d != "" { diff --git a/fleetspeak/src/client/daemonservice/execution/execution.go b/fleetspeak/src/client/daemonservice/execution/execution.go index 013d9c8b..69e036d6 100644 --- a/fleetspeak/src/client/daemonservice/execution/execution.go +++ b/fleetspeak/src/client/daemonservice/execution/execution.go @@ -193,6 +193,10 @@ func New(ctx context.Context, daemonServiceName string, cfg *dspb.Config, sc ser ret.samplePeriod = time.Duration(cfg.ResourceMonitoringSamplePeriodSeconds) * time.Second } + if info := sc.GetLocalInfo(); info != nil { + ret.cmd.AddEnvVar("FLEETSPEAK_CLIENT_ID=" + info.ClientID.String()) + } + var err error ret.channel, err = ret.cmd.SetupCommsChannel() if err != nil {