From c4c52ed2fc7f82cdf258c315c72c7267c64d69d8 Mon Sep 17 00:00:00 2001
From: Jacob Gillespie <jacobwgillespie@gmail.com>
Date: Fri, 19 Apr 2024 21:15:45 +0100
Subject: [PATCH] Add support for using CNI

---
 proto/depot/cloud/v3/machine.proto           |  1 +
 src/gen/ts/depot/cloud/v3/machine_connect.ts |  2 +-
 src/gen/ts/depot/cloud/v3/machine_pb.ts      |  6 ++++++
 src/tasks/buildkit.ts                        | 16 ++++++++++++++++
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/proto/depot/cloud/v3/machine.proto b/proto/depot/cloud/v3/machine.proto
index dc0f886..7c8ce21 100644
--- a/proto/depot/cloud/v3/machine.proto
+++ b/proto/depot/cloud/v3/machine.proto
@@ -86,6 +86,7 @@ message RegisterMachineResponse {
     // when all customers are using context logging.
     optional bool enable_context_logging = 17;
     optional string vector_config = 18;
+    optional bool enable_cni = 19;
   }
 
   // EngineTask represents an instruction to start an engine daemon
diff --git a/src/gen/ts/depot/cloud/v3/machine_connect.ts b/src/gen/ts/depot/cloud/v3/machine_connect.ts
index 1f45ed8..a4c6298 100644
--- a/src/gen/ts/depot/cloud/v3/machine_connect.ts
+++ b/src/gen/ts/depot/cloud/v3/machine_connect.ts
@@ -1,4 +1,4 @@
-// @generated by protoc-gen-connect-es v0.13.0 with parameter "target=ts,import_extension=none"
+// @generated by protoc-gen-connect-es v0.12.0 with parameter "target=ts,import_extension=none"
 // @generated from file depot/cloud/v3/machine.proto (package depot.cloud.v3, syntax proto3)
 /* eslint-disable */
 // @ts-nocheck
diff --git a/src/gen/ts/depot/cloud/v3/machine_pb.ts b/src/gen/ts/depot/cloud/v3/machine_pb.ts
index b875cb1..f4fcb89 100644
--- a/src/gen/ts/depot/cloud/v3/machine_pb.ts
+++ b/src/gen/ts/depot/cloud/v3/machine_pb.ts
@@ -500,6 +500,11 @@ export class RegisterMachineResponse_BuildKitTask extends Message<RegisterMachin
    */
   vectorConfig?: string
 
+  /**
+   * @generated from field: optional bool enable_cni = 19;
+   */
+  enableCni?: boolean
+
   constructor(data?: PartialMessage<RegisterMachineResponse_BuildKitTask>) {
     super()
     proto3.util.initPartial(data, this)
@@ -525,6 +530,7 @@ export class RegisterMachineResponse_BuildKitTask extends Message<RegisterMachin
     {no: 16, name: 'enable_debug_logging', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
     {no: 17, name: 'enable_context_logging', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
     {no: 18, name: 'vector_config', kind: 'scalar', T: 9 /* ScalarType.STRING */, opt: true},
+    {no: 19, name: 'enable_cni', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
   ])
 
   static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): RegisterMachineResponse_BuildKitTask {
diff --git a/src/tasks/buildkit.ts b/src/tasks/buildkit.ts
index da0a659..43b038e 100644
--- a/src/tasks/buildkit.ts
+++ b/src/tasks/buildkit.ts
@@ -102,6 +102,22 @@ keepBytes = ${cacheSizeBytes}
 `
   await fsp.writeFile('/etc/buildkit/buildkitd.toml', config, {mode: 0o644})
 
+  if (task.enableCni) {
+    const cniConfig = {
+      cniVersion: '1.0.0',
+      name: 'buildkit',
+      type: 'bridge',
+      bridge: 'buildkit0',
+      isDefaultGateway: true,
+      forceAddress: false,
+      ipMasq: true,
+      hairpinMode: true,
+      ipam: {type: 'host-local', ranges: [[{subnet: '192.168.0.0/16'}]]},
+    }
+    await fsp.mkdir('/etc/buildkit', {recursive: true})
+    await fsp.writeFile('/etc/buildkit/cni.json', JSON.stringify(cniConfig, null, 2), {mode: 0o644})
+  }
+
   const controller = new AbortController()
   const signal = controller.signal