From e6f4dd54bb1111174119fadf2944bbc2c7824f3f Mon Sep 17 00:00:00 2001
From: Zachary Wasserman <zachwass2000@gmail.com>
Date: Tue, 16 Oct 2018 11:13:55 -0700
Subject: [PATCH] Fix loading of network interfaces with high host counts
 (#1942)

This should fix the loading of the all hosts page in cases where there are many
hosts and it overwhelms the number of parameters allowed in a prepared
statement. May also make that page load slightly quicker as it removes the
constraint from the query, but should return the same number of results.

Fixes #1939
---
 server/datastore/datastore_hosts_test.go | 11 ++++++
 server/datastore/mysql/hosts.go          | 43 +++++++++++++++++++++---
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/server/datastore/datastore_hosts_test.go b/server/datastore/datastore_hosts_test.go
index 67762a44ef..bc078080cd 100644
--- a/server/datastore/datastore_hosts_test.go
+++ b/server/datastore/datastore_hosts_test.go
@@ -208,6 +208,17 @@ func testListHost(t *testing.T, ds kolide.Datastore) {
 	assert.Equal(t, "en1", hosts2[1].NetworkInterfaces[1].Interface)
 	assert.Equal(t, "en2", hosts2[3].NetworkInterfaces[0].Interface)
 
+	// Test with logic for only a few hosts
+	hosts2, err = ds.ListHosts(kolide.ListOptions{PerPage: 4, Page: 0})
+	require.Nil(t, err)
+	assert.Equal(t, 4, len(hosts2))
+
+	require.Equal(t, 2, len(hosts2[1].NetworkInterfaces))
+	require.Equal(t, 0, len(hosts2[2].NetworkInterfaces))
+	require.Equal(t, 1, len(hosts2[3].NetworkInterfaces))
+	assert.Equal(t, "en1", hosts2[1].NetworkInterfaces[1].Interface)
+	assert.Equal(t, "en2", hosts2[3].NetworkInterfaces[0].Interface)
+
 	err = ds.DeleteHost(hosts[0].ID)
 	require.Nil(t, err)
 	hosts2, err = ds.ListHosts(kolide.ListOptions{})
diff --git a/server/datastore/mysql/hosts.go b/server/datastore/mysql/hosts.go
index 07fe8fd0cd..4e88b23d5c 100644
--- a/server/datastore/mysql/hosts.go
+++ b/server/datastore/mysql/hosts.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/jmoiron/sqlx"
 	"github.com/kolide/fleet/server/kolide"
-	"github.com/patrickmn/sortutil"
 	"github.com/pkg/errors"
 )
 
@@ -303,8 +302,16 @@ func (d *Datastore) ListHosts(opt kolide.ListOptions) ([]*kolide.Host, error) {
 		return nil, errors.Wrap(err, "list hosts")
 	}
 
-	if err := d.getNetInterfacesForHosts(hosts); err != nil {
-		return nil, err
+	if opt.PerPage == 0 || (opt.Page == 0 && uint(len(hosts)) < opt.PerPage) {
+		// If all hosts, we can use the optimized network interface retrieval function
+		if err := d.getNetInterfacesForAllHosts(hosts); err != nil {
+			return nil, err
+		}
+
+	} else {
+		if err := d.getNetInterfacesForHosts(hosts); err != nil {
+			return nil, err
+		}
 	}
 
 	return hosts, nil
@@ -383,7 +390,35 @@ func (d *Datastore) getNetInterfacesForHosts(hosts []*kolide.Host) error {
 		return errors.Wrap(err, "select nics for hosts, rebound query")
 	}
 
-	sortutil.AscByField(hosts, "ID")
+	for _, host := range hosts {
+		for i := 0; i < len(nics); i++ {
+			if host.ID == nics[i].HostID {
+				host.NetworkInterfaces = append(host.NetworkInterfaces, nics[i])
+			}
+		}
+	}
+
+	return nil
+}
+
+// When we know we're loading the network interfaces for all hosts, we can skip
+// the IN clause and load them all. This allows us to load net interfaces
+// without error for larger sets of hosts.
+func (d *Datastore) getNetInterfacesForAllHosts(hosts []*kolide.Host) error {
+	if len(hosts) == 0 {
+		return nil
+	}
+
+	sqlStatement := `
+		SELECT *
+		FROM network_interfaces
+		ORDER BY host_id ASC
+	`
+	nics := []*kolide.NetworkInterface{}
+	err := d.db.Select(&nics, sqlStatement)
+	if err != nil {
+		return errors.Wrap(err, "select nics for all hosts")
+	}
 
 	for _, host := range hosts {
 		for i := 0; i < len(nics); i++ {