Nix Flake for Aya eBPF project

flake.nix

{
  description = "Krait eBPF mesh VPN";

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
    flake-utils.url = "github:numtide/flake-utils";
  };

  outputs = { self, nixpkgs, flake-utils, ... }:
    flake-utils.lib.eachDefaultSystem (system:
      let
        pkgs = nixpkgs.legacyPackages.${system};
        krait-ebpf = pkgs.rustPlatform.buildRustPackage {
          pname = "krait-ebpf";
          version = "0.1.0";
          src = ./.;

          cargoHash = "sha256-l77zkJqvURwdpk5W4Cy6gB2KcGqc0L/pFut+roCXUVs=";

          buildPhase = ''
            runHook preBuild
            cd krait-ebpf
            cargo build --release --target bpfel-unknown-none --offline
            runHook postBuild
          '';

          installPhase = ''
            runHook preInstall
            mkdir -p $out/lib
            cp ../target/bpfel-unknown-none/release/krait $out/lib/
            runHook postInstall
          '';

          nativeBuildInputs = [ pkgs.bpf-linker ];
          doCheck = false;

          RUSTC_BOOTSTRAP = 1;
          RUSTFLAGS = "-Cpanic=abort";
        };
      in
      {
        packages.krait-ebpf = krait-ebpf;
        
        # pkgsStatic to build as static binaries using musl libc.
        packages.krait-agent = pkgs.pkgsStatic.rustPlatform.buildRustPackage {
          pname = "krait-agent";
          version = "0.1.0";
          src = ./.;
          buildAndTestSubDir = "krait-agent";
          cargoHash = "sha256-l77zkJqvURwdpk5W4Cy6gB2KcGqc0L/pFut+roCXUVs=";
          cargoBuildFlags = [ "-p" "krait-agent" ];
          doCheck = false;
          AYA_BUILD_SKIP = 1; # skip Aya build.rs logic
          KRAIT_EBPF_PATH = "${krait-ebpf}/lib/krait";
        };

        packages.krait-tests = pkgs.pkgsStatic.rustPlatform.buildRustPackage {
          pname = "krait-tests";
          version = "0.1.0";
          src = ./.;
          buildAndTestSubDir = "krait-tests";
          cargoHash = "sha256-l77zkJqvURwdpk5W4Cy6gB2KcGqc0L/pFut+roCXUVs=";
          cargoBuildFlags = [ "-p" "krait-tests" ];
          doCheck = false;
          AYA_BUILD_SKIP = 1;
          KRAIT_EBPF_PATH = "${krait-ebpf}/lib/krait";
        };

        # NixOS integration tests
        checks = {
          wireguardTest = pkgs.testers.nixosTest (import ./tests/nixos/wireguard.nix {
            inherit pkgs;
            krait-agent = self.packages.${system}.krait-agent;
            krait-tests = self.packages.${system}.krait-tests;
          });
        };

        devShells.default = pkgs.mkShell {
          buildInputs = [
            pkgs.bpftools
            pkgs.bpf-linker
            pkgs.cargo
            pkgs.clippy
            pkgs.rust-analyzer
            pkgs.rustup
            pkgs.wireguard-tools
          ];
        };
      });
}

wireguard.nix

# Sample NixOS integration for Wireguard + eBPF

{ pkgs, krait-agent, krait-tests, ... }:

let
  # Krait agent systemd service definition
  kraitAgentService = {
    description = "Krait eBPF Agent";
    after = [ "network.target" "wireguard-wg0.service" ];
    wants = [ "wireguard-wg0.service" ];
    wantedBy = [ "multi-user.target" ];
    path = [ pkgs.iproute2 ];

    serviceConfig = {
      Type = "simple";
      ExecStart = "${krait-agent}/bin/krait --iface wg0";
      Restart = "on-failure";
      RestartSec = "5s";

      # eBPF requires some capabilities
      CapabilityBoundingSet = [ "CAP_SYS_ADMIN" "CAP_NET_ADMIN" "CAP_SYS_RESOURCE" ];
      AmbientCapabilities = [ "CAP_SYS_ADMIN" "CAP_NET_ADMIN" "CAP_SYS_RESOURCE" ];
    };

    environment = {
      RUST_LOG = "info";
    };
  };
in
{
  name = "wireguard-connectivity-test";

  nodes = {
    # --- Server VM ---
    server = { pkgs, ... }: {
      networking.firewall.allowedUDPPorts = [ 51820 ];
      networking.firewall.allowedTCPPorts = [ 5201 ]; # iperf3 port
      networking.wireguard.interfaces.wg0 = {
        ips = [ "10.10.0.1/24" ];
        listenPort = 51820;
        privateKey = builtins.readFile ./server.priv;
        peers = [{
          publicKey = builtins.readFile ./client.pub;
          allowedIPs = [ "10.10.0.2/32" ];
        }];
      };
      environment.systemPackages = with pkgs; [
        bpftools
        iperf3
        iproute2
        tcpdump
      ];

      # Krait agent systemd service
      systemd.services.krait-agent = kraitAgentService;
    };

    # --- Client VM ---
    client = { pkgs, ... }: {
      networking.wireguard.interfaces.wg0 = {
        ips = [ "10.10.0.2/24" ];
        privateKey = builtins.readFile ./client.priv;
        peers = [{
          publicKey = builtins.readFile ./server.pub;
          # In NixOS tests, nodes can resolve each other by their hostname.
          endpoint = "server:51820";
          allowedIPs = [ "10.10.0.0/24" ];
          persistentKeepalive = 25;
        }];
      };
      environment.systemPackages = with pkgs; [
        bpftools
        iperf3
        iproute2
        tcpdump
      ];

      # Krait agent systemd service
      systemd.services.krait-agent = kraitAgentService;
    };
  };

  # The testScript is written in Python
  testScript = ''
    start_all()

    with subtest("Wait for Wireguard services to start"):
        server.wait_for_unit("wireguard-wg0.service")
        client.wait_for_unit("wireguard-wg0.service")
        # Give additional time for interfaces to come up
        server.sleep(2)
        client.sleep(2)

    with subtest("Wait for Krait agent to start"):
        server.wait_for_unit("krait-agent.service")
        client.wait_for_unit("krait-agent.service")
        # Verify agents are running
        server.succeed("systemctl is-active krait-agent.service")
        client.succeed("systemctl is-active krait-agent.service")

    with subtest("Verify interface configuration"):
        # Verify interfaces have correct IP addresses
        server.succeed("ip addr show wg0 | grep '10.10.0.1/24'")
        client.succeed("ip addr show wg0 | grep '10.10.0.2/24'")

        # Show interface status for debugging
        print("Server wg0:", server.succeed("ip addr show wg0"))
        print("Client wg0:", client.succeed("ip addr show wg0"))

    with subtest("Establish Wireguard tunnel"):
        # Trigger handshake by having client send a packet first
        # This establishes the tunnel since client has the endpoint configured
        print("Initiating Wireguard handshake...")

        # Try ping until it succeeds to establish the tunnel
        client.wait_until_succeeds("ping -c 1 10.10.0.1", timeout=30)

        # Verify handshake occurred on both sides
        server.succeed("wg show | grep 'latest handshake'")
        client.succeed("wg show | grep 'latest handshake'")

    with subtest("Test tunnel connectivity"):
        # Test bidirectional connectivity
        client.succeed("ping -c 3 10.10.0.1")
        server.succeed("ping -c 3 10.10.0.2")

    with subtest("Test tunnel throughput"):
        server.succeed("iperf3 -s -D")
        server.sleep(1)
        client.succeed("iperf3 -c 10.10.0.1 -t 5")

    with subtest("Run Krait integration test"):
        # Run integration test on client side - it will send traffic to server
        print("Running Krait integration test suite on client...")
        client.succeed("RUST_LOG=info ${krait-tests}/bin/krait-tests --interface wg0 --target-ip 10.10.0.1 --num-bytes 30M --output-file /tmp/krait-test-results.json --test-threads 1")

        # Show test results
        results = client.succeed("cat /tmp/krait-test-results.json")
        print("Krait integration test results:")
        print(results)

    print("Wireguard tunnel test completed successfully!")
  '';
}