Use opentelemetry to add trace ids to request logs

package.json

@@ -20,6 +20,10 @@
     "@balena/env-parsing": "^1.2.0",
     "@balena/es-version": "^1.0.3",
     "@balena/node-metrics-gatherer": "^6.0.3",
+    "@opentelemetry/api": "^1.9.0",
+    "@opentelemetry/instrumentation-express": "^0.47.1",
+    "@opentelemetry/instrumentation-http": "^0.57.2",
+    "@opentelemetry/sdk-node": "^0.57.2",
     "@sentry/node": "^8.30.0",
     "bluebird": "^3.7.2",
     "compression": "^1.7.4",

src/api.ts

@@ -22,7 +22,7 @@ import express from 'express';
 import memoize from 'memoizee';
 import morgan from 'morgan';
 
-import { getLogger } from './utils/index.js';
+import { getLogger, getPassthrough } from './utils/index.js';
 import {
 	BALENA_API_INTERNAL_HOST,
 	DELAY_ON_AUTH_FAIL,
@@ -35,6 +35,7 @@ import { setTimeout } from 'timers/promises';
 import { pooledRequest } from './utils/request.js';
 import { Metrics } from './utils/metrics.js';
 import { setConnected } from './utils/clients.js';
+import { trace } from '@opentelemetry/api';
 
 // Private endpoints should use the `fromLocalHost` middleware.
 const fromLocalHost: express.RequestHandler = (req, res, next) => {
@@ -50,7 +51,7 @@ const checkDeviceAuth = memoize(
 	async (username: string, password: string) => {
 		const { statusCode } = await pooledRequest.get({
 			url: `${BALENA_API_INTERNAL_HOST}/services/vpn/auth/${username}`,
-			headers: { Authorization: `Bearer ${password}` },
+			...getPassthrough(`Bearer ${password}`),
 		});
 		if ([200, 401, 403].includes(statusCode)) {
 			return statusCode;
@@ -178,11 +179,12 @@ export const apiServer = (serviceId: number) => {
 		morgan((tokens, req, res) => {
 			const date = new Date().toISOString();
 			const ip = tokens['remote-addr'](req, res);
+			const traceId = trace.getActiveSpan()?.spanContext().traceId ?? '-';
 			const url = tokens.url(req, res);
 			const statusCode = tokens.status(req, res) ?? '-';
 			const responseTime = tokens['response-time'](req, res) ?? '-';
 
-			return `${date} ${ip} ${
+			return `${date} ${ip} ${traceId} ${
 				req.method
 			} ${url} ${statusCode} ${responseTime}ms`;
 		}),

src/app.ts

@@ -15,9 +15,7 @@
 	along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
-import { set } from '@balena/es-version';
-// Set the desired es version for downstream modules that support it, before we import any
-set('es2021');
+import './init.js';
 
 import { metrics } from '@balena/node-metrics-gatherer';
 import cluster from 'cluster';

src/init.ts

@@ -0,0 +1,29 @@
+/*
+	Copyright (C) 2025 Balena Ltd.
+
+	This program is free software: you can redistribute it and/or modify
+	it under the terms of the GNU Affero General Public License as published
+	by the Free Software Foundation, either version 3 of the License, or
+	(at your option) any later version.
+
+	This program is distributed in the hope that it will be useful,
+	but WITHOUT ANY WARRANTY; without even the implied warranty of
+	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+	GNU Affero General Public License for more details.
+
+	You should have received a copy of the GNU Affero General Public License
+	along with this program.  If not, see <https://www.gnu.org/licenses/>.
+*/
+
+import { set } from '@balena/es-version';
+// Set the desired es version for downstream modules that support it, before we import any
+set('es2021');
+
+import { NodeSDK } from '@opentelemetry/sdk-node';
+import { ExpressInstrumentation } from '@opentelemetry/instrumentation-express';
+import { HttpInstrumentation } from '@opentelemetry/instrumentation-http';
+
+const sdk = new NodeSDK({
+	instrumentations: [new ExpressInstrumentation(), new HttpInstrumentation()],
+});
+sdk.start();

src/proxy-worker.ts

@@ -34,6 +34,7 @@ import {
 	getDeviceByUUID,
 	getDeviceVpnHost,
 } from './utils/device.js';
+import { context, propagation, trace } from '@opentelemetry/api';
 
 const HTTP_500 = 'HTTP/1.0 500 Internal Server Error\r\n\r\n';
 
@@ -57,7 +58,27 @@ class Tunnel extends nodeTunnel.Tunnel {
 			}
 		});
 
-		this.use(this.tunnelToDevice);
+		const tunnelTracer = trace.getTracer('node-tunnel');
+		this.use(async (...args) => {
+			const [req] = args;
+			const extractedContext = propagation.extract(
+				context.active(),
+				req.headers,
+			);
+
+			await tunnelTracer.startActiveSpan(
+				req.method ?? 'Unknown',
+				{},
+				extractedContext,
+				async (span) => {
+					try {
+						await this.tunnelToDevice(...args);
+					} finally {
+						span.end();
+					}
+				},
+			);
+		});
 	}
 
 	public connect = async (
@@ -152,8 +173,8 @@ class Tunnel extends nodeTunnel.Tunnel {
 			throw new errors.BadRequestError();
 		}
 
-		const match = req.url.match(
-			/^([a-fA-F0-9]+)\.(balena|resin|vpn)(?::([0-9]+))?$/,
+		const match = /^([a-fA-F0-9]+)\.(balena|resin|vpn)(?::([0-9]+))?$/.exec(
+			req.url,
 		);
 		if (match == null) {
 			throw new errors.InvalidHostnameError(`invalid hostname: ${req.url}`);
@@ -171,12 +192,7 @@ class Tunnel extends nodeTunnel.Tunnel {
 		return { uuid, port: parseInt(port, 10), auth };
 	};
 
-	private tunnelToDevice: nodeTunnel.Middleware = async (
-		req,
-		cltSocket,
-		_head,
-		next,
-	) => {
+	private tunnelToDevice = (async (req, cltSocket, _head, next) => {
 		try {
 			const { uuid, port, auth } = this.parseRequest(req);
 			this.logger.info(`tunnel requested to ${uuid}:${port}`);
@@ -215,7 +231,7 @@ class Tunnel extends nodeTunnel.Tunnel {
 				cltSocket.end(HTTP_500);
 			}
 		}
-	};
+	}) satisfies nodeTunnel.Middleware;
 
 	private forwardRequest = (
 		vpnHost: string,

src/utils/clients.ts

@@ -37,6 +37,7 @@ import {
 	VPN_SERVICE_API_KEY,
 } from './config.js';
 import { captureException } from './errors.js';
+import { getPassthrough } from './index.js';
 
 interface DeviceStateTracker {
 	currentConnected?: boolean;
@@ -72,7 +73,7 @@ export const setConnected = (() => {
 								uuids: uuidChunk,
 								connected,
 							},
-							headers: { Authorization: `Bearer ${VPN_SERVICE_API_KEY}` },
+							...getPassthrough(`Bearer ${VPN_SERVICE_API_KEY}`),
 						})
 						.promise()
 						.timeout(REQUEST_TIMEOUT);

src/utils/device.ts

@@ -18,19 +18,17 @@
 import { optionalVar } from '@balena/env-parsing';
 import memoize from 'memoizee';
 
-import { balenaApi, StatusError } from './index.js';
+import { balenaApi, getPassthrough, StatusError } from './index.js';
 import { APIError, captureException } from './errors.js';
 
 const VPN_GUEST_API_KEY = optionalVar('VPN_GUEST_API_KEY');
 
-const authHeader = (auth?: Buffer): { Authorization?: string } => {
-	const headers: { Authorization?: string } = {};
+const authHeader = (auth?: Buffer): string | undefined => {
 	if (auth != null) {
-		headers.Authorization = `Bearer ${auth}`;
+		return `Bearer ${auth}`;
 	} else if (VPN_GUEST_API_KEY != null) {
-		headers.Authorization = `Bearer ${VPN_GUEST_API_KEY}`;
+		return `Bearer ${VPN_GUEST_API_KEY}`;
 	}
-	return headers;
 };
 
 export interface DeviceInfo {
@@ -55,9 +53,11 @@ export const getDeviceByUUID = async (
 	auth?: Buffer,
 ): Promise<DeviceInfo> => {
 	try {
-		const devices = await getDeviceByUUIDQuery({ uuid }, undefined, {
-			headers: authHeader(auth),
-		});
+		const devices = await getDeviceByUUIDQuery(
+			{ uuid },
+			undefined,
+			getPassthrough(authHeader(auth)),
+		);
 		if (!Array.isArray(devices) || devices.length === 0) {
 			throw new Error('invalid api response');
 		}
@@ -88,7 +88,7 @@ const $canAccessDevice = async (
 			{
 				action: { or: ['tunnel-any', `tunnel-${port}`] },
 			},
-			{ headers: authHeader(auth) },
+			getPassthrough(authHeader(auth)),
 		)) as { d?: Array<{ id: number }> };
 		return Array.isArray(d) && d.length === 1 && d[0].id === device.id;
 	} catch {
@@ -124,7 +124,7 @@ export const getDeviceVpnHost = async (
 					},
 				},
 			},
-			passthrough: { headers: authHeader(auth) },
+			passthrough: getPassthrough(authHeader(auth)),
 		})) as Array<{ id: number; ip_address: string }>;
 		return services[0];
 	} catch (err) {

src/utils/index.ts

@@ -22,12 +22,24 @@ export { StatusError } from 'pinejs-client-request';
 import { BALENA_API_INTERNAL_HOST } from './config.js';
 
 import packageJSON from '../../package.json' with { type: 'json' };
+import { context, propagation } from '@opentelemetry/api';
 export const VERSION = packageJSON.version;
 
 export const balenaApi = new PinejsClientRequest({
 	apiPrefix: `${BALENA_API_INTERNAL_HOST}/v6/`,
 });
 
+export const getPassthrough = (auth: string | undefined) => {
+	const headers: Record<string, string> = {};
+	// Propogate the active trace context to the api
+	propagation.inject(context.active(), headers);
+
+	if (auth != null) {
+		headers.Authorization = auth;
+	}
+	return { headers };
+};
+
 export const getLogger = (
 	service: string,
 	serviceId?: number,

src/utils/service.ts

@@ -16,7 +16,7 @@
 */
 
 import { setTimeout } from 'timers/promises';
-import { balenaApi } from './index.js';
+import { balenaApi, getPassthrough } from './index.js';
 import { VPN_SERVICE_API_KEY } from './config.js';
 import { captureException, ServiceRegistrationError } from './errors.js';
 
@@ -40,9 +40,7 @@ class ServiceInstance {
 		try {
 			const { id } = (await balenaApi.post({
 				resource: 'service_instance',
-				passthrough: {
-					headers: { Authorization: `Bearer ${VPN_SERVICE_API_KEY}` },
-				},
+				passthrough: getPassthrough(`Bearer ${VPN_SERVICE_API_KEY}`),
 				body: ipAddress != null ? { ip_address: ipAddress } : {},
 			})) as { id?: number };
 			if (id == null) {
@@ -77,9 +75,7 @@ class ServiceInstance {
 					// Just indicate being online, api handles the timestamp with hooks
 					is_alive: true,
 				},
-				passthrough: {
-					headers: { Authorization: `Bearer ${VPN_SERVICE_API_KEY}` },
-				},
+				passthrough: getPassthrough(`Bearer ${VPN_SERVICE_API_KEY}`),
 			});
 			return true;
 		} catch (err) {