feat(toolchain): build args support for remote builds

cloud/packages/ci-manager/src/build-store.ts

@@ -20,7 +20,13 @@ export class BuildStore {
 		await mkdir(this.tempDir, { recursive: true });
 	}
 
-	createBuild(buildName: string, dockerfilePath: string, environmentId: string): string {
+	createBuild(
+		buildName: string,
+		dockerfilePath: string,
+		environmentId: string,
+		buildArgs: Record<string, string>,
+		buildTarget: string | undefined
+	): string {
 		const id = randomUUID();
 		const contextPath = join(this.tempDir, id, "context.tar.gz");
 		const outputPath = join(this.tempDir, id, "output.tar.gz");
@@ -32,6 +38,8 @@ export class BuildStore {
 			dockerfilePath,
 			environmentId,
 			contextPath,
+			buildArgs,
+			buildTarget,
 			outputPath,
 			events: [],
 			createdAt: new Date(),
@@ -86,7 +94,7 @@ export class BuildStore {
 		}
 	}
 
-	getContextPath(id: string): string | undefined {
+	getContextPath(id: string): string | undefined{
 		return this.builds.get(id)?.contextPath;
 	}
 
@@ -118,17 +126,15 @@ export class BuildStore {
 			}
 
 			// Remove build directory and all files
-			if (build.contextPath) {
-				const buildDir = dirname(build.contextPath);
-				try {
-					await rm(buildDir, { recursive: true, force: true });
-					console.log(`Removed build directory: ${buildDir}`);
-				} catch (error) {
-					console.warn(
-						`Failed to remove build directory ${buildDir}:`,
-						error,
-					);
-				}
+			const buildDir = dirname(build.contextPath);
+			try {
+				await rm(buildDir, { recursive: true, force: true });
+				console.log(`Removed build directory: ${buildDir}`);
+			} catch (error) {
+				console.warn(
+					`Failed to remove build directory ${buildDir}:`,
+					error,
+				);
 			}
 
 			// Remove from memory

cloud/packages/ci-manager/src/common.ts

@@ -0,0 +1,44 @@
+// Keep in sync with ci-runner/entry.sh
+export const UNIT_SEP_CHAR = "\x1F";
+export const NO_SEP_CHAR_REGEX = /^[^\x1F]+$/;
+
+interface KanikoArguments {
+	contextUrl: string;
+	outputUrl: string;
+	destination: string;
+	dockerfilePath: string;
+	buildArgs: Record<string, string>;
+	buildTarget?: string;
+}
+
+// SAFETY: buildArgs keys never have equal signs or spaces
+function convertBuildArgsToArgs(
+	buildArgs: Record<string, string>,
+): string[] {
+	return Object.entries(buildArgs).flatMap(([key, value]) => [
+		`--build-arg`,
+		`${key}=${value}`,
+	]);
+}
+
+export function serializeKanikoArguments(args: KanikoArguments): string {
+	// SAFETY: Nothing needed to be escaped, as values are already sanitized,
+	// and are joined by IFS=UNIT_SEP_CHAR (see entry.sh of ci-runner).
+	const preparedArgs = [
+		...convertBuildArgsToArgs(args.buildArgs),
+		`--context=${args.contextUrl}`,
+		`--destination=${args.destination}`,
+		`--upload-tar=${args.outputUrl}`,
+		`--dockerfile=${args.dockerfilePath}`,
+		...(args.buildTarget ? [`--target='${args.buildTarget}'`] : []),
+		"--no-push",
+		"--single-snapshot",
+		"--verbosity=info",
+	].map(arg => {
+		// Args should never contain UNIT_SEP_CHAR, but we can
+		// escape it if they do.
+		return arg.replaceAll(UNIT_SEP_CHAR, "\\" + UNIT_SEP_CHAR)
+	});
+
+	return preparedArgs.join(UNIT_SEP_CHAR);
+}

cloud/packages/ci-manager/src/executors/docker.ts

@@ -1,5 +1,6 @@
 import { spawn } from "node:child_process";
 import { BuildStore } from "../build-store";
+import { serializeKanikoArguments, UNIT_SEP_CHAR } from "../common";
 
 export async function runDockerBuild(
 	buildStore: BuildStore,
@@ -19,13 +20,16 @@ export async function runDockerBuild(
 		"--rm",
 		"--network=host",
 		"-e",
-		`CONTEXT_URL=${contextUrl}`,
-		"-e",
-		`OUTPUT_URL=${outputUrl}`,
-		"-e",
-		`DESTINATION=${buildId}:latest`,
-		"-e",
-		`DOCKERFILE_PATH=${build.dockerfilePath}`,
+		`KANIKO_ARGS=${
+			serializeKanikoArguments({
+				contextUrl,
+				outputUrl,
+				destination: `${buildId}:latest`,
+				dockerfilePath: build.dockerfilePath,
+				buildArgs: build.buildArgs,
+				buildTarget: build.buildTarget,
+			})
+		}`,
 		"ci-runner",
 	];
 
@@ -36,7 +40,7 @@ export async function runDockerBuild(
 
 	return new Promise<void>((resolve, reject) => {
 		const dockerProcess = spawn("docker", kanikoArgs, {
-			stdio: ["pipe", "pipe", "pipe"],
+			stdio: ["pipe", "pipe", "pipe"]
 		});
 
 		buildStore.setContainerProcess(buildId, dockerProcess);

cloud/packages/ci-manager/src/executors/rivet.ts

@@ -1,5 +1,6 @@
 import { RivetClient } from "@rivet-gg/api";
 import { BuildStore } from "../build-store";
+import { serializeKanikoArguments } from "../common";
 
 export async function runRivetBuild(
 	buildStore: BuildStore,
@@ -49,10 +50,14 @@ export async function runRivetBuild(
 				build: kanikoBuildId,
 				runtime: {
 					environment: {
-						CONTEXT_URL: contextUrl,
-						OUTPUT_URL: outputUrl,
-						DESTINATION: `${buildId}:latest`,
-						DOCKERFILE_PATH: build.dockerfilePath!,
+						KANIKO_ARGS: serializeKanikoArguments({
+							contextUrl,
+							outputUrl,
+							destination: `${buildId}:latest`,
+							dockerfilePath: build.dockerfilePath,
+							buildArgs: build.buildArgs,
+							buildTarget: build.buildTarget,
+						})
 					},
 				},
 				network: {

cloud/packages/ci-manager/src/kaniko-runner.ts

@@ -37,10 +37,10 @@ async function validateBuildUpload(
 	buildId: string,
 ): Promise<void> {
 	const build = buildStore.getBuild(buildId);
-	if (!build || !build.outputPath) {
+	if (!build) {
 		buildStore.updateStatus(buildId, {
 			type: "failure",
-			data: { reason: "Build not found or missing output path" },
+			data: { reason: "Build not found" },
 		});
 		return;
 	}

cloud/packages/ci-manager/src/server.ts

@@ -17,14 +17,16 @@ import {
 	uploadOCIBundleToRivet,
 	type RivetUploadConfig,
 } from "./rivet-uploader";
+import { UNIT_SEP_CHAR } from "./common";
+import { BuildRequestSchema } from "./types";
 
 async function processRivetUpload(
 	buildStore: BuildStore,
 	buildId: string,
 ): Promise<void> {
 	const build = buildStore.getBuild(buildId);
-	if (!build || !build.outputPath) {
-		throw new Error(`Build ${buildId} not found or missing output path`);
+	if (!build) {
+		throw new Error(`Build ${buildId} not found`);
 	}
 
 	try {
@@ -63,7 +65,7 @@ async function processRivetUpload(
 
 			const uploadResult = await uploadOCIBundleToRivet(
 				conversionResult.bundleTarPath,
-				build.buildName!,
+				build.buildName,
 				`${buildId}:latest`, // Match kaniko destination format
 				rivetConfig,
 				new Date().toISOString(), // Use timestamp as version for now
@@ -100,30 +102,31 @@ export async function createServer(port: number = 3000) {
 
 	app.post("/builds", async (c) => {
 		try {
-			const formData = await c.req.formData();
-			const buildName = formData.get("buildName") as string;
-			const dockerfilePath = formData.get("dockerfilePath") as string;
-			const environmentId = formData.get("environmentId") as string;
-			const contextFile = formData.get("context") as File;
-
-			if (!buildName) {
-				return c.json({ error: "buildName is required" }, 400);
-			}
-
-			if (!dockerfilePath) {
-				return c.json({ error: "dockerfilePath is required" }, 400);
-			}
-
-			if (!environmentId) {
-				return c.json({ error: "environmentId is required" }, 400);
-			}
-
-			if (!contextFile) {
-				return c.json({ error: "context file is required" }, 400);
+			const body = await c.req.parseBody();
+			const parseResult = BuildRequestSchema.safeParse(body);
+			if (!parseResult.success) {
+				return c.json(
+					{ error: "Invalid build request format" },
+					400,
+				);
 			}
+			const {
+				buildName,
+				dockerfilePath,
+				environmentId,
+				buildArgs,
+				buildTarget,
+				context: contextFile
+			} = parseResult.data;
 
 			// Create the build
-			const buildId = buildStore.createBuild(buildName, dockerfilePath, environmentId);
+			const buildId = buildStore.createBuild(
+				buildName, 
+				dockerfilePath, 
+				environmentId, 
+				buildArgs,
+				buildTarget
+			);
 			const contextPath = buildStore.getContextPath(buildId);
 
 			if (!contextPath) {
@@ -169,6 +172,7 @@ export async function createServer(port: number = 3000) {
 
 			return c.json({ buildId });
 		} catch (error) {
+			console.error("Error processing build request:", error);
 			return c.json({ error: "Failed to process build request" }, 500);
 		}
 	});

cloud/packages/ci-manager/src/types.ts

@@ -1,43 +1,83 @@
 import { z } from "zod";
+import { NO_SEP_CHAR_REGEX, UNIT_SEP_CHAR } from "./common";
 
 export const StatusSchema = z.discriminatedUnion("type", [
-  z.object({ type: z.literal("starting"), data: z.object({}) }),
-  z.object({ type: z.literal("running"), data: z.object({}) }),
-  z.object({ type: z.literal("finishing"), data: z.object({}) }),
-  z.object({ type: z.literal("converting"), data: z.object({}) }),
-  z.object({ type: z.literal("uploading"), data: z.object({}) }),
-  z.object({ type: z.literal("failure"), data: z.object({ reason: z.string() }) }),
-  z.object({ type: z.literal("success"), data: z.object({ buildId: z.string() }) }),
+    z.object({ type: z.literal("starting"), data: z.object({}) }),
+    z.object({ type: z.literal("running"), data: z.object({}) }),
+    z.object({ type: z.literal("finishing"), data: z.object({}) }),
+    z.object({ type: z.literal("converting"), data: z.object({}) }),
+    z.object({ type: z.literal("uploading"), data: z.object({}) }),
+    z.object({ type: z.literal("failure"), data: z.object({ reason: z.string() }) }),
+    z.object({ type: z.literal("success"), data: z.object({ buildId: z.string() }) }),
 ]);
 
 export type Status = z.infer<typeof StatusSchema>;
 
+const ILLEGAL_BUILD_ARG_KEY = /[\s'"\\]/g;
+const BuildArgsSchema = z.string()
+    .transform((str) => JSON.parse(str))
+    .pipe(z.array(z.string()))
+    .refine((arr) => {
+        // Check each key=value pair to ensure keys have no spaces
+        return arr.every(item => {
+            const [key] = item.split('=');
+            if (!key) return false;
+            if (ILLEGAL_BUILD_ARG_KEY.test(key)) return false;
+            if (item.includes(UNIT_SEP_CHAR)) return false;
+            return true;
+        });
+    }, { message: "Argument key/value contains invalid character" })
+    .transform((arr) => {
+        const result: Record<string, string> = Object.create(null);
+        // Convert array of strings to an object
+        for (const item of arr) {
+            const [key, ...valueParts] = item.split('=');
+            const value = valueParts.join('=');
+
+            if (key && value !== undefined) {
+                result[key] = value;
+            }
+        }
+
+        return result;
+    });
+
 export const BuildRequestSchema = z.object({
-  buildName: z.string(),
-  dockerfilePath: z.string(),
-  environmentId: z.string(),
+    buildName: z.string()
+        .regex(NO_SEP_CHAR_REGEX, "buildName cannot contain special characters"),
+    dockerfilePath: z.string()
+        .regex(NO_SEP_CHAR_REGEX, "dockerfilePath cannot contain special characters"),
+    environmentId: z.string()
+        .regex(NO_SEP_CHAR_REGEX, "environmentId cannot contain special characters"),
+    buildArgs: BuildArgsSchema,
+    buildTarget: z.string()
+        .regex(NO_SEP_CHAR_REGEX, "buildTarget cannot contain special characters")
+        .optional(),
+    context: z.instanceof(File)
 });
 
 export type BuildRequest = z.infer<typeof BuildRequestSchema>;
 
 export const BuildEventSchema = z.discriminatedUnion("type", [
-  z.object({ type: z.literal("status"), data: StatusSchema }),
-  z.object({ type: z.literal("log"), data: z.object({ line: z.string() }) }),
+    z.object({ type: z.literal("status"), data: StatusSchema }),
+    z.object({ type: z.literal("log"), data: z.object({ line: z.string() }) }),
 ]);
 
 export type BuildEvent = z.infer<typeof BuildEventSchema>;
 
 export interface BuildInfo {
-  id: string;
-  status: Status;
-  buildName?: string;
-  dockerfilePath?: string;
-  environmentId?: string;
-  contextPath?: string;
-  outputPath?: string;
-  events: BuildEvent[];
-  containerProcess?: any;
-  createdAt: Date;
-  downloadedAt?: Date;
-  cleanupTimeout?: NodeJS.Timeout;
+    id: string;
+    status: Status;
+    buildName: string;
+    dockerfilePath: string;
+    environmentId: string;
+    contextPath: string;
+    buildArgs: Record<string, string>;
+    buildTarget?: string;
+    outputPath: string;
+    events: BuildEvent[];
+    containerProcess?: any;
+    createdAt: Date;
+    downloadedAt?: Date;
+    cleanupTimeout?: NodeJS.Timeout;
 }

cloud/packages/ci-runner/Dockerfile

@@ -5,5 +5,6 @@ FROM ghcr.io/rivet-gg/executor@sha256:439d4dbb0f3f8c1c6c2195e144d29195b4930b8716
 COPY --from=builder /bin/sh /bin/sh
 COPY --from=builder /lib/ld-musl-x86_64.so.1 /lib/ld-musl-x86_64.so.1
 
-# HACK: Use env vars to interpolate args bc Rivet doesn't support passing args
-ENTRYPOINT ["/bin/sh", "-c", "/kaniko/executor --context=${CONTEXT_URL} --destination=${DESTINATION} --upload-tar=${OUTPUT_URL} --dockerfile=${DOCKERFILE_PATH} --no-push --single-snapshot --verbosity=info"]
+COPY entry.sh ~/entry.sh
+
+ENTRYPOINT ["/bin/sh", "~/entry.sh"]