Skip to content

Commit

Permalink
feat: use pino logging lib for S3 scan object (#183)
Browse files Browse the repository at this point in the history 
Update to use the `pino` logging lib for the S3 scan object lambda.
This allows for lambda invocation specfic parameters like `RequestId`
to be added automatically to all log messages.
  • Loading branch information
@patheard
patheard authored Jul 12, 2022
1 parent 2df099a commit 568935b
Show file tree
Hide file tree
Showing 4 changed files with 188 additions and 40 deletions.
47 changes: 31 additions & 16 deletions module/s3-scan-object/app.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,17 @@
*/

const axios = require("axios");
const pino = require("pino");
const util = require("util");
const { lambdaRequestTracker, pinoLambdaDestination } = require("pino-lambda");
const { S3Client, PutObjectTaggingCommand } = require("@aws-sdk/client-s3");
const { SecretsManagerClient, GetSecretValueCommand } = require("@aws-sdk/client-secrets-manager");
const { STSClient, AssumeRoleCommand } = require("@aws-sdk/client-sts");

const AWS_ROLE_TO_ASSUME = process.env.AWS_ROLE_TO_ASSUME ? process.env.AWS_ROLE_TO_ASSUME : "ScanFilesGetObjects";
const REGION = process.env.REGION;
const ENDPOINT_URL = process.env.AWS_SAM_LOCAL ? "http://host.docker.internal:3001" : undefined;
const LOGGING_LEVEL = process.env.LOGGING_LEVEL ? process.env.LOGGING_LEVEL : "info";
const SCAN_FILES_URL = process.env.SCAN_FILES_URL;
const SCAN_FILES_API_KEY_SECRET_ARN = process.env.SCAN_FILES_API_KEY_SECRET_ARN;
const SCAN_IN_PROGRESS = "in_progress";
Expand All @@ -27,6 +30,16 @@ const EVENT_SNS = "aws:sns";
const stsClient = new STSClient({ region: REGION, endpoint: ENDPOINT_URL });
const secretsManagerClient = new SecretsManagerClient({ region: REGION, endpoint: ENDPOINT_URL });

// Setup logging and add a custom requestId attribute to all log messages
const logger = pino({ level: LOGGING_LEVEL }, pinoLambdaDestination());
const withRequest = lambdaRequestTracker({
requestMixin: (event) => {
return {
requestId: event.RequestId ? event.RequestId : undefined,
};
},
});

/**
* Performs function initialization outside of the Lambda handler so that
* it only occurs once per cold start of the function rather than on
Expand All @@ -42,7 +55,7 @@ const initConfig = async () => {
const response = await secretsManagerClient.send(command);
return { apiKey: response.SecretString };
} catch (error) {
console.error(`Unable to get '${SCAN_FILES_API_KEY_SECRET_ARN}' secret: ${error}`);
logger.error(`Unable to get '${SCAN_FILES_API_KEY_SECRET_ARN}' secret: ${error}`);
throw error;
}
})();
Expand All @@ -58,7 +71,9 @@ const configPromise = initConfig();
* is received with an update scan status.
* @param {Object} event Lambda invocation event
*/
exports.handler = async (event) => {
exports.handler = async (event, context) => {
withRequest(event, context);

const config = await configPromise;
const s3Clients = {};
let errorCount = 0;
Expand All @@ -80,6 +95,9 @@ exports.handler = async (event) => {
sns: "request-id",
});

// Make sure SNS events have a top-level RequestId attribute for logging
event.RequestId = requestId || event.RequestId;

// Do not scan S3 folder objects
if (isS3Folder(s3Object)) {
continue;
Expand All @@ -104,7 +122,7 @@ exports.handler = async (event) => {
scanStatus = record.Sns.MessageAttributes["av-status"].Value;
}
} else {
console.error(`[${requestId}] Unsupported event record: ${util.inspect(record)}`);
logger.error(`Unsupported event record: ${util.inspect(record)}`);
}

// Tag the S3 object if we've got a scan status
Expand All @@ -124,8 +142,8 @@ exports.handler = async (event) => {
}

roleArn = `arn:aws:iam::${awsAccountId}:role/${AWS_ROLE_TO_ASSUME}`;
s3Clients[awsAccountId] = await getS3Client(s3Clients[awsAccountId], stsClient, roleArn, requestId);
isObjectTagged = await tagS3Object(s3Clients[awsAccountId], s3Object, tags, requestId);
s3Clients[awsAccountId] = await getS3Client(s3Clients[awsAccountId], stsClient, roleArn);
isObjectTagged = await tagS3Object(s3Clients[awsAccountId], s3Object, tags);
}

// Track if there were any errors processing this record
Expand Down Expand Up @@ -176,10 +194,9 @@ const getRecordEventSource = (record) => {
* be used by other SDK clients.
* @param {STSClient} stsClient AWS SDK STS client used to assume the role
* @param {string} roleArn ARN of the role to assume
* @param {string} requestId Request ID of the scan
* @returns Credentials object for the role
*/
const getRoleCredentials = async (stsClient, roleArn, requestId) => {
const getRoleCredentials = async (stsClient, roleArn) => {
let credentials = null;

try {
Expand All @@ -192,7 +209,7 @@ const getRoleCredentials = async (stsClient, roleArn, requestId) => {
sessionToken: response.Credentials.SessionToken,
};
} catch (error) {
console.error(`[${requestId}] Failed to assume role ${roleArn}: ${error}`);
logger.error(`Failed to assume role ${roleArn}: ${error}`);
}

return credentials;
Expand All @@ -204,12 +221,11 @@ const getRoleCredentials = async (stsClient, roleArn, requestId) => {
* @param {S3Client} s3Client Initialized S3 client or null
* @param {STSClient} stsClient STS client used to assume the role
* @param {string} roleArn ARN of the role to assume to get tempoary credentials
* @param {string} requestId Request ID of the scan
* @returns S3 client
*/
const getS3Client = async (s3Client, stsClient, roleArn, requestId) => {
const getS3Client = async (s3Client, stsClient, roleArn) => {
if (!s3Client) {
const credentials = await getRoleCredentials(stsClient, roleArn, requestId);
const credentials = await getRoleCredentials(stsClient, roleArn);
return new S3Client({
region: REGION,
endpoint: ENDPOINT_URL,
Expand Down Expand Up @@ -301,10 +317,10 @@ const startS3ObjectScan = async (apiEndpoint, apiKey, s3Object, awsAccountId, sn
},
}
);
console.info(`[${requestId}] Scan response ${response.status}: ${util.inspect(response.data)}`);
logger.info(`Scan response ${response.status}: ${util.inspect(response.data)}`);
return response;
} catch (error) {
console.error(`[${requestId}] Could not start scan for ${util.inspect(s3Object)}: ${util.inspect(error.response)}`);
logger.error(`Could not start scan for ${util.inspect(s3Object)}: ${util.inspect(error.response)}`);
return error.response;
}
};
Expand All @@ -314,9 +330,8 @@ const startS3ObjectScan = async (apiEndpoint, apiKey, s3Object, awsAccountId, sn
* @param {S3Client} s3Client AWS SDK S3 client used to tag the object
* @param {{Bucket: string, Key: string}} s3Object S3 object to tag
* @param {Array<{Key: string, Value: string}>} tags Array of Key/Value pairs to tag the S3 object with
* @param {string} requestId Request ID of the scan
*/
const tagS3Object = async (s3Client, s3Object, tags, requestId) => {
const tagS3Object = async (s3Client, s3Object, tags) => {
const tagging = {
Tagging: {
TagSet: tags,
Expand All @@ -329,7 +344,7 @@ const tagS3Object = async (s3Client, s3Object, tags, requestId) => {
const response = await s3Client.send(command);
isSuccess = response.VersionId !== undefined;
} catch (error) {
console.error(`[${requestId}] Failed to tag S3 object: ${error}`);
logger.error(`Failed to tag S3 object: ${error}`);
}

return isSuccess;
Expand Down
36 changes: 15 additions & 21 deletions module/s3-scan-object/app.test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,6 @@ const {
} = helpers;

jest.mock("axios");
global.console = {
...console,
error: jest.fn(),
info: jest.fn(),
};

const TEST_TIME = new Date(1978, 3, 30).getTime();
beforeAll(() => {
Expand Down Expand Up @@ -70,6 +65,7 @@ describe("handler", () => {
EventSource: "aws:sns",
Sns: {
MessageAttributes: {
"request-id": { Value: "zxcv9584" },
"av-filepath": { Value: "s3://bam/baz" },
"av-status": { Value: "SPIFY" },
"av-checksum": { Value: "42" },
Expand Down Expand Up @@ -104,7 +100,7 @@ describe("handler", () => {
mockS3Client.on(PutObjectTaggingCommand).resolves({ VersionId: "yeet" });
mockSTSClient.on(AssumeRoleCommand).resolves({ Credentials: {} });

const response = await handler(event);
const response = await handler(event, {});
expect(response).toEqual(expectedResponse);
expect(mockS3Client).toHaveReceivedNthCommandWith(1, PutObjectTaggingCommand, {
Bucket: "foo",
Expand All @@ -127,6 +123,7 @@ describe("handler", () => {
{ Key: "av-status", Value: "SPIFY" },
{ Key: "av-timestamp", Value: TEST_TIME },
{ Key: "av-checksum", Value: "42" },
{ Key: "request-id", Value: "zxcv9584" },
],
},
});
Expand All @@ -138,7 +135,7 @@ describe("handler", () => {
{ Key: "av-scanner", Value: "clamav" },
{ Key: "av-status", Value: "in_progress" },
{ Key: "av-timestamp", Value: TEST_TIME },
{ Key: "request-id", Value: "1234asdf" },
{ Key: "request-id", Value: "zxcv9584" },
],
},
});
Expand Down Expand Up @@ -187,7 +184,7 @@ describe("handler", () => {
mockS3Client.on(PutObjectTaggingCommand).resolves({ VersionId: "yeet" });
mockSTSClient.on(AssumeRoleCommand).resolves({ Credentials: {} });

const response = await handler(event);
const response = await handler(event, {});
expect(response).toEqual(expectedResponse);
expect(mockS3Client).toHaveReceivedNthCommandWith(1, PutObjectTaggingCommand, {
Bucket: "foo",
Expand Down Expand Up @@ -230,7 +227,7 @@ describe("handler", () => {
mockS3Client.on(PutObjectTaggingCommand).resolves({ VersionId: "yeet" });
mockSTSClient.on(AssumeRoleCommand).resolves({ Credentials: {} });

const response = await handler(event);
const response = await handler(event, {});
expect(response).toEqual(expectedResponse);
expect(mockS3Client).toHaveReceivedNthCommandWith(1, PutObjectTaggingCommand, {
Bucket: "foo",
Expand Down Expand Up @@ -266,7 +263,7 @@ describe("handler", () => {
axios.post.mockResolvedValue({ status: 200 });
mockS3Client.on(PutObjectTaggingCommand).resolves({ VersionId: "yeet" });

const response = await handler(event);
const response = await handler(event, {});
expect(response).toEqual(expectedResponse);
});
});
Expand Down Expand Up @@ -331,7 +328,7 @@ describe("getRoleCredentials", () => {
},
};
mockSTSClient.on(AssumeRoleCommand).resolves(response);
const credentials = await getRoleCredentials(mockSTSClient, "foo", "123");
const credentials = await getRoleCredentials(mockSTSClient, "foo");

expect(credentials).toEqual({
accessKeyId: "why",
Expand All @@ -346,15 +343,15 @@ describe("getRoleCredentials", () => {

test("fails to assume role", async () => {
mockSTSClient.on(AssumeRoleCommand).rejects(new Error("nope"));
const credentials = await getRoleCredentials(mockSTSClient, "foo", "123");
const credentials = await getRoleCredentials(mockSTSClient, "foo");
expect(credentials).toBe(null);
});
});

describe("getS3Client", () => {
test("successfully gets new client", async () => {
mockSTSClient.on(AssumeRoleCommand).resolves({ Credentials: { foo: "bar" } });
const s3Client = await getS3Client(null, mockSTSClient, "bar", "bam");
const s3Client = await getS3Client(null, mockSTSClient, "bar");
expect(s3Client).toBeInstanceOf(S3Client);
expect(mockSTSClient).toHaveReceivedNthCommandWith(1, AssumeRoleCommand, {
RoleArn: "bar",
Expand All @@ -363,7 +360,7 @@ describe("getS3Client", () => {
});

test("successfully returns cached client", async () => {
const s3Client = await getS3Client("mellow", mockSTSClient, "bar", "baz");
const s3Client = await getS3Client("mellow", mockSTSClient, "bar");
expect(s3Client).toBe("mellow");
expect(mockSTSClient.calls().length).toBe(0);
});
Expand Down Expand Up @@ -543,19 +540,16 @@ describe("tagS3Object", () => {
TagSet: [{ Key: "some-tag", Value: "some-value" }],
},
};
const response = await tagS3Object(
mockS3Client,
{ Bucket: "foo", Key: "bar" },
[{ Key: "some-tag", Value: "some-value" }],
"foo"
);
const response = await tagS3Object(mockS3Client, { Bucket: "foo", Key: "bar" }, [
{ Key: "some-tag", Value: "some-value" },
]);
expect(response).toBe(true);
expect(mockS3Client).toHaveReceivedCommandWith(PutObjectTaggingCommand, input);
});

test("fails to tag", async () => {
mockS3Client.on(PutObjectTaggingCommand).resolvesOnce({});
const response = await tagS3Object(mockS3Client, {}, [], null);
const response = await tagS3Object(mockS3Client, {}, []);
expect(response).toBe(false);
});
});
4 changes: 3 additions & 1 deletion module/s3-scan-object/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,9 @@
"@aws-sdk/client-s3": "^3.113.0",
"@aws-sdk/client-secrets-manager": "^3.118.1",
"@aws-sdk/client-sts": "^3.121.0",
"axios": "^0.27.2"
"axios": "^0.27.2",
"pino": "7.11.0",
"pino-lambda": "^4.0.0"
},
"devDependencies": {
"aws-sdk-client-mock": "1.0.0",
Expand Down
Loading

0 comments on commit 568935b

Please sign in to comment.