[Asset Management] Osquery telemetry updates (#100754) (#102557)

* first pass of basic osquery usage stats collection

* updates, linting

* updated exported metrics

* clean up comments, add description fields to metric fields

* reworked types

* actually use the updated types

* added tests around the route usage recoder functions

* review comments

* update aggregate types

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

Co-authored-by: Bryan Clement <bclement01@gmail.com>

x-pack/plugins/osquery/common/types.ts

@@ -9,7 +9,8 @@ import { PackagePolicy, PackagePolicyInput, PackagePolicyInputStream } from '../
 
 export const savedQuerySavedObjectType = 'osquery-saved-query';
 export const packSavedObjectType = 'osquery-pack';
-export type SavedObjectType = 'osquery-saved-query' | 'osquery-pack';
+export const usageMetricSavedObjectType = 'osquery-usage-metric';
+export type SavedObjectType = 'osquery-saved-query' | 'osquery-pack' | 'osquery-usage-metric';
 
 /**
  * This makes any optional property the same as Required<T> would but also has the

x-pack/plugins/osquery/kibana.json

@@ -10,6 +10,7 @@
   "kibanaVersion": "kibana",
   "optionalPlugins": [
     "home",
+    "usageCollection",
     "lens"
   ],
   "requiredBundles": [

x-pack/plugins/osquery/server/plugin.ts

@@ -18,6 +18,7 @@ import { OsqueryPluginSetup, OsqueryPluginStart, SetupPlugins, StartPlugins } fr
 import { defineRoutes } from './routes';
 import { osquerySearchStrategyProvider } from './search_strategy/osquery';
 import { initSavedObjects } from './saved_objects';
+import { initUsageCollectors } from './usage';
 import { OsqueryAppContext, OsqueryAppContextService } from './lib/osquery_app_context_services';
 import { ConfigType } from './config';
 
@@ -48,6 +49,11 @@ export class OsqueryPlugin implements Plugin<OsqueryPluginSetup, OsqueryPluginSt
     };
 
     initSavedObjects(core.savedObjects, osqueryContext);
+    initUsageCollectors({
+      core,
+      osqueryContext,
+      usageCollection: plugins.usageCollection,
+    });
     defineRoutes(router, osqueryContext);
 
     core.getStartServices().then(([, depsStart]) => {

x-pack/plugins/osquery/server/routes/action/create_action_route.ts

@@ -18,6 +18,8 @@ import {
   CreateActionRequestBodySchema,
 } from '../../../common/schemas/routes/action/create_action_request_body_schema';
 
+import { incrementCount } from '../usage';
+
 export const createActionRoute = (router: IRouter, osqueryContext: OsqueryAppContext) => {
   router.post(
     {
@@ -39,34 +41,43 @@ export const createActionRoute = (router: IRouter, osqueryContext: OsqueryAppCon
         osqueryContext,
         agentSelection
       );
-
+      incrementCount(soClient, 'live_query');
       if (!selectedAgents.length) {
+        incrementCount(soClient, 'live_query', 'errors');
         return response.badRequest({ body: new Error('No agents found for selection') });
       }
 
-      const action = {
-        action_id: uuid.v4(),
-        '@timestamp': moment().toISOString(),
-        expiration: moment().add(1, 'days').toISOString(),
-        type: 'INPUT_ACTION',
-        input_type: 'osquery',
-        agents: selectedAgents,
-        data: {
-          id: uuid.v4(),
-          query: request.body.query,
-        },
-      };
-      const actionResponse = await esClient.index<{}, {}>({
-        index: '.fleet-actions',
-        body: action,
-      });
+      try {
+        const action = {
+          action_id: uuid.v4(),
+          '@timestamp': moment().toISOString(),
+          expiration: moment().add(1, 'days').toISOString(),
+          type: 'INPUT_ACTION',
+          input_type: 'osquery',
+          agents: selectedAgents,
+          data: {
+            id: uuid.v4(),
+            query: request.body.query,
+          },
+        };
+        const actionResponse = await esClient.index<{}, {}>({
+          index: '.fleet-actions',
+          body: action,
+        });
 
-      return response.ok({
-        body: {
-          response: actionResponse,
-          actions: [action],
-        },
-      });
+        return response.ok({
+          body: {
+            response: actionResponse,
+            actions: [action],
+          },
+        });
+      } catch (error) {
+        incrementCount(soClient, 'live_query', 'errors');
+        return response.customError({
+          statusCode: 500,
+          body: new Error(`Error occurred whlie processing ${error}`),
+        });
+      }
     }
   );
 };

x-pack/plugins/osquery/server/routes/usage/index.ts

@@ -0,0 +1,8 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export * from './recorder';

x-pack/plugins/osquery/server/routes/usage/recorder.test.ts

@@ -0,0 +1,135 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { savedObjectsClientMock } from '../../../../../../src/core/server/mocks';
+
+import { usageMetricSavedObjectType } from '../../../common/types';
+
+import {
+  CounterValue,
+  createMetricObjects,
+  getRouteMetric,
+  incrementCount,
+  RouteString,
+  routeStrings,
+} from './recorder';
+
+const savedObjectsClient = savedObjectsClientMock.create();
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function checkGetCalls(calls: any[]) {
+  expect(calls.length).toEqual(routeStrings.length);
+  for (let i = 0; i < routeStrings.length; ++i) {
+    expect(calls[i]).toEqual([usageMetricSavedObjectType, routeStrings[i]]);
+  }
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function checkCreateCalls(calls: any[], expectedCallRoutes: string[] = routeStrings) {
+  expect(calls.length).toEqual(expectedCallRoutes.length);
+  for (let i = 0; i < expectedCallRoutes.length; ++i) {
+    expect(calls[i][0]).toEqual(usageMetricSavedObjectType);
+    expect(calls[i][2].id).toEqual(expectedCallRoutes[i]);
+  }
+}
+
+describe('Usage metric recorder', () => {
+  describe('Metric initalizer', () => {
+    const get = savedObjectsClient.get as jest.Mock;
+    const create = savedObjectsClient.create as jest.Mock;
+    afterEach(() => {
+      get.mockClear();
+      create.mockClear();
+    });
+    it('should seed route metrics objects', async () => {
+      get.mockRejectedValueOnce('stub value');
+      create.mockReturnValueOnce('stub value');
+      const result = await createMetricObjects(savedObjectsClient);
+      checkGetCalls(get.mock.calls);
+      checkCreateCalls(create.mock.calls);
+      expect(result).toBe(true);
+    });
+
+    it('should handle previously seeded objects properly', async () => {
+      get.mockReturnValueOnce('stub value');
+      create.mockRejectedValueOnce('stub value');
+      const result = await createMetricObjects(savedObjectsClient);
+      checkGetCalls(get.mock.calls);
+      checkCreateCalls(create.mock.calls, []);
+      expect(result).toBe(true);
+    });
+
+    it('should report failure to create the metrics object', async () => {
+      get.mockRejectedValueOnce('stub value');
+      create.mockRejectedValueOnce('stub value');
+      const result = await createMetricObjects(savedObjectsClient);
+      checkGetCalls(get.mock.calls);
+      checkCreateCalls(create.mock.calls);
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('Incrementation', () => {
+    let counterMap: { [key: string]: CounterValue };
+    const get = savedObjectsClient.get as jest.Mock;
+    const update = savedObjectsClient.update as jest.Mock;
+    update.mockImplementation(
+      async (objectType: string, route: RouteString, newVal: CounterValue) => {
+        counterMap[`${objectType}-${route}`] = newVal;
+      }
+    );
+    get.mockImplementation(async (objectType: string, route: RouteString) => ({
+      attributes: counterMap[`${objectType}-${route}`],
+    }));
+    beforeEach(() => {
+      counterMap = routeStrings.reduce((acc, route) => {
+        acc[`${usageMetricSavedObjectType}-${route}`] = {
+          count: 0,
+          errors: 0,
+        };
+        return acc;
+      }, {} as { [key: string]: CounterValue });
+      get.mockClear();
+      update.mockClear();
+    });
+    it('should increment the route counter', async () => {
+      expect(await getRouteMetric(savedObjectsClient, 'live_query')).toEqual({
+        count: 0,
+        errors: 0,
+      });
+      await incrementCount(savedObjectsClient, 'live_query');
+      expect(await getRouteMetric(savedObjectsClient, 'live_query')).toEqual({
+        count: 1,
+        errors: 0,
+      });
+    });
+
+    it('should allow incrementing the error counter', async () => {
+      expect(await getRouteMetric(savedObjectsClient, 'live_query')).toEqual({
+        count: 0,
+        errors: 0,
+      });
+      await incrementCount(savedObjectsClient, 'live_query', 'errors');
+      expect(await getRouteMetric(savedObjectsClient, 'live_query')).toEqual({
+        count: 0,
+        errors: 1,
+      });
+    });
+
+    it('should allow adjustment of the increment', async () => {
+      expect(await getRouteMetric(savedObjectsClient, 'live_query')).toEqual({
+        count: 0,
+        errors: 0,
+      });
+      await incrementCount(savedObjectsClient, 'live_query', 'count', 2);
+      expect(await getRouteMetric(savedObjectsClient, 'live_query')).toEqual({
+        count: 2,
+        errors: 0,
+      });
+    });
+  });
+});

x-pack/plugins/osquery/server/routes/usage/recorder.ts

@@ -0,0 +1,65 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { SavedObjectsClientContract } from 'kibana/server';
+import { usageMetricSavedObjectType } from '../../../common/types';
+import { LiveQuerySessionUsage } from '../../usage/types';
+
+export interface RouteUsageMetric {
+  queries: number;
+  errors: number;
+}
+
+export type RouteString = 'live_query';
+
+export const routeStrings: RouteString[] = ['live_query'];
+
+export async function createMetricObjects(soClient: SavedObjectsClientContract) {
+  const res = await Promise.allSettled(
+    routeStrings.map(async (route) => {
+      try {
+        await soClient.get(usageMetricSavedObjectType, route);
+      } catch (e) {
+        await soClient.create(
+          usageMetricSavedObjectType,
+          {
+            errors: 0,
+            count: 0,
+          },
+          {
+            id: route,
+          }
+        );
+      }
+    })
+  );
+  return !res.some((e) => e.status === 'rejected');
+}
+
+export async function getCount(soClient: SavedObjectsClientContract, route: RouteString) {
+  return await soClient.get<LiveQuerySessionUsage>(usageMetricSavedObjectType, route);
+}
+
+export interface CounterValue {
+  count: number;
+  errors: number;
+}
+
+export async function incrementCount(
+  soClient: SavedObjectsClientContract,
+  route: RouteString,
+  key: keyof CounterValue = 'count',
+  increment = 1
+) {
+  const metric = await soClient.get<CounterValue>(usageMetricSavedObjectType, route);
+  metric.attributes[key] += increment;
+  await soClient.update(usageMetricSavedObjectType, route, metric.attributes);
+}
+
+export async function getRouteMetric(soClient: SavedObjectsClientContract, route: RouteString) {
+  return (await getCount(soClient, route)).attributes;
+}

x-pack/plugins/osquery/server/routes/usage/saved_object_mappings.ts

@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { SavedObjectsType } from '../../../../../../src/core/server';
+
+import { usageMetricSavedObjectType } from '../../../common/types';
+
+export const usageMetricSavedObjectMappings: SavedObjectsType['mappings'] = {
+  properties: {
+    count: {
+      type: 'long',
+    },
+    errors: {
+      type: 'long',
+    },
+  },
+};
+
+export const usageMetricType: SavedObjectsType = {
+  name: usageMetricSavedObjectType,
+  hidden: false,
+  namespaceType: 'single',
+  mappings: usageMetricSavedObjectMappings,
+};

x-pack/plugins/osquery/server/saved_objects.ts

@@ -9,6 +9,7 @@ import { CoreSetup } from '../../../../src/core/server';
 
 import { OsqueryAppContext } from './lib/osquery_app_context_services';
 import { savedQueryType, packType } from './lib/saved_query/saved_object_mappings';
+import { usageMetricType } from './routes/usage/saved_object_mappings';
 
 const types = [savedQueryType, packType];
 
@@ -20,6 +21,8 @@ export const initSavedObjects = (
 ) => {
   const config = osqueryContext.config();
 
+  savedObjects.registerType(usageMetricType);
+
   if (config.savedQueries) {
     savedObjects.registerType(savedQueryType);
   }

x-pack/plugins/osquery/server/types.ts

@@ -11,6 +11,7 @@ import {
   PluginStart as DataPluginStart,
 } from '../../../../src/plugins/data/server';
 import { FleetStartContract } from '../../fleet/server';
+import { UsageCollectionSetup } from '../../../../src/plugins/usage_collection/server';
 import { PluginSetupContract } from '../../features/server';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
@@ -19,6 +20,7 @@ export interface OsqueryPluginSetup {}
 export interface OsqueryPluginStart {}
 
 export interface SetupPlugins {
+  usageCollection?: UsageCollectionSetup;
   actions: ActionsPlugin['setup'];
   data: DataPluginSetup;
   features: PluginSetupContract;