Randomize initial URL (#1009)

When clients are configured with multiple frontends, the initial url should be randomly chosen so they randomly distribute across frontends.
Velocidex · Apr 6, 2021 · 5b29c3f · 5b29c3f
1 parent 4c14162
commit 5b29c3f
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 15 deletions.
diff --git a/actions/events_test.go b/actions/events_test.go
@@ -25,6 +25,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/services/launcher"
 	"www.velocidex.com/golang/velociraptor/services/notifications"
 	"www.velocidex.com/golang/velociraptor/services/repository"
+	"www.velocidex.com/golang/velociraptor/utils"
 	"www.velocidex.com/golang/velociraptor/vtesting"
 )
 
@@ -49,6 +50,8 @@ type EventsTestSuite struct {
 	sm         *services.Service
 	responder  *responder.Responder
 	writeback  string
+
+	Clock utils.Clock
 }
 
 func (self *EventsTestSuite) SetupTest() {
@@ -58,6 +61,7 @@ func (self *EventsTestSuite) SetupTest() {
 	self.config_obj = config.GetDefaultConfig()
 	self.config_obj.Frontend.DoNotCompressArtifacts = true
 	self.config_obj.Datastore.Implementation = "Test"
+	self.Clock = &utils.IncClock{}
 
 	tmpfile, err := ioutil.TempFile("", "")
 	require.NoError(t, err)
@@ -124,6 +128,7 @@ var server_state = &flows_proto.ClientEventTable{
 
 func (self *EventsTestSuite) TestEventTableUpdate() {
 	client_manager := services.ClientEventManager()
+	client_manager.(*client_monitoring.ClientEventTable).Clock = self.Clock
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
 	defer cancel()
@@ -177,6 +182,8 @@ func (self *EventsTestSuite) TestEventTableUpdate() {
 	// be the same as the old one, except the version will be
 	// advanced.
 	label_manager := services.GetLabeler()
+	label_manager.(*labels.Labeler).Clock = self.Clock
+
 	require.NoError(self.T(),
 		label_manager.SetClientLabel(self.config_obj, self.client_id,
 			"Foobar"))

diff --git a/go.mod b/go.mod
@@ -98,6 +98,7 @@ require (
 	golang.org/x/net v0.0.0-20210119194325-5f4716e94777
 	golang.org/x/oauth2 v0.0.0-20210220000619-9bb904979d93
 	golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c
+	golang.org/x/text v0.3.6 // indirect
 	golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba // indirect
 	golang.org/x/tools v0.0.0-20200828161849-5deb26317202 // indirect
 	google.golang.org/api v0.30.0
@@ -118,7 +119,7 @@ require (
 	www.velocidex.com/golang/go-prefetch v0.0.0-20200722101157-37e4751dd5ca
 	www.velocidex.com/golang/oleparse v0.0.0-20190327031422-34195d413196
 	www.velocidex.com/golang/regparser v0.0.0-20190625082115-b02dc43c2500
-	www.velocidex.com/golang/vfilter v0.0.0-20210220121641-879064f4499e
+	www.velocidex.com/golang/vfilter v0.0.0-20210406162709-d40800885aff
 	www.velocidex.com/golang/vtypes v0.0.0-20210323032031-b61f37170666
 )
 

diff --git a/go.sum b/go.sum
@@ -712,6 +712,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.4/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.5 h1:i6eZZ+zk0SOf0xgBpEpPD18qWcJda6q1sxt3S0kzyUQ=
 golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.3.6 h1:aRYxNxv6iGQlyVaZmk6ZgYEDa+Jg18DxebPSrd6bg1M=
+golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
@@ -916,5 +918,7 @@ www.velocidex.com/golang/regparser v0.0.0-20190625082115-b02dc43c2500/go.mod h1:
 www.velocidex.com/golang/vfilter v0.0.0-20210108051106-c18c13c24eff/go.mod h1:EdP5LDT3l9khZVjDVD2YKoDvECi3AW0e0074quDPNpA=
 www.velocidex.com/golang/vfilter v0.0.0-20210220121641-879064f4499e h1:h7uErYK4QOhsxTlle+6jxXgjCIwZ/wx44h9397fd3TM=
 www.velocidex.com/golang/vfilter v0.0.0-20210220121641-879064f4499e/go.mod h1:KB724xBNYh4lgipyGwsvx0/5hXRqsKjmrMrkSjGESvU=
+www.velocidex.com/golang/vfilter v0.0.0-20210406162709-d40800885aff h1:O+DyRkA83hZMKBzX8FSLo/BzRmvqoMj9W2fYLkg8BaM=
+www.velocidex.com/golang/vfilter v0.0.0-20210406162709-d40800885aff/go.mod h1:KB724xBNYh4lgipyGwsvx0/5hXRqsKjmrMrkSjGESvU=
 www.velocidex.com/golang/vtypes v0.0.0-20210323032031-b61f37170666 h1:4VjIpQYv3WWXizcMMwAHi8hp5B6pbWPD1jsNxygWhRE=
 www.velocidex.com/golang/vtypes v0.0.0-20210323032031-b61f37170666/go.mod h1:34AZRfhNvJ1QAwPpYrDxjCyOFys+NbSmH6LLVSjsAEg=
diff --git a/http_comms/comms.go b/http_comms/comms.go
@@ -47,6 +47,8 @@ import (
 
 var (
 	RedirectError = errors.New("RedirectError")
+
+	Rand func(int) int = rand.Intn
 )
 
 // Responsible for maybe enrolling the client. Enrollments should not
@@ -194,6 +196,11 @@ func NewHTTPConnector(
 		logger:     logger,
 		clock:      clock,
 
+		// Start with a random URL from the set of
+		// preconfigured URLs. This should distribute clients
+		// randomly to all frontends.
+		current_url_idx: Rand(len(urls)),
+
 		minPoll:    time.Duration(1) * time.Second,
 		maxPoll:    time.Duration(max_poll) * time.Second,
 		maxPollDev: maxPollDev,
@@ -310,7 +317,7 @@ func (self *HTTPConnector) Post(handler string, data []byte, urgent bool) (
 			// For safety we wait after redirect in case we end up
 			// in a redirect loop.
 			wait := self.maxPoll + time.Duration(
-				rand.Intn(int(self.maxPollDev)))*time.Second
+				Rand(int(self.maxPollDev)))*time.Second
 			self.logger.Info("Waiting after redirect: %v", wait)
 			<-self.clock.After(wait)
 		}
@@ -352,7 +359,7 @@ func (self *HTTPConnector) advanceToNextServer() {
 	// sleep to back off.
 	if self.current_url_idx == self.last_success_idx {
 		wait := self.maxPoll + time.Duration(
-			rand.Intn(int(self.maxPollDev)))*time.Second
+			Rand(int(self.maxPollDev)))*time.Second
 
 		self.logger.Info(
 			"Waiting for a reachable server: %v", wait)
@@ -542,7 +549,7 @@ func (self *NotificationReader) sendMessageList(
 			// Add random wait between polls to avoid
 			// synchronization of endpoints.
 			wait := self.maxPoll + time.Duration(
-				rand.Intn(int(self.maxPollDev)))*time.Second
+				Rand(int(self.maxPollDev)))*time.Second
 			self.logger.Info("Sleeping for %v", wait)
 
 			select {

diff --git a/http_comms/comms_test.go b/http_comms/comms_test.go
@@ -157,6 +157,9 @@ func (self *CommsTestSuite) SetupTest() {
 	cm := &crypto.NullCryptoManager{}
 	self.empty_response, _ = cm.EncryptMessageList(
 		&crypto_proto.MessageList{}, "C.1234")
+
+	// Disable randomness for the test.
+	Rand = func(int) int { return 0 }
 }
 
 func (self *CommsTestSuite) TearDownTest() {

diff --git a/services/client_monitoring/client_monitoring.go b/services/client_monitoring/client_monitoring.go
@@ -64,7 +64,7 @@ type ClientEventTable struct {
 	// protobufs in memory.
 	state *flows_proto.ClientEventTable
 
-	clock utils.Clock
+	Clock utils.Clock
 
 	id string
 }
@@ -188,7 +188,7 @@ func (self *ClientEventTable) setClientMonitoringState(
 	}
 
 	self.state = state
-	state.Version = uint64(self.clock.Now().UnixNano())
+	state.Version = uint64(self.Clock.Now().UnixNano())
 
 	// Store the new table in the data store.
 	db, err := datastore.GetDB(config_obj)
@@ -231,7 +231,7 @@ func (self *ClientEventTable) GetClientUpdateEventTableMessage(
 	self.mu.Unlock()
 
 	result := &actions_proto.VQLEventTable{
-		Version: uint64(self.clock.Now().UnixNano()),
+		Version: uint64(self.Clock.Now().UnixNano()),
 	}
 
 	if state.Artifacts == nil {
@@ -380,7 +380,7 @@ func StartClientMonitoringService(
 	config_obj *config_proto.Config) error {
 
 	event_table := &ClientEventTable{
-		clock: &utils.RealClock{},
+		Clock: &utils.RealClock{},
 		id:    uuid.New().String(),
 	}
 	services.RegisterClientEventManager(event_table)

diff --git a/services/client_monitoring/client_monitoring_test.go b/services/client_monitoring/client_monitoring_test.go
@@ -80,7 +80,7 @@ sources:
 `, "")
 
 	manager := services.ClientEventManager().(*ClientEventTable)
-	manager.clock = current_clock
+	manager.Clock = current_clock
 
 	err := manager.SetClientMonitoringState(context.Background(), self.config_obj,
 		&flows_proto.ClientEventTable{
@@ -120,7 +120,7 @@ sources:
 `, "")
 
 	manager := services.ClientEventManager().(*ClientEventTable)
-	manager.clock = current_clock
+	manager.Clock = current_clock
 
 	// Set the initial table.
 	err := manager.SetClientMonitoringState(context.Background(), self.config_obj,
@@ -163,7 +163,7 @@ sources:
 `, "")
 
 	manager1 := services.ClientEventManager().(*ClientEventTable)
-	manager1.clock = current_clock
+	manager1.Clock = current_clock
 
 	// Set the initial table.
 	err := manager1.SetClientMonitoringState(context.Background(),
@@ -180,7 +180,7 @@ sources:
 	// Now another frontend sets the client monitoring state
 	require.NoError(self.T(), self.sm.Start(StartClientMonitoringService))
 	manager2 := services.ClientEventManager().(*ClientEventTable)
-	manager2.clock = current_clock
+	manager2.Clock = current_clock
 
 	// Now update the monitoring state
 	err = manager2.SetClientMonitoringState(context.Background(),
@@ -209,7 +209,7 @@ func (self *ClientMonitoringTestSuite) TestClientMonitoringCompiling() {
 
 	// If no table exists, we will get a default table.
 	manager := services.ClientEventManager().(*ClientEventTable)
-	manager.clock = current_clock
+	manager.Clock = current_clock
 
 	// Install an initial monitoring table: Everyone gets ServiceCreation.
 	manager.SetClientMonitoringState(context.Background(),
@@ -324,7 +324,7 @@ func (self *ClientMonitoringTestSuite) TestClientMonitoringCompilingMultipleArti
 
 	// If no table exists, we will get a default table.
 	manager := services.ClientEventManager().(*ClientEventTable)
-	manager.clock = current_clock
+	manager.Clock = current_clock
 
 	// Install an initial monitoring table: Everyone gets ServiceCreation.
 	manager.SetClientMonitoringState(context.Background(),
@@ -377,7 +377,7 @@ func (self *ClientMonitoringTestSuite) TestClientMonitoring() {
 
 	// If no table exists, we will get a default table.
 	manager := services.ClientEventManager().(*ClientEventTable)
-	manager.clock = current_clock
+	manager.Clock = current_clock
 
 	test_utils.GetMemoryDataStore(self.T(), self.config_obj).Clear()
 	assert.NoError(self.T(), manager.LoadFromFile(context.Background(), self.config_obj))