diff --git a/server/service/endpoint_campaigns.go b/server/service/endpoint_campaigns.go index 2a6bfc7fcd..2c376c7956 100644 --- a/server/service/endpoint_campaigns.go +++ b/server/service/endpoint_campaigns.go @@ -5,12 +5,12 @@ import ( "encoding/json" "net/http" - "github.com/go-kit/kit/endpoint" - kitlog "github.com/go-kit/kit/log" - "github.com/igm/sockjs-go/v3/sockjs" "github.com/fleetdm/fleet/server/contexts/viewer" "github.com/fleetdm/fleet/server/kolide" "github.com/fleetdm/fleet/server/websocket" + "github.com/go-kit/kit/endpoint" + kitlog "github.com/go-kit/kit/log" + "github.com/igm/sockjs-go/v3/sockjs" ) //////////////////////////////////////////////////////////////////////////////// @@ -79,9 +79,14 @@ func makeStreamDistributedQueryCampaignResultsHandler(svc kolide.Service, jwtKey opt.Websocket = true opt.RawWebsocket = true return sockjs.NewHandler("/api/v1/kolide/results", opt, func(session sockjs.Session) { - defer session.Close(0, "none") - conn := &websocket.Conn{Session: session} + defer func() { + if p := recover(); p != nil { + logger.Log("err", p, "msg", "panic in result handler") + conn.WriteJSONError("panic in result handler") + } + session.Close(0, "none") + }() // Receive the auth bearer token token, err := conn.ReadAuthToken() diff --git a/server/service/service_campaigns.go b/server/service/service_campaigns.go index 45d2b82938..89a5ad4244 100644 --- a/server/service/service_campaigns.go +++ b/server/service/service_campaigns.go @@ -163,11 +163,18 @@ func (svc service) StreamCampaignResults(ctx context.Context, conn *websocket.Co lastTotals := targetTotals{} // to improve performance of the frontend rendering the results table, we - // add the "host_hostname" field to every row. - mapHostnameRows := func(hostname string, rows []map[string]string) { - for _, row := range rows { - row["host_hostname"] = hostname + // add the "host_hostname" field to every row and clean null rows. + mapHostnameRows := func(res *kolide.DistributedQueryResult) { + filteredRows := []map[string]string{} + for _, row := range res.Rows { + if row == nil { + continue + } + row["host_hostname"] = res.Host.HostName + filteredRows = append(filteredRows, row) } + + res.Rows = filteredRows } hostIDs, labelIDs, err := svc.ds.DistributedQueryCampaignTargetIDs(campaign.ID) @@ -230,7 +237,7 @@ func (svc service) StreamCampaignResults(ctx context.Context, conn *websocket.Co // Receive a result and push it over the websocket switch res := res.(type) { case kolide.DistributedQueryResult: - mapHostnameRows(res.Host.HostName, res.Rows) + mapHostnameRows(&res) err = conn.WriteJSONMessage("result", res) if errors.Cause(err) == sockjs.ErrSessionNotOpen { // return and stop sending the query if the session was closed