diff --git a/flake.nix b/flake.nix
index a185f0f..419a05e 100644
--- a/flake.nix
+++ b/flake.nix
@@ -25,6 +25,7 @@
             inherit (gomod2nix.legacyPackages.${system}) buildGoApplication;
           };
           devShells.default = callPackage ./shell.nix {
+            inherit source;
             inherit (gomod2nix.legacyPackages.${system}) mkGoEnv gomod2nix;
           };
         })
diff --git a/shell.nix b/shell.nix
index d4f21a8..52939bb 100644
--- a/shell.nix
+++ b/shell.nix
@@ -1,24 +1,26 @@
-{ pkgs ? (
+{
+  source,
+  pkgs ? (
     let
       inherit (builtins) fetchTree fromJSON readFile;
       inherit ((fromJSON (readFile ./flake.lock)).nodes) nixpkgs gomod2nix;
     in
-    import (fetchTree nixpkgs.locked) {
-      overlays = [
-        (import "${fetchTree gomod2nix.locked}/overlay.nix")
-      ];
-    }
-  )
-, mkGoEnv ? pkgs.mkGoEnv
-, gomod2nix ? pkgs.gomod2nix
-}:
-
-let
-  goEnv = mkGoEnv { pwd = ./.; };
+      import (fetchTree nixpkgs.locked) {
+        overlays = [
+          (import "${fetchTree gomod2nix.locked}/overlay.nix")
+        ];
+      }
+  ),
+  mkGoEnv ? pkgs.mkGoEnv,
+  gomod2nix ? pkgs.gomod2nix,
+}: let
+  update_gomod2nix = pkgs.writeShellScriptBin "update_gomod2nix" ''
+    ${gomod2nix}/bin/gomod2nix --dir ${source} --outdir ./
+  '';
 in
-pkgs.mkShell {
-  packages = [
-    goEnv
-    gomod2nix
-  ];
-}
+  pkgs.mkShell {
+    packages = [
+      gomod2nix
+      update_gomod2nix
+    ];
+  }