make collector update endpoints on registration

JDragovichAlertLogic · JDragovichAlertLogic · commit 83286acf85d4 · 2019-05-16T11:09:13.000+01:00
diff --git a/.travis.yml b/.travis.yml
@@ -1,7 +1,7 @@
 sudo: false
 language: node_js
 node_js:
-  - "6"
+  - "8"
 install:
   - make deps
 script:
diff --git a/al_aws.js b/al_aws.js
@@ -150,13 +150,27 @@ var arnToAccId = function (arn) {
 
 var setEnv = function(vars, callback) {
     const lambda = new AWS.Lambda();
-    var params = {
-        FunctionName : process.env.AWS_LAMBDA_FUNCTION_NAME,
-        Environment : {
-            Variables : vars
-        }
+
+    const getConfigParams = {
+      FunctionName: process.env.AWS_LAMBDA_FUNCTION_NAME
     };
-    return lambda.updateFunctionConfiguration(params, callback);
+
+    lambda.getFunctionConfiguration(getConfigParams, (err, config) => {
+        if(err){
+            console.error('Error getting function config, environment variables were not updated', err);
+            return callback(err);
+        }
+        var params = {
+            FunctionName: process.env.AWS_LAMBDA_FUNCTION_NAME,
+            Environment : {
+                Variables : {
+                    ...config.Environment.Variables,
+                    ...vars
+                }
+            }
+        };
+        return lambda.updateFunctionConfiguration(params, callback);
+    });
 };
 
 module.exports = {
diff --git a/al_aws_collector.js b/al_aws_collector.js
@@ -145,13 +145,55 @@ class AlAwsCollector {
         const context = this._invokeContext;
         const regValues = Object.assign(this._getAttrs(), custom);
 
-        this._azcollectc.register(regValues)
-            .then(resp => {
+        async.waterfall([
+            (asyncCallback) => {
+                const {
+                    azcollect_api,
+                    ingest_api
+                } = process.env;
+
+                if(!azcollect_api || !ingest_api){
+                    // handling errors like this because the other unit tests seem to indicat that the collectorshould
+                    // register even if there is an error in getting the endpoints.
+                    this.updateEndpoints((err, newConfig) => {
+                        if(err){
+                            console.warn('Error updating endpoints', err);
+                        } else {
+                            // reassign env vars because the config change occurs in the same run in registration.
+                            const {
+                                Environment: {
+                                    Variables
+                                }
+                            } = newConfig;
+
+                            Object.assign(process.env, Variables);
+                            this._azcollectc = new m_alCollector.AzcollectC(process.env.azollect_api, this._aimsc, this._collectorType);
+                            this._ingestc = new m_alCollector.IngestC(process.env.ingest_api, this._aimsc, 'lambda_function');
+                        }
+
+                        asyncCallback(null);
+                    });
+                } else{
+                    asyncCallback(null);
+                }
+            },
+            (asyncCallback) => {
+                this._azcollectc.register(regValues)
+                    .then(resp => {
+                        asyncCallback(null);
+                    })
+                    .catch(exception => {
+                        asyncCallback("registration error: " + exception);
+                    });
+            }
+        ],
+        (err)=> {
+            if(err){
+                return response.send(event, context, response.FAILED, {Error: err});
+            } else {
                 return response.send(event, context, response.SUCCESS);
-            })
-            .catch(exception => {
-                return response.send(event, context, response.FAILED, {Error: exception});
-            });
+            }
+        });
     }
 
     checkin(callback) {
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@alertlogic/al-aws-collector-js",
-  "version": "1.1.3",
+  "version": "1.2.0",
   "license": "MIT",
   "description": "Alert Logic AWS Collector Common Library",
   "repository": {
diff --git a/test/al_aws_collector_test.js b/test/al_aws_collector_test.js
@@ -5,6 +5,7 @@ const m_response = require('cfn-response');
 const deepEqual = require('deep-equal');
 
 const AlAwsCollector = require('../al_aws_collector');
+const m_al_aws = require('../al_aws');
 var m_alCollector = require('@alertlogic/al-collector-js');
 const m_healthChecks = require('../health_checks');
 var AWS = require('aws-sdk-mock');
@@ -17,6 +18,7 @@ const context = {
 
 var alserviceStub = {};
 var responseStub = {};
+var setEnvStub = {};
 
 function setAlServiceStub() {
     alserviceStub.get = sinon.stub(m_alCollector.AlServiceC.prototype, 'get').callsFake(
@@ -116,6 +118,23 @@ function mockLambdaEndpointsUpdateConfiguration() {
     });
 }
 
+function mockSetEnvStub() {
+    setEnvStub = sinon.stub(m_al_aws, 'setEnv').callsFake((vars, callback)=>{
+        const {
+            ingest_api,
+            azcollect_api
+        } = vars;
+        process.env.ingest_api = ingest_api ? ingest_api : process.env.ingest_api;
+        process.env.azollect_api = azcollect_api ? azcollect_api : process.env.azollect_api;
+        const returnBody = {
+            Environment: {
+                Varaibles: vars
+            }
+        };
+        return callback(null, returnBody);
+    });
+}
+
 function mockLambdaMetricStatistics() {
     AWS.mock('CloudWatch', 'getMetricStatistics', function (params, callback) {
         var ret = colMock.CLOUDWATCH_GET_METRIC_STATS_OK;
@@ -144,18 +163,41 @@ describe('al_aws_collector tests', function() {
             });
 
         setAlServiceStub();
+        mockSetEnvStub();
     });
 
     afterEach(function(){
         restoreAlServiceStub();
+        setEnvStub.restore();
         responseStub.restore();
     });
 
-    it('register success', function(done) {
+    it('register success with env vars set', function(done) {
+        var mockContext = {
+            invokedFunctionArn : colMock.FUNCTION_ARN,
+            succeed : () => {
+                sinon.assert.calledWith(alserviceStub.post, colMock.REG_URL, colMock.REG_AZCOLLECT_QUERY);
+                sinon.assert.neverCalledWithMatch(alserviceStub.get, colMock.GET_INGEST_URL);
+                sinon.assert.neverCalledWithMatch(alserviceStub.get, colMock.GET_AZCOLLECT_URL);
+                done();
+            }
+        };
+        AlAwsCollector.load().then(function(creds) {
+            var collector = new AlAwsCollector(
+            mockContext, 'cwe', AlAwsCollector.IngestTypes.SECMSGS, '1.0.0', creds, function() {});
+            collector.register(colMock.REGISTRATION_TEST_EVENT, colMock.REG_PARAMS);
+        });
+    });
+
+    it('register success with env vars not set', function(done) {
+        mockLambdaEndpointsUpdateConfiguration();
+        delete process.env.ingest_api;
+        delete process.env.azcollect_api;
         var mockContext = {
             invokedFunctionArn : colMock.FUNCTION_ARN,
             succeed : () => {
                 sinon.assert.calledWith(alserviceStub.post, colMock.REG_URL, colMock.REG_AZCOLLECT_QUERY);
+                sinon.assert.calledTwice(alserviceStub.get);
                 done();
             }
         };
@@ -752,11 +794,13 @@ describe('al_aws_collector error tests', function() {
             return callback(null, colMock.CF_DESCRIBE_STACKS_RESPONSE);
         });
         mockLambdaMetricStatistics();
+        mockSetEnvStub();
     });
 
     afterEach(function(){
         restoreAlServiceStub();
         responseStub.restore();
+        setEnvStub.restore();
         AWS.restore('CloudFormation', 'describeStacks');
         AWS.restore('CloudWatch', 'getMetricStatistics');
     });
@@ -766,7 +810,7 @@ describe('al_aws_collector error tests', function() {
             invokedFunctionArn : colMock.FUNCTION_ARN,
             done : () => {
                 sinon.assert.calledWith(alserviceStub.post, colMock.REG_URL, colMock.REG_AZCOLLECT_QUERY);
-                sinon.assert.calledWith(responseStub, sinon.match.any, sinon.match.any, m_response.FAILED, {Error: 'post error'});
+                sinon.assert.calledWith(responseStub, sinon.match.any, sinon.match.any, m_response.FAILED, {Error: 'registration error: post error'});
                 done();
             }
         };
diff --git a/test/collector_mock.js b/test/collector_mock.js
@@ -8,7 +8,7 @@ process.env.AWS_REGION = 'us-east-1';
 process.env.AWS_LAMBDA_FUNCTION_NAME = FUNCTION_NAME;
 process.env.al_api = 'api.global-services.global.alertlogic.com';
 process.env.ingest_api = 'ingest.global-services.global.alertlogic.com';
-process.env.azollect_api = 'azcollect.global-services.global.alertlogic.com';
+process.env.azcollect_api = 'azcollect.global-services.global.alertlogic.com';
 process.env.aims_access_key_id = 'aims-key-id';
 process.env.aims_secret_key = 'aims-secret-key-encrypted';
 process.env.aws_lambda_s3_bucket = S3_BUCKET;
@@ -60,6 +60,8 @@ const DEREGISTRATION_TEST_EVENT = {
 };
 
 const REG_URL = '/aws/cwe/123456789012/us-east-1/' + encodeURIComponent(FUNCTION_NAME);
+const GET_INGEST_URL = '/residency/default/serevices/ingest/endpoint';
+const GET_AZCOLLECT_URL = '/residency/default/services/azcollect/endpoint';
 const REG_PARAMS = {
     stackName : STACK_NAME,
     custom_fields: {
@@ -253,7 +255,7 @@ const LAMBDA_FUNCTION_CONFIGURATION = {
             al_api: process.env.al_api,
             aws_lambda_s3_bucket: S3_BUCKET,
             aws_lambda_zipfile_name: S3_ZIPFILE,
-            azollect_api: process.env.azollect_api,
+            azcollect_api: process.env.azcollect_api,
             ingest_api: process.env.ingest_api 
         } 
     },
@@ -282,7 +284,7 @@ const LAMBDA_FUNCTION_CONFIGURATION_CHANGED = {
             al_api: 'new al_api value',
             aws_lambda_s3_bucket: S3_BUCKET,
             aws_lambda_zipfile_name: S3_ZIPFILE,
-            azollect_api: process.env.azollect_api,
+            azcollect_api: process.env.azcollect_api,
             ingest_api: process.env.ingest_api,
             x: 'XXXX'
         } 
@@ -303,6 +305,8 @@ module.exports = {
 
     REGISTRATION_TEST_EVENT : REGISTRATION_TEST_EVENT,
     REG_URL : REG_URL,
+    GET_AZCOLLECT_URL: GET_AZCOLLECT_URL,
+    GET_INGEST_URL: GET_INGEST_URL,
     REG_PARAMS : REG_PARAMS,
     REG_AZCOLLECT_QUERY : REG_AZCOLLECT_QUERY,
 

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "@alertlogic/al-aws-collector-js",`
`3`		`- "version": "1.1.3",`
	`3`	`+ "version": "1.2.0",`
`4`	`4`	`"license": "MIT",`
`5`	`5`	`"description": "Alert Logic AWS Collector Common Library",`
`6`	`6`	`"repository": {`