feat: add support for socket.dev firewall client

17 files changed, 226 insertions, 17 deletions

diff --git a/lua/mason-core/async/init.lua b/lua/mason-core/async/init.lua
index 03963264..e3a2e850 100644
--- a/lua/mason-core/async/init.lua
+++ b/lua/mason-core/async/init.lua
@@ -78,7 +78,8 @@ local function new_execution_context(suspend_fn, callback, ...)
         if cancelled or not thread then
             return
         end
-        local ok, promise_or_result = co.resume(thread, ...)
+        local results = { co.resume(thread, ...) }
+        local ok, promise_or_result = results[1], results[2]
         if cancelled or not thread then
             return
         end
@@ -88,7 +89,7 @@ local function new_execution_context(suspend_fn, callback, ...)
                     promise_or_result(step)
                 else
                     -- yield to parent coroutine
-                    step(coroutine.yield(promise_or_result))
+                    step(coroutine.yield(promise_or_result, unpack(results, 3)))
                 end
             else
                 callback(true, promise_or_result)
diff --git a/lua/mason-core/installer/InstallHandle.lua b/lua/mason-core/installer/InstallHandle.lua
index 3846659e..6492acd9 100644
--- a/lua/mason-core/installer/InstallHandle.lua
+++ b/lua/mason-core/installer/InstallHandle.lua
@@ -20,6 +20,7 @@ local uv = vim.loop
 ---@field pid integer
 ---@field cmd string
 ---@field args string[]
+---@field firewall boolean
 local InstallHandleSpawnHandle = {}
 InstallHandleSpawnHandle.__index = InstallHandleSpawnHandle
 
@@ -27,7 +28,8 @@ InstallHandleSpawnHandle.__index = InstallHandleSpawnHandle
 ---@param pid integer
 ---@param cmd string
 ---@param args string[]
-function InstallHandleSpawnHandle:new(luv_handle, pid, cmd, args)
+---@param firewall boolean
+function InstallHandleSpawnHandle:new(luv_handle, pid, cmd, args, firewall)
     ---@type InstallHandleSpawnHandle
     local instance = {}
     setmetatable(instance, InstallHandleSpawnHandle)
@@ -35,6 +37,7 @@ function InstallHandleSpawnHandle:new(luv_handle, pid, cmd, args)
     instance.pid = pid
     instance.cmd = cmd
     instance.args = args
+    instance.firewall = firewall
     return instance
 end
 
@@ -73,8 +76,9 @@ end
 ---@param pid integer
 ---@param cmd string
 ---@param args string[]
-function InstallHandle:register_spawn_handle(luv_handle, pid, cmd, args)
-    local spawn_handles = InstallHandleSpawnHandle:new(luv_handle, pid, cmd, args)
+---@param firewall boolean
+function InstallHandle:register_spawn_handle(luv_handle, pid, cmd, args, firewall)
+    local spawn_handles = InstallHandleSpawnHandle:new(luv_handle, pid, cmd, args, firewall)
     log.fmt_trace("Pushing spawn_handles stack for %s: %s (pid: %s)", self, spawn_handles, pid)
     self.spawn_handles[#self.spawn_handles + 1] = spawn_handles
     self:emit "spawn_handles:change"
diff --git a/lua/mason-core/installer/InstallRunner.lua b/lua/mason-core/installer/InstallRunner.lua
index 3eba879a..336f0c93 100644
--- a/lua/mason-core/installer/InstallRunner.lua
+++ b/lua/mason-core/installer/InstallRunner.lua
@@ -39,7 +39,14 @@ function InstallRunner:execute(opts, callback)
     local handle = self.handle
     log.fmt_info("Executing installer for %s %s", handle.package, opts)
 
-    local context = InstallContext:new(handle, opts)
+    local context = InstallContext:new(handle, opts, {
+        suspend = function()
+            self:suspend()
+        end,
+        resume = function()
+            self:resume()
+        end,
+    })
 
     local tailed_output = {}
 
@@ -212,6 +219,20 @@ function InstallRunner:acquire_permit()
     return Result.success(channel)
 end
 
+function InstallRunner:suspend()
+    if self.global_permit then
+        self.global_permit:forget()
+        self.global_permit = nil
+    end
+end
+
+---@async
+function InstallRunner:resume()
+    self.handle:set_state "QUEUED"
+    self.global_permit = self.global_semaphore:acquire()
+    self.handle:set_state "ACTIVE"
+end
+
 ---@private
 function InstallRunner:release_permit()
     if self.global_permit then
diff --git a/lua/mason-core/installer/context/InstallContextSpawn.lua b/lua/mason-core/installer/context/InstallContextSpawn.lua
index 29e62101..25419980 100644
--- a/lua/mason-core/installer/context/InstallContextSpawn.lua
+++ b/lua/mason-core/installer/context/InstallContextSpawn.lua
@@ -22,7 +22,7 @@ end
 
 ---@param cmd string
 function InstallContextSpawn:__index(cmd)
-    ---@param args JobSpawnOpts
+    ---@param args SpawnArgs
     return function(args)
         args.cwd = args.cwd or self.cwd:get()
         args.stdio_sink = args.stdio_sink or self.handle.stdio_sink
@@ -30,7 +30,7 @@ function InstallContextSpawn:__index(cmd)
         local captured_handle
         args.on_spawn = function(handle, stdio, pid, ...)
             captured_handle = handle
-            self.handle:register_spawn_handle(handle, pid, cmd, spawn._flatten_cmd_args(args))
+            self.handle:register_spawn_handle(handle, pid, cmd, spawn._flatten_cmd_args(args), args.firewall == true)
             if on_spawn then
                 on_spawn(handle, stdio, pid, ...)
             end
diff --git a/lua/mason-core/installer/context/init.lua b/lua/mason-core/installer/context/init.lua
index 9225315f..69738911 100644
--- a/lua/mason-core/installer/context/init.lua
+++ b/lua/mason-core/installer/context/init.lua
@@ -21,19 +21,22 @@ local receipt = require "mason-core.receipt"
 ---@field cwd InstallContextCwd
 ---@field opts PackageInstallOpts
 ---@field stdio_sink StdioSink
+---@field runner { suspend: fun(), resume: async fun() }
 ---@field links { bin: table<string, string>, share: table<string, string>, opt: table<string, string> }
 local InstallContext = {}
 InstallContext.__index = InstallContext
 
 ---@param handle InstallHandle
 ---@param opts PackageInstallOpts
-function InstallContext:new(handle, opts)
+---@param runner { suspend: fun(), resume: async fun() }
+function InstallContext:new(handle, opts, runner)
     local cwd = InstallContextCwd:new(handle)
     local spawn = InstallContextSpawn:new(handle, cwd, false)
     local fs = InstallContextFs:new(cwd)
     return setmetatable({
         cwd = cwd,
         spawn = spawn,
+        runner = runner,
         handle = handle,
         location = handle.location, -- for convenience
         package = handle.package, -- for convenience
@@ -264,6 +267,7 @@ function InstallContext:link_bin(executable, rel_path)
 end
 
 InstallContext.CONTEXT_REQUEST = {}
+InstallContext.ABORT = {}
 
 ---@generic T
 ---@param fn fun(context: InstallContext): T
@@ -277,14 +281,17 @@ function InstallContext:execute(fn)
     local step
     local ret_val
     step = function(...)
-        local ok, result = coroutine.resume(thread, ...)
+        local results = { coroutine.resume(thread, ...) }
+        local ok, result = results[1], results[2]
         if not ok then
             error(result, 0)
         elseif result == InstallContext.CONTEXT_REQUEST then
             step(self)
+        elseif result == InstallContext.ABORT then
+            ret_val = Result.failure(results[3])
         elseif coroutine.status(thread) == "suspended" then
             -- yield to parent coroutine
-            step(coroutine.yield(result))
+            step(coroutine.yield(result, unpack(results, 3)))
         else
             ret_val = result
         end
@@ -294,6 +301,31 @@ function InstallContext:execute(fn)
 end
 
 ---@async
+---@param system_pkg SystemPackage
+function InstallContext:require(system_pkg)
+    local result = Result.try(function(try)
+        if try(system_pkg:needs_install()) then
+            self.stdio_sink:stdout("Installing dependency " .. system_pkg.name .. ".\n")
+            self.runner.suspend()
+            try(system_pkg:install():on_failure(function()
+                if self.opts.force then
+                    self.runner.resume()
+                end
+            end))
+            self.runner.resume()
+        end
+    end)
+    if result:is_failure() then
+        if not self.opts.force then
+            self.stdio_sink:stderr "Run with :MasonInstall --force to attempt installation anyway.\n"
+            coroutine.yield(InstallContext.ABORT, result:err_or_nil())
+        else
+            self.stdio_sink:stderr(result:err_or_nil() .. "\n")
+        end
+    end
+end
+
+---@async
 function InstallContext:build_receipt()
     log.fmt_debug("Building receipt for %s", self.package)
     return Result.pcall(function()
diff --git a/lua/mason-core/installer/managers/cargo.lua b/lua/mason-core/installer/managers/cargo.lua
index 22ec9ed6..a6116f9a 100644
--- a/lua/mason-core/installer/managers/cargo.lua
+++ b/lua/mason-core/installer/managers/cargo.lua
@@ -1,4 +1,5 @@
 local Result = require "mason-core.result"
+local SystemPackage = require "mason-core.system-package"
 local _ = require "mason-core.functional"
 local installer = require "mason-core.installer"
 local log = require "mason-core.log"
@@ -15,6 +16,7 @@ function M.install(crate, version, opts)
     opts = opts or {}
     log.fmt_debug("cargo: install %s %s %s", crate, version, opts)
     local ctx = installer.context()
+    ctx:require(SystemPackage.sfw)
     ctx.stdio_sink:stdout(("Installing crate %s@%s…\n"):format(crate, version))
     return ctx.spawn.cargo {
         "install",
@@ -29,6 +31,7 @@ function M.install(crate, version, opts)
         opts.features and { "--features", opts.features } or vim.NIL,
         opts.locked and "--locked" or vim.NIL,
         crate,
+        firewall = true,
     }
 end
 
diff --git a/lua/mason-core/installer/managers/npm.lua b/lua/mason-core/installer/managers/npm.lua
index 93af3a85..8a8d1582 100644
--- a/lua/mason-core/installer/managers/npm.lua
+++ b/lua/mason-core/installer/managers/npm.lua
@@ -1,4 +1,5 @@
 local Result = require "mason-core.result"
+local SystemPackage = require "mason-core.system-package"
 local _ = require "mason-core.functional"
 local installer = require "mason-core.installer"
 local log = require "mason-core.log"
@@ -62,12 +63,14 @@ function M.install(pkg, version, opts)
     opts = opts or {}
     log.fmt_debug("npm: install %s %s %s", pkg, version, opts)
     local ctx = installer.context()
+    ctx:require(SystemPackage.sfw)
     ctx.stdio_sink:stdout(("Installing npm package %s@%s…\n"):format(pkg, version))
     return ctx.spawn.npm {
         "install",
         ("%s@%s"):format(pkg, version),
         opts.extra_packages or vim.NIL,
         opts.install_extra_args or vim.NIL,
+        firewall = true,
     }
 end
 
diff --git a/lua/mason-core/installer/managers/pypi.lua b/lua/mason-core/installer/managers/pypi.lua
index 9b39b0d6..a8efd0dd 100644
--- a/lua/mason-core/installer/managers/pypi.lua
+++ b/lua/mason-core/installer/managers/pypi.lua
@@ -1,5 +1,6 @@
 local Optional = require "mason-core.optional"
 local Result = require "mason-core.result"
+local SystemPackage = require "mason-core.system-package"
 local _ = require "mason-core.functional"
 local a = require "mason-core.async"
 local installer = require "mason-core.installer"
@@ -173,6 +174,8 @@ end
 ---@param pkgs string[]
 ---@param extra_args? string[]
 local function pip_install(pkgs, extra_args)
+    local ctx = installer.context()
+    ctx:require(SystemPackage.sfw)
     return venv_python {
         "-m",
         "pip",
@@ -182,6 +185,7 @@ local function pip_install(pkgs, extra_args)
         "--ignore-installed",
         extra_args or vim.NIL,
         pkgs,
+        firewall = true,
     }
 end
 
diff --git a/lua/mason-core/result.lua b/lua/mason-core/result.lua
index e98a11b3..27c46862 100644
--- a/lua/mason-core/result.lua
+++ b/lua/mason-core/result.lua
@@ -189,7 +189,8 @@ function Result.try(fn)
     local thread = coroutine.create(fn)
     local step
     step = function(...)
-        local ok, result = coroutine.resume(thread, ...)
+        local results = { coroutine.resume(thread, ...) }
+        local ok, result = results[1], results[2]
         if not ok then
             return Result.failure(result)
         end
@@ -207,7 +208,7 @@ function Result.try(fn)
             end
         else
             -- yield to parent coroutine
-            return step(coroutine.yield(result))
+            return step(coroutine.yield(result, unpack(results, 3)))
         end
     end
     return step(coroutine.yield)
diff --git a/lua/mason-core/spawn.lua b/lua/mason-core/spawn.lua
index 3c9c645d..78be4fd7 100644
--- a/lua/mason-core/spawn.lua
+++ b/lua/mason-core/spawn.lua
@@ -4,6 +4,7 @@ local a = require "mason-core.async"
 local log = require "mason-core.log"
 local platform = require "mason-core.platform"
 local process = require "mason-core.process"
+local settings = require "mason.settings"
 
 local is_not_nil = _.complement(_.equals(vim.NIL))
 
@@ -64,6 +65,7 @@ local get_path_from_env_list =
 ---@field stdio_sink StdioSink? If provided, will be used to write to stdout and stderr.
 ---@field cwd string?
 ---@field on_spawn (fun(handle: luv_handle, stdio: luv_pipe[], pid: integer))? Will be called when the process successfully spawns.
+---@field firewall boolean?
 
 setmetatable(spawn, {
     ---@param canonical_cmd string
@@ -102,6 +104,20 @@ setmetatable(spawn, {
                 end
             end
 
+            if args.firewall and settings.current.firewall.enabled then
+                a.scheduler()
+                table.insert(spawn_args.args, 1, cmd)
+                local expanded_cmd = exepath(
+                    "sfw",
+                    settings.current.firewall.auto_managed and vim.fn.expand "$MASON/opt/mason/system/bin" or nil
+                )
+                if expanded_cmd == "" then
+                    return Failure({ stderr = "Failed to find sfw (Socket Firewall) in PATH." }, "sfw")
+                else
+                    cmd = expanded_cmd
+                end
+            end
+
             local _, exit_code, signal = a.wait(function(resolve)
                 local handle, stdio, pid = process.spawn(cmd, spawn_args, resolve)
                 if args.on_spawn and handle and stdio and pid then
diff --git a/lua/mason-core/system-package.lua b/lua/mason-core/system-package.lua
new file mode 100644
index 00000000..1eb14e3f
--- /dev/null
+++ b/lua/mason-core/system-package.lua
@@ -0,0 +1,95 @@
+local Result = require "mason-core.result"
+local _ = require "mason-core.functional"
+local a = require "mason-core.async"
+local registry = require "mason-registry"
+local settings = require "mason.settings"
+local OneShotChannel = require("mason-core.async.control").OneShotChannel
+
+---@class SystemPackage
+---@field name string
+---@field condition? fun(): bool
+local SystemPackage = {}
+SystemPackage.__index = SystemPackage
+
+---@type table<string, OneShotChannel>
+SystemPackage.channels = {}
+
+function SystemPackage:new(system_pkg_name)
+    ---@type SystemPackage
+    local instance = {}
+    setmetatable(instance, self)
+    instance.name = system_pkg_name
+    return instance
+end
+
+function SystemPackage:conditional(fn)
+    self.condition = fn
+    return self
+end
+
+---@async
+function SystemPackage:get_package()
+    pcall(a.wait, registry.refresh_system)
+    if not registry.has_system_package(self.name) then
+        -- Force update to the very latest registry version
+        pcall(a.wait, registry.update)
+    end
+    if not registry.has_system_package(self.name) then
+        return Result.failure("Unable to find system package " .. self.name)
+    end
+    return Result.pcall(registry.get_system_package, self.name)
+end
+
+---@async
+---@return Result<boolean>
+function SystemPackage:needs_install()
+    return Result.try(function(try)
+        if self.condition and not self.condition() then
+            return false
+        end
+        local pkg = try(self:get_package())
+        if not pkg:is_installed() or pkg:is_installing() then
+            return true
+        end
+        if pkg:get_installed_version() ~= pkg:get_latest_version() then
+            return true
+        end
+        return false
+    end)
+end
+
+---@async
+function SystemPackage:await_channel()
+    assert(SystemPackage.channels[self.name], "Tried to await non-existing channel.")
+    local success, result = SystemPackage.channels[self.name]:receive()
+    if not success then
+        return Result.failure("Failed to install system package " .. self.name .. ". Error: " .. result)
+    end
+    return Result.success()
+end
+
+---@async
+---@return Result
+function SystemPackage:install()
+    return Result.try(function(try)
+        local pkg = try(self:get_package())
+        if not pkg:is_installing() then
+            local channel = OneShotChannel:new()
+            SystemPackage.channels[self.name] = channel
+            pkg:install({}, function(success, result)
+                channel:send(success, result)
+            end)
+        end
+        return self:await_channel()
+    end)
+end
+
+function SystemPackage:__tostring()
+    return ("SystemPackage(name=%s)"):format(self.name)
+end
+
+SystemPackage.sfw = SystemPackage:new("sfw@latest"):conditional(function()
+    return settings.current.firewall.enabled and settings.current.firewall.auto_managed
+end)
+
+return SystemPackage
diff --git a/lua/mason-registry/init.lua b/lua/mason-registry/init.lua
index 1e59175c..a5d0ed36 100644
--- a/lua/mason-registry/init.lua
+++ b/lua/mason-registry/init.lua
@@ -1,5 +1,6 @@
 local EventEmitter = require "mason-core.EventEmitter"
 local InstallLocation = require "mason-core.installer.InstallLocation"
+local _ = require "mason-core.functional"
 local log = require "mason-core.log"
 local path = require "mason-core.path"
 local uv = vim.loop
@@ -56,6 +57,11 @@ function Registry.get_system_package(pkg_name)
     error(("Cannot find system package %q."):format(pkg_name))
 end
 
+function Registry.has_system_package(pkg_name)
+    local ok = pcall(Registry.get_system_package, pkg_name)
+    return ok
+end
+
 function Registry.get_installed_package_names()
     local fs = require "mason-core.fs"
     if not fs.sync.dir_exists(InstallLocation.global():package()) then
@@ -169,7 +175,7 @@ local REGISTRY_STORE_TTL = 86400 -- 24 hrs
 
 ---@param sources LazySourceCollection
 ---@param callback? RegistryUpdateCallback
-local function refresh(sources, callback)
+local refresh = _.scheduler_wrap(function(sources, callback)
     local a = require "mason-core.async"
 
     local state = sources:get_install_state()
@@ -196,7 +202,7 @@ local function refresh(sources, callback)
     else
         a.run(update, callback, sources)
     end
-end
+end)
 
 ---@param callback? RegistryUpdateCallback
 function Registry.refresh(callback)
diff --git a/lua/mason-registry/sources/init.lua b/lua/mason-registry/sources/init.lua
index b69fa8b1..7d688cd4 100644
--- a/lua/mason-registry/sources/init.lua
+++ b/lua/mason-registry/sources/init.lua
@@ -1,4 +1,5 @@
 local _ = require "mason-core.functional"
+local _ = require "mason-core.functional"
 local log = require "mason-core.log"
 
 ---@class RegistrySource
diff --git a/lua/mason-test/helpers.lua b/lua/mason-test/helpers.lua
index 88354046..a056406e 100644
--- a/lua/mason-test/helpers.lua
+++ b/lua/mason-test/helpers.lua
@@ -12,7 +12,10 @@ local M = {}
 function M.create_context(opts)
     local pkg = registry.get_package(opts and opts.package or "dummy")
     local handle = InstallHandle:new(pkg, InstallLocation.global())
-    local context = InstallContext:new(handle, opts and opts.install_opts or {})
+    local context = InstallContext:new(handle, opts and opts.install_opts or {}, {
+        suspend = function() end,
+        resume = function() end,
+    })
     context.spawn = setmetatable({}, {
         __index = function(s, cmd)
             s[cmd] = spy.new(function()
diff --git a/lua/mason/settings.lua b/lua/mason/settings.lua
index 1788c035..40289e8d 100644
--- a/lua/mason/settings.lua
+++ b/lua/mason/settings.lua
@@ -42,6 +42,18 @@ local DEFAULT_SETTINGS = {
         "github:mason-org/mason-system-registry",
     },
 
+    firewall = {
+        ---@since 2.3.0
+        -- Whether to enable the socket.dev firewall (sfw) for supported package sources.
+        -- For more information, refer to https://socket.dev.
+        enabled = false,
+
+        ---@since 2.3.0
+        -- Whether mason.nvim should automatically install and update the Socket Firewall client.
+        -- If false, the sfw binary must exist in PATH if the firewall is enabled.
+        auto_managed = true,
+    },
+
     ---@since 1.0.0
     -- The provider implementations to use for resolving supplementary package metadata (e.g., all available versions).
     -- Accepts multiple entries, where later entries will be used as fallback should prior providers fail.
diff --git a/lua/mason/ui/components/main/package_list.lua b/lua/mason/ui/components/main/package_list.lua
index 11f0d174..67805088 100644
--- a/lua/mason/ui/components/main/package_list.lua
+++ b/lua/mason/ui/components/main/package_list.lua
@@ -271,6 +271,7 @@ local function InstallingPackageComponent(pkg, state)
                 pkg_state.has_failed and p.error(settings.current.ui.icons.package_uninstalled)
                     or p.highlight(settings.current.ui.icons.package_pending),
                 p.none(" " .. pkg.name),
+                pkg_state.firewall_active and p.highlight " (firewall active)" or p.none "",
                 current_state,
                 pkg_state.latest_spawn and p.Comment((" $ %s"):format(pkg_state.latest_spawn)) or p.none "",
             },
diff --git a/lua/mason/ui/instance.lua b/lua/mason/ui/instance.lua
index 476bdf8c..fc035318 100644
--- a/lua/mason/ui/instance.lua
+++ b/lua/mason/ui/instance.lua
@@ -46,6 +46,7 @@ end
 ---@field is_log_expanded boolean
 ---@field has_failed boolean
 ---@field latest_spawn string?
+---@field firewall_active boolean?
 ---@field linked_executables table<string, string>?
 ---@field installed_purl string?
 ---@field lsp_settings_schema table?
@@ -239,8 +240,13 @@ local function setup_handle(handle)
 
     local function handle_spawnhandle_change()
         mutate_state(function(state)
+            local spawn_handle = handle:peek_spawn_handle()
             state.packages.states[handle.package.name].latest_spawn =
-                handle:peek_spawn_handle():map(tostring):map(_.gsub("\n", "\\n ")):or_else(nil)
+                spawn_handle:map(tostring):map(_.gsub("\n", "\\n ")):or_else(nil)
+            if settings.current.firewall.enabled then
+                state.packages.states[handle.package.name].firewall_active =
+                    spawn_handle:map(_.prop "firewall"):or_else(false)
+            end
         end)
     end