Implement importable env

src/cloudflare/internal/env.d.ts

@@ -0,0 +1,3 @@
+// Get the current environment, if any
+export function getCurrent(): object | undefined;
+export function withEnv(newEnv: unknown, fn: () => unknown): unknown;

src/cloudflare/workers.ts

@@ -6,9 +6,92 @@
 //   wrapper module that simply re-exports the classes from the built-in module.
 
 import entrypoints from 'cloudflare-internal:workers';
+import innerEnv from 'cloudflare-internal:env';
 
 export const WorkerEntrypoint = entrypoints.WorkerEntrypoint;
 export const DurableObject = entrypoints.DurableObject;
 export const RpcStub = entrypoints.RpcStub;
 export const RpcTarget = entrypoints.RpcTarget;
 export const WorkflowEntrypoint = entrypoints.WorkflowEntrypoint;
+
+// We turn off these specific linting rules here intentionally as mixing
+// typescript and proxies can just be... awkward. The use of any in these
+// is intentional.
+/* eslint-disable @typescript-eslint/no-unsafe-assignment */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+export function withEnv(newEnv: unknown, fn: () => unknown): unknown {
+  return innerEnv.withEnv(newEnv, fn);
+}
+
+// A proxy for the workers env/bindings. The proxy is io-context
+// aware in that it will only return values when there is an active
+// IoContext. Mutations to the env via this proxy propagate to the
+// underlying env and do not follow the async context.
+export const env: object = new Proxy(
+  {},
+  {
+    get(_: any, prop: string | symbol): unknown {
+      const inner = innerEnv.getCurrent();
+      if (inner) {
+        return Reflect.get(inner, prop);
+      }
+      return undefined;
+    },
+
+    set(_: any, prop: string | symbol, newValue: any): boolean {
+      const inner = innerEnv.getCurrent();
+      if (inner) {
+        return Reflect.set(inner, prop, newValue);
+      }
+      return true;
+    },
+
+    has(_: any, prop: string | symbol): boolean {
+      const inner = innerEnv.getCurrent();
+      if (inner) {
+        return Reflect.has(inner, prop);
+      }
+      return false;
+    },
+
+    ownKeys(_: any): any {
+      const inner = innerEnv.getCurrent();
+      if (inner) {
+        return Reflect.ownKeys(inner);
+      }
+      return [];
+    },
+
+    deleteProperty(_: any, prop: string | symbol): boolean {
+      const inner = innerEnv.getCurrent();
+      if (inner) {
+        return Reflect.deleteProperty(inner, prop);
+      }
+      return true;
+    },
+
+    defineProperty(
+      _: any,
+      prop: string | symbol,
+      attr: PropertyDescriptor
+    ): boolean {
+      const inner = innerEnv.getCurrent();
+      if (inner) {
+        return Reflect.defineProperty(inner, prop, attr);
+      }
+      return true;
+    },
+
+    getOwnPropertyDescriptor(
+      _: any,
+      prop: string | symbol
+    ): PropertyDescriptor | undefined {
+      const inner = innerEnv.getCurrent();
+      if (inner) {
+        return Reflect.getOwnPropertyDescriptor(inner, prop);
+      }
+      return undefined;
+    },
+  }
+);

src/workerd/api/BUILD.bazel

@@ -19,6 +19,7 @@ filegroup(
             "pyodide/pyodide.c++",
             "pyodide/setup-emscripten.c++",
             "memory-cache.c++",
+            "modules.c++",
             "r2*.c++",
             "rtti.c++",
             "url.c++",
@@ -84,7 +85,10 @@ wd_cc_library(
 
 wd_cc_library(
     name = "rtti",
-    srcs = ["rtti.c++"],
+    srcs = [
+        "modules.c++",
+        "rtti.c++",
+    ],
     hdrs = [
         "modules.h",
         "rtti.h",
@@ -650,3 +654,15 @@ wd_test(
     args = ["--experimental"],
     data = ["tests/fetch-test.js"],
 )
+
+wd_test(
+    src = "tests/importable-env-test.wd-test",
+    args = ["--experimental"],
+    data = ["tests/importable-env-test.js"],
+)
+
+wd_test(
+    src = "tests/disable-importable-env-test.wd-test",
+    args = ["--experimental"],
+    data = ["tests/disable-importable-env-test.js"],
+)

src/workerd/api/modules.c++

@@ -0,0 +1,36 @@
+#include "modules.h"
+
+#include <workerd/jsg/setup.h>
+
+namespace workerd::api {
+
+kj::Maybe<jsg::JsObject> EnvModule::getCurrent(jsg::Lock& js) {
+  auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
+  KJ_IF_SOME(frame, jsg::AsyncContextFrame::current(js)) {
+    KJ_IF_SOME(value, frame.get(key)) {
+      auto handle = value.getHandle(js);
+      if (handle->IsObject()) {
+        return jsg::JsObject(handle.As<v8::Object>());
+      }
+    }
+  }
+  // If the compat flag is set to disable importable env, then this
+  // will return nothing.
+  if (FeatureFlags::get(js).getDisableImportableEnv()) return kj::none;
+
+  // Otherwise, fallback to provide the stored environment.
+  return js.getWorkerEnv().map([&](const jsg::Value& val) -> jsg::JsObject {
+    auto handle = val.getHandle(js);
+    JSG_REQUIRE(handle->IsObject(), TypeError, "Expected environment to be an object.");
+    return jsg::JsObject(handle.As<v8::Object>());
+  });
+}
+
+jsg::JsValue EnvModule::withEnv(
+    jsg::Lock& js, jsg::Value newEnv, jsg::Function<jsg::JsValue()> fn) {
+  auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
+  jsg::AsyncContextFrame::StorageScope storage(js, key, kj::mv(newEnv));
+  return fn(js);
+}
+
+}  // namespace workerd::api

src/workerd/api/modules.h

@@ -10,13 +10,26 @@
 #include <workerd/api/sockets.h>
 #include <workerd/api/unsafe.h>
 #include <workerd/api/worker-rpc.h>
+#include <workerd/io/features.h>
 #include <workerd/io/worker.h>
 #include <workerd/jsg/modules-new.h>
 
 #include <cloudflare/cloudflare.capnp.h>
 
 namespace workerd::api {
 
+class EnvModule final: public jsg::Object {
+ public:
+  kj::Maybe<jsg::JsObject> getCurrent(jsg::Lock& js);
+
+  jsg::JsValue withEnv(jsg::Lock& js, jsg::Value newEnv, jsg::Function<jsg::JsValue()> fn);
+
+  JSG_RESOURCE_TYPE(EnvModule) {
+    JSG_METHOD(getCurrent);
+    JSG_METHOD(withEnv);
+  }
+};
+
 template <class Registry>
 void registerModules(Registry& registry, auto featureFlags) {
   node::registerNodeJsCompatModules(registry, featureFlags);
@@ -30,6 +43,8 @@ void registerModules(Registry& registry, auto featureFlags) {
   registerSocketsModule(registry, featureFlags);
   registry.addBuiltinBundle(CLOUDFLARE_BUNDLE);
   registerRpcModules(registry, featureFlags);
+  registry.template addBuiltinModule<EnvModule>(
+      "cloudflare-internal:env", workerd::jsg::ModuleRegistry::Type::INTERNAL);
 }
 
 template <class TypeWrapper>
@@ -61,6 +76,8 @@ void registerBuiltinModules(jsg::modules::ModuleRegistry::Builder& builder, auto
     jsg::modules::ModuleBundle::getBuiltInBundleFromCapnp(builtinsBuilder, CLOUDFLARE_BUNDLE);
     builder.add(builtinsBuilder.finish());
   }
+
+  // TODO(later): Add the internal env module also.
 }
 
 }  // namespace workerd::api

src/workerd/api/node/util.c++

@@ -242,7 +242,7 @@ jsg::JsValue UtilModule::getBuiltinModule(jsg::Lock& js, kj::String specifier) {
 }
 
 jsg::JsObject UtilModule::getEnvObject(jsg::Lock& js) {
-  return js.getEnv(true);
+  return js.getProcessEnv(true);
 }
 
 namespace {

src/workerd/api/tests/disable-importable-env-test.js

@@ -0,0 +1,25 @@
+import { env, withEnv } from 'cloudflare:workers';
+
+// This test runs with the disallow-importable-env flag set, meaning that
+// the env imported from cloudflare:workers exists but is not-populated
+// with the primary environment. Using withEnv, however, would still work.
+
+import { strictEqual, notDeepStrictEqual } from 'node:assert';
+
+strictEqual(env.FOO, undefined);
+
+export const test = {
+  async test(_, argEnv) {
+    strictEqual(env.FOO, undefined);
+    notDeepStrictEqual(argEnv, env);
+
+    // withEnv still works as expected tho
+    const { env: otherEnv2 } = await withEnv({ BAZ: 1 }, async () => {
+      await scheduler.wait(0);
+      return import('child');
+    });
+    strictEqual(otherEnv2.FOO, undefined);
+    strictEqual(otherEnv2.BAZ, 1);
+    strictEqual(argEnv.BAZ, undefined);
+  },
+};

src/workerd/api/tests/disable-importable-env-test.wd-test

@@ -0,0 +1,22 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "disable-importable-env-test",
+      worker = (
+        modules = [
+          (name = "worker", esModule = embed "disable-importable-env-test.js"),
+          (name = "child", esModule = "import {env as live} from 'cloudflare:workers'; export const env = {...live};"),
+        ],
+        compatibilityDate = "2025-02-01",
+        compatibilityFlags = [
+          "nodejs_compat_v2",
+          "disallow_importable_env",
+        ],
+        bindings = [
+          (name = "FOO", text = "BAR"),
+        ],
+      )
+    ),
+  ],
+);

src/workerd/api/tests/importable-env-test.js

@@ -0,0 +1,101 @@
+import {
+  strictEqual,
+  deepStrictEqual,
+  notStrictEqual,
+  ok,
+  rejects,
+} from 'node:assert';
+import { env, withEnv } from 'cloudflare:workers';
+
+import { AsyncLocalStorage } from 'node:async_hooks';
+const check = new AsyncLocalStorage();
+
+// The env is populated at the top level scope.
+strictEqual(env.FOO, 'BAR');
+
+// Cache exists and is accessible.
+ok(env.CACHE);
+// But fails when used because we're not in an io-context
+await rejects(
+  env.CACHE.read('hello', async () => {}),
+  {
+    message: /Disallowed operation called within global scope./,
+  }
+);
+
+export const importableEnv = {
+  async test(_, argEnv) {
+    // Accessing the cache initially at the global scope didn't break anything
+    const cached = await argEnv.CACHE.read('hello', async () => {
+      return {
+        value: 123,
+        expiration: Date.now() + 10000,
+      };
+    });
+    strictEqual(cached, 123);
+
+    // They aren't the same objects...
+    notStrictEqual(env, argEnv);
+    // But have all the same stuff...
+    deepStrictEqual(env, argEnv);
+
+    // It is populated inside a request
+    strictEqual(env.FOO, 'BAR');
+
+    // And following async operations.
+    await scheduler.wait(10);
+    strictEqual(env.FOO, 'BAR');
+
+    // Mutations to the env carry through as expected...
+    env.BAR = 123;
+    const { env: otherEnv } = await import('child');
+    strictEqual(otherEnv.FOO, 'BAR');
+    strictEqual(otherEnv.BAR, 123);
+    deepStrictEqual(argEnv, otherEnv);
+
+    // Using withEnv to replace the env...
+    const { env: otherEnv2 } = await withEnv({ BAZ: 1 }, async () => {
+      await scheduler.wait(0);
+      return import('child2');
+    });
+    strictEqual(otherEnv2.FOO, undefined);
+    strictEqual(otherEnv2.BAZ, 1);
+
+    // Original env is unmodified
+    strictEqual(env.BAZ, undefined);
+
+    // Verify that JSRPC calls appropriately see the environment.
+    const remote = await argEnv.RPC.rpcTarget(undefined);
+    await Promise.all([remote.test(), remote.test2()]);
+  },
+
+  get fetch() {
+    // It's a weird edge case, yes, but let's make sure that the env is
+    // available in getters when extracting the default handlers.
+    strictEqual(env.FOO, 'BAR');
+    return async () => {};
+  },
+
+  // Verifies that the environment is available and correctly propagated
+  // in custom events like jsrpc.
+  rpcTarget() {
+    strictEqual(env.FOO, 'BAR');
+    return withEnv({ FOO: 'BAZ' }, () => {
+      // Arguably, returned RPC targets should probably automatically capture
+      // the current async context and propagate that on subsequent calls,
+      // but they currently do not... therefore we have to manually capture
+      // the async context and be sure to enter it ourselves below.
+      const runInAsyncContext = AsyncLocalStorage.snapshot();
+      return {
+        test() {
+          // This one runs with the modified env
+          runInAsyncContext(() => strictEqual(env.FOO, 'BAZ'));
+        },
+        test2() {
+          // This one runs with the original env
+          strictEqual(env.FOO, 'BAR');
+        },
+      };
+    });
+  },
+};