Add support for nix store remote copying

.github/nix-server/Dockerfile

@@ -0,0 +1,24 @@
+FROM ubuntu:23.04
+
+RUN apt-get update -qq && \
+    apt-get install openssh-server curl xz-utils sudo locales ca-certificates -y && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN mkdir -m 0755 /nix && \
+    groupadd -r nixbld && \
+    chown root /nix && \
+    for n in $(seq 1 10); do useradd -c "Nix build user $n" -d /var/empty -g nixbld -G nixbld -M -N -r -s "$(command -v nologin)" "nixbld$n"; done
+
+RUN curl -L https://nixos.org/nix/install | bash
+
+COPY .github/nix-server/keys .
+
+RUN cat ci.pub > $HOME/.ssh/authorized_keys
+
+RUN echo 'PermitRootLogin yes' >> /etc/ssh/sshd_config && \
+    echo ". /root/.nix-profile/etc/profile.d/nix.sh" >> $HOME/.bashrc && \
+    ln -sf /root/.nix-profile/bin/nix-store /usr/bin/ && \
+    ln -sf /root/.nix-profile/bin/nix-daemon /usr/bin/
+
+CMD ["/usr/sbin/sshd", "-D"]

.github/nix-server/config

@@ -0,0 +1,9 @@
+Host nix-server
+  Hostname localhost
+  Port 2222
+  User root
+  IdentityFile /home/runner/.ssh/ci
+
+Host *
+    StrictHostKeyChecking no
+    UserKnownHostsFile=/dev/null

.github/workflows/workflow.yaml

@@ -17,7 +17,7 @@ jobs:
         with:
           access_token: ${{ github.token }}
   test-nixpkgs:
-    name: Build & Test - Nixpkgs - ${{ matrix.bzlmodEnabled && 'bzlmod' || 'workspace' }} - ${{ matrix.os }}
+    name: Build & Test - Nixpkgs - ${{ matrix.bzlmodEnabled && 'bzlmod' || 'workspace' }} ${{ matrix.withNixRemote && '- NixRemote' || '' }} - ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
@@ -27,6 +27,9 @@ jobs:
         bzlmodEnabled:
           - true
           - false
+        withNixRemote:
+          - true
+          - false
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/[email protected]
@@ -49,13 +52,34 @@ jobs:
           # no-op flag to avoid "ERROR: Config value 'ci' is not defined in any .rc file"
           common:ci --announce_rc=false
           EOF
+      - name: Start remote Nix server
+        if: matrix.withNixRemote && matrix.os == 'ubuntu-latest'
+        run: |
+          # Generate temporary SSH keys.
+          mkdir -p $HOME/.ssh
+          mkdir -p .github/nix-server/keys
+          ssh-keygen -t ed25519 -f .github/nix-server/keys/ci -C ci-nix-server -q -N ""
+
+          docker build -t nix-server -f .github/nix-server/Dockerfile .
+          docker run -d -p 2222:22 nix-server
+
+          cp .github/nix-server/keys/* $HOME/.ssh/
+
+          sudo cp .github/nix-server/config /etc/ssh/ssh_config
       - name: Build & test
         env:
           BZLMOD_ENABLED: ${{ matrix.bzlmodEnabled }}
+          NIX_REMOTE_ENABLED: matrix.withNixRemote && matrix.os == 'ubuntu-latest'
         run: |
+          if [ "$NIX_REMOTE_ENABLED" = "true" ]; then
+            echo "Setting BAZEL_NIX_REMOTE env variable"
+            export BAZEL_NIX_REMOTE=nix-server
+          fi
+
           nix-shell --pure \
             --keep GITHUB_REPOSITORY \
             --keep BZLMOD_ENABLED \
+            --keep BAZEL_NIX_REMOTE \
             --run 'bash .github/build-and-test'
   test-examples:
     name: Build & Test - Examples

.gitignore

@@ -1,3 +1,5 @@
 /**/bazel-*
 
 /**/node_modules
+
+.github/nix-server/keys

core/nixpkgs.bzl

@@ -396,22 +396,65 @@ def _nixpkgs_build_file_content(repository_ctx):
     else:
         return None
 
-def _nixpkgs_build_and_symlink(repository_ctx, nix_build, build_file_content):
+def _nixpkgs_build_and_symlink(repository_ctx, nix_build_cmd, expr_args, build_file_content):
     # Large enough integer that Bazel can still parse. We don't have
     # access to MAX_INT and 0 is not a valid timeout so this is as good
     # as we can do. The value shouldn't be too large to avoid errors on
     # macOS, see https://github.com/tweag/rules_nixpkgs/issues/92.
     timeout = 8640000
     repository_ctx.report_progress("Building Nix derivation")
+
+    nix_path = executable_path(
+        repository_ctx,
+        "nix",
+        extra_msg = "See: https://nixos.org/nix/",
+    )
+
+    nix_host = repository_ctx.os.environ.get('BAZEL_NIX_REMOTE', '')
+    if nix_host:
+        nix_store = "ssh-ng://{host}?max-connections=1".format(host = nix_host)
+        repository_ctx.report_progress("Remote-building Nix derivation")
+        exec_result = execute_or_fail(
+             repository_ctx,
+             nix_build_cmd + ["--store", nix_store, "--eval-store", "auto"] + expr_args,
+             failure_message = "Cannot build Nix attribute '{}'.".format(
+                 repository_ctx.attr.name,
+             ),
+             quiet = repository_ctx.attr.quiet,
+             timeout = timeout,
+        )
+        output_path = exec_result.stdout.splitlines()[-1]
+
+        ssh_path = repository_ctx.which("ssh")
+        repository_ctx.report_progress("Creating remote store root")
+        exec_result = execute_or_fail(
+            repository_ctx,
+            [ssh_path] + [nix_host, "nix-store --add-root ~/rules_nixpkgs_gcroots/{root} -r {path}".format(root = output_path.split('/')[-1], path = output_path) ],
+            failure_message = "Cannot create remote store root for Nix attribute '{}'.".format(
+                repository_ctx.attr.name,
+            ),
+            quiet = repository_ctx.attr.quiet,
+            timeout = 10000,
+        )
+
+        repository_ctx.report_progress("Downloading Nix derivation")
+        exec_result = repository_ctx.execute(
+            [nix_path, "copy", "--from", nix_store, output_path],
+            quiet = repository_ctx.attr.quiet,
+            timeout = 10000,
+        )
+
     exec_result = execute_or_fail(
         repository_ctx,
-        nix_build,
+        nix_build_cmd + expr_args,
         failure_message = "Cannot build Nix derivation for package '@{}'.".format(repository_ctx.name),
         quiet = repository_ctx.attr.quiet,
         timeout = timeout,
     )
     output_path = exec_result.stdout.splitlines()[-1]
 
+    repository_ctx.report_progress("Creating local folders")
+
     # ensure that the output is a directory
     test_path = repository_ctx.which("test")
     execute_or_fail(
@@ -542,9 +585,8 @@ def _nixpkgs_package_impl(repository_ctx):
         "nix-build",
         extra_msg = "See: https://nixos.org/nix/",
     )
-    nix_build = [nix_build_path] + expr_args
 
-    _nixpkgs_build_and_symlink(repository_ctx, nix_build, build_file_content)
+    _nixpkgs_build_and_symlink(repository_ctx, [nix_build_path], expr_args, build_file_content)
 
 _nixpkgs_package = repository_rule(
     implementation = _nixpkgs_package_impl,
@@ -741,9 +783,8 @@ def _nixpkgs_flake_package_impl(repository_ctx):
         "nix",
         extra_msg = "See: https://nixos.org/nix/",
     )
-    nix_build = [nix_path, "build"] + expr_args
 
-    _nixpkgs_build_and_symlink(repository_ctx, nix_build, build_file_content)
+    _nixpkgs_build_and_symlink(repository_ctx, [nix_path, "build"], expr_args, build_file_content)
 
 _nixpkgs_flake_package = repository_rule(
     implementation = _nixpkgs_flake_package_impl,

flake.nix

@@ -15,7 +15,7 @@
       {
         devShells.default = with pkgs; mkShell {
           name = "rules_nixpkgs_shell";
-          packages = [ bazel_6 bazel-buildtools cacert gcc nix git ];
+          packages = [ bazel_6 bazel-buildtools cacert gcc nix git openssh ];
         };
       });
 }