feat: add token generation (fixes #1)

packages/api/serverless.yml

@@ -10,6 +10,7 @@ params:
     DATA_RETENTION_DAYS: ${env:DATA_RETENTION_DAYS, 30}
     CUSTOM_DOMAIN: ${env:CUSTOM_DOMAIN, ""}
     HAS_CUSTOM_DOMAIN: ${strToBool(${env:HAS_CUSTOM_DOMAIN, "false"})}
+    TRACER_TOKEN: ${env:TRACER_TOKEN, ""}
 
 custom:
   customDomain:
@@ -45,6 +46,7 @@ provider:
     NODE_ENV: ${sls:stage}
     TABLE_NAME: ${self:service}-${sls:stage}
     CUSTOM_DOMAIN: ${param:CUSTOM_DOMAIN}
+    TRACER_TOKEN: ${param:TRACER_TOKEN}
     LAMBDA_LAYER_ARN: !Ref TracerLambdaLayer
     AUTO_TRACE_EXCLUDE: 1
   httpApi:

packages/api/src/events/auto-trace.js

@@ -7,6 +7,7 @@ import {
   ApiGatewayV2Client,
   GetApisCommand,
 } from "@aws-sdk/client-apigatewayv2";
+import { acquireLock, releaseLock } from "../lib/locks";
 
 const supportedRuntimes = ["nodejs16.x", "nodejs18.x", "nodejs20.x"];
 const lambdaExecWrapper = "/opt/nodejs/tracer_wrapper";
@@ -57,6 +58,13 @@ export const autoTrace = async () => {
   // Get our API Gateway endpoint for the collector
   const edgeEndpoint = await getApiEndpoint();
 
+  // Make sure we lock so that only one process is updating lambdas
+  const lockAcquired = await acquireLock("auto-trace");
+  if (!lockAcquired) {
+    console.log("Lock not acquired, skipping");
+    return;
+  }
+
   // List all the lambda functions in the AWS account
   const lambdas = await getAccountLambdas();
 
@@ -118,4 +126,6 @@ export const autoTrace = async () => {
 
     // TODO: save function info in DynamoDB
   }
+
+  await releaseLock("auto-trace");
 };

packages/api/src/index.js

@@ -5,11 +5,13 @@ import fs from "fs";
 
 import collector from "./routes/collector";
 import explore from "./routes/explore";
+import autoTraceRoute from "./routes/auto-trace";
 import { autoTrace } from "./events/auto-trace";
 
 const app = new Hono();
 app.route("/api/spans", collector);
 app.route("/api/explore", explore);
+app.route("/api/auto-trace", autoTraceRoute);
 
 let html = "";
 app.use("/assets/*", serveStatic({ root: "./dist" }));

packages/api/src/lib/database.js

@@ -1,5 +1,6 @@
 import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
 import {
+  DeleteCommand,
   DynamoDBDocumentClient,
   PutCommand,
   QueryCommand,
@@ -22,7 +23,7 @@ export const time = () => Math.floor(Date.now() / 1000);
 export const getExpiryTime = () =>
   time() + 86400 * Number(process.env.DATA_RETENTION_DAYS);
 
-export const put = async (item, expires = false) => {
+export const put = async (item, expires = false, options = {}) => {
   item = {
     ...item,
     _created: time(),
@@ -37,6 +38,7 @@ export const put = async (item, expires = false) => {
     new PutCommand({
       TableName: process.env.TABLE_NAME,
       Item: item,
+      ...options
     }),
   );
 };
@@ -63,6 +65,15 @@ export const update = async (
   );
 };
 
+export const deleteItem = async (Key) => {
+  return await dynamo.send(
+    new DeleteCommand({
+      TableName: process.env.TABLE_NAME,
+      Key,
+    }),
+  );
+}
+
 export const queryAll = async (
   /** @type {Omit<import("@aws-sdk/lib-dynamodb").QueryCommandInput, "TableName">} */ params,
 ) => {

packages/api/src/lib/locks.js

@@ -0,0 +1,39 @@
+import { deleteItem, put } from "./database";
+
+export const acquireLock = async (key, ttl = 900) => {
+  const lockKey = `lock#${key}`;
+  const expires = Math.floor(Date.now() / 1000) + ttl;
+  try {
+    await put(
+      {
+        pk: lockKey,
+        sk: lockKey,
+        _expires: expires,
+      },
+      true,
+      {
+        ConditionExpression: "attribute_not_exists(pk)",
+      },
+    );
+    return true;
+  } catch (error) {
+    if (error.name === "ConditionalCheckFailedException") {
+      return false;
+    }
+    throw error;
+  }
+};
+
+export const releaseLock = async (key) => {
+  const lockKey = `lock#${key}`;
+  try {
+    await deleteItem({
+      pk: lockKey,
+      sk: lockKey,
+    });
+    return true;
+  } catch (e) {
+    console.log(e);
+    return false;
+  }
+};

packages/api/src/routes/auto-trace/index.js

@@ -0,0 +1,17 @@
+import { Hono } from "hono";
+import { autoTrace } from "../../events/auto-trace";
+
+const app = new Hono();
+
+app.post("/", async (c) => {
+  const body = await c.req.json();
+  if (body.token !== process.env.TRACER_TOKEN) {
+    return c.json({ error: "Invalid token" }, 401);
+  }
+
+  await autoTrace();
+
+  return c.json({ success: true });
+});
+
+export default app;

packages/api/src/routes/collector/index.js

@@ -7,6 +7,10 @@ const app = new Hono();
 app.post("/", async (c) => {
   const body = await c.req.json();
 
+  if (process.env.TRACER_TOKEN && body.token !== process.env.TRACER_TOKEN) {
+    return c.json({ error: "Invalid token" }, 401);
+  }
+
   for (const span of body) {
     console.log(span);
 

packages/deploy-script/src/index.js

@@ -2,6 +2,7 @@
 
 import chalk from "chalk";
 import boxen from "boxen";
+import crypto from "crypto";
 import inquirer from "inquirer";
 import degit from "degit";
 import child_process from "child_process";
@@ -12,6 +13,9 @@ import {
   GetApisCommand,
 } from "@aws-sdk/client-apigatewayv2";
 
+// TODO: use previous token if it exists
+const tracerToken = crypto.randomBytes(16).toString("hex");
+
 const exec = (command, options = {}) => {
   const child = child_process.exec(command, {
     ...options,
@@ -96,25 +100,40 @@ await cloner.clone("/tmp/trace-stack");
 console.log(chalk.blue("Installing dependencies..."));
 await exec("yarn install", { cwd: "/tmp/trace-stack" });
 
-// Write .env file
+// Write config files
+console.log(chalk.blue("Writing tracer config file..."));
+await writeFile(
+  "/tmp/trace-stack/packages/lambda-layer/config.json",
+  JSON.stringify({ token: tracerToken }),
+);
+
 console.log(chalk.blue("Writing .env file..."));
 await writeFile(
   "/tmp/trace-stack/packages/api/.env",
   `RETENTION_DAYS=${answers.RETENTION_DAYS}\n` +
     `CUSTOM_DOMAIN=${answers.CUSTOM_DOMAIN}\n` +
+    `TRACER_TOKEN=${tracerToken}\n` +
     `HAS_CUSTOM_DOMAIN=${answers.CUSTOM_DOMAIN ? "true" : "false"}\n`,
 );
 
 // Deploy
 console.log(chalk.blue("Deploying..."));
 await exec("yarn deploy", { cwd: "/tmp/trace-stack" });
 
+// Run auto-trace
+const endpoint = await getApiEndpoint();
+await fetch(`https://${endpoint}/api/auto-trace`, {
+  method: "POST",
+  headers: { "Content-Type": "application/json" },
+  body: JSON.stringify({ token: tracerToken }),
+});
+
 // Create user
 console.log(chalk.blue("Creating user..."));
 // TODO: create user account in DDB
 
 // Done
-const domain = answers.CUSTOM_DOMAIN || (await getApiEndpoint());
+const domain = answers.CUSTOM_DOMAIN || endpoint;
 console.log(
   "\n\n" +
     boxen(

packages/lambda-layer/index.js

@@ -1,6 +1,16 @@
+let config;
+try {
+  config = require("./config.json");
+} catch (e) {
+  config = {
+    token: "t_0000000000000000",
+    edgeHost: process.env.AUTO_TRACE_HOST,
+  };
+}
+
 const tracer = require("@lumigo/tracer")({
-  token: "t_0000000000000000",
-  edgeHost: process.env.AUTO_TRACE_HOST,
+  token: config.token,
+  edgeHost: process.env.AUTO_TRACE_HOST || config.edgeHost,
 });
 
 const { load } = require("./lib/aws/aws-user-function.js");