arbi/flake.nix

{
  inputs = {
    flake-parts.url = "github:hercules-ci/flake-parts";
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    treefmt-nix = {
      url = "github:numtide/treefmt-nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    git-hooks = {
      url = "github:cachix/git-hooks.nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    agenix-shell.url = "github:aciceri/agenix-shell";
    flake-root.url = "github:srid/flake-root";
    nix-github-actions = {
      url = "github:nix-community/nix-github-actions";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    forge-std = {
      flake = false;
      url = "github:foundry-rs/forge-std/v1.9.6";
    };
  };

  outputs = inputs:
    inputs.flake-parts.lib.mkFlake { inherit inputs; } (flake@{ config, lib, moduleWithSystem, withSystem, ... }: {
      systems = [ "x86_64-linux" "aarch64-linux" ];

      imports = with inputs; [
        git-hooks.flakeModule
        treefmt-nix.flakeModule
        flake-root.flakeModule
        agenix-shell.flakeModules.agenix-shell
      ];

      agenix-shell = {
        secrets = {
          ALCHEMY_KEY.file = ./secrets/alchemy_key.age;
          WALLET_PRIVATE_KEY.file = ./secrets/wallet_private_key.age;
        };
      };

      perSystem = { pkgs, config, ... }: {
        treefmt.config = {
          flakeFormatter = true;
          flakeCheck = true;
          programs = {
            nixpkgs-fmt.enable = true;
            rustfmt.enable = true;
            actionlint.enable = true;
          };
        };

        pre-commit = {
          check.enable = false;
          settings.hooks = {
            treefmt = {
              enable = true;
              package = config.treefmt.build.wrapper;
            };
          };
        };

        devShells.default = pkgs.mkShell {
          packages = with pkgs; [ cargo rustc rust-analyzer clippy foundry typst tinymist age ragenix ];
          inputsFrom = [ config.flake-root.devShell ];
          shellHook = ''
            source ${lib.getExe config.agenix-shell.installationScript}

            # forge will use this directory to download the solc compilers
            mkdir -p $HOME/.svm

            # forge needs forge-std to work
            mkdir -p $FLAKE_ROOT/onchain/lib/
            ln -sf ${inputs.forge-std.outPath} $FLAKE_ROOT/onchain/lib/forge-std

            if [ ! -f "$FLAKE_ROOT/offchain/config.kdl" ]; then \
              cp ${config.packages.arbi_sample_config_kdl} $FLAKE_ROOT/offchain/config.kdl
            fi
            export ARBI_CONFIG="$FLAKE_ROOT/offchain/config.kdl"

            ${config.pre-commit.installationScript}
          '';
          env = {
            OPENSSL_DIR = pkgs.openssl.dev;
            OPENSSL_NO_VENDOR = true;
            OPENSSL_LIB_DIR = "${pkgs.openssl.out}/lib";
            OPENSSL_INCLUDE_DIR = "${pkgs.openssl.dev}/include";
            PKG_CONFIG_PATH = "${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH";
            RUST_BACKTRACE = true;
            ARBI_LOG_LEVEL = "debug";
          };
        };

        packages = {
          default = config.packages.arbi;

          arbi = pkgs.rustPlatform.buildRustPackage {
            pname = "arbi";
            version = "0.1.0";
            cargoLock.lockFile = ./offchain/Cargo.lock;
            src = ./offchain;
            env = {
              OPENSSL_DIR = pkgs.openssl.dev;
              OPENSSL_NO_VENDOR = true;
              OPENSSL_LIB_DIR = "${pkgs.openssl.out}/lib";
              OPENSSL_INCLUDE_DIR = "${pkgs.openssl.dev}/include";
              PKG_CONFIG_PATH = "${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH";
            };
            meta.mainProgram = "arbi";
          };

          arbi_sample_config_kdl = pkgs.writeText "arbi-sample-config.kdl" ''
            endpoint "wss://eth-mainnet.g.alchemy.com/v2/<REDACTED>"
            pairs_file "pairs.json"
            concurrency 5
          '';

          run-forge-tests = pkgs.writeShellScriptBin "run-forge-tests" ''
            pushd "$FLAKE_ROOT/onchain"
            forge test \
              --fork-url "wss://mainnet.infura.io/ws/v3/$ALCHEMY_KEY" \
              --via-ir \
              -vvv
            popd
          '';

          run-vm = pkgs.writeShellScriptBin "run-vm" (lib.getExe flake.config.flake.nixosConfigurations.vm.config.system.build.vm);
        } // lib.genAttrs [ "polygon-mainnet" ] (network: pkgs.writeShellScriptBin "deploy-${network}" ''
          pushd "$FLAKE_ROOT/onchain"
          forge create \
            --rpc-url "wss://${network}.infura.io/ws/v3/$ALCHEMY_KEY" \
            --private-key "$WALLET_PRIVATE_KEY" \
            --via-ir \
            --broadcast \
            src/ArbitrageManager.sol:ArbitrageManager
          popd
        '');


        checks = {
          inherit (config.packages) arbi;
        };
      };

      flake = {
        githubActions = inputs.nix-github-actions.lib.mkGithubMatrix {
          checks = lib.getAttrs [ "x86_64-linux" ] config.flake.checks;
        };

        nixosConfigurations.vm = withSystem "x86_64-linux" (ctx: inputs.nixpkgs.lib.nixosSystem {
          system = "x86_64-linux";
          modules = [
            ({ pkgs, modulesPath, ... }: {
              imports = [
                "${modulesPath}/virtualisation/qemu-vm.nix"
                config.flake.nixosModules.arbi
              ];

              services.getty.autologinUser = "root";
              services.openssh.settings.PasswordAuthentication = lib.mkForce true;
              services.openssh.settings.PermitRootLogin = lib.mkForce "yes";
              users.users.root.password = "";

              virtualisation = {
                graphics = false;
                memorySize = 2048;
                diskSize = 10000;
                forwardPorts = [
                  {
                    from = "host";
                    host.port = 2222;
                    guest.port = 22;
                  }
                ];
              };

              system.stateVersion = "25.05";

              services.arbi = {
                enable = true;
                log_level = "debug";
                configFile = pkgs.writeText "arbi-config.kdl" ''
                  endpoint "wss://eth-mainnet.g.alchemy.com/v2/kkDMaLVYpWQA0GsCYNFvAODnAxCCiamv"
                  pairs_file "pairs.json"
                  concurrency 5
                '';
              };
            })
          ];
        });

        nixosModules = {
          arbi = moduleWithSystem ({ config }: nixos@{ lib, utils, ... }:
            let
              cfg = nixos.config.services.arbi;
            in
            {
              options.services.arbi = {
                enable = lib.mkEnableOption "arbi";
                package = lib.mkOption {
                  type = lib.types.package;
                  default = config.packages.arbi;
                };
                log_level = lib.mkOption {
                  type = lib.types.enum [ "debug" "trace" "warn" "error" "info" ];
                  default = "info";
                };
                configFile = lib.mkOption {
                  type = lib.types.path;
                };
                dataDir = lib.mkOption {
                  type = lib.types.path;
                  default = "/var/lib/arbi";
                };
                user = lib.mkOption {
                  type = lib.types.str;
                  default = "arbi";
                };
                group = lib.mkOption {
                  type = lib.types.str;
                  default = "arbi";
                };
              };
              config = lib.mkIf cfg.enable {
                environment.systemPackages = [ cfg.package ];

                users.users.arbi = lib.mkIf (cfg.user == "arbi") {
                  isSystemUser = true;
                  group = cfg.group;
                };

                users.groups.arbi = lib.mkIf (cfg.group == "arbi") { };

                systemd.tmpfiles.settings."10-arbi" = {
                  ${cfg.dataDir}.d = {
                    inherit (cfg) user group;
                    mode = "0755";
                  };
                };

                systemd.services.arbi = {
                  description = "Arbitrage bot";

                  after = [ "network.target" ];
                  wantedBy = [ "multi-user.target" ];

                  environment.ARBI_LOG_LEVEL = cfg.log_level;

                  serviceConfig = {
                    ExecStart = utils.escapeSystemdExecArgs [
                      (lib.getExe cfg.package)
                      "--config"
                      cfg.configFile
                      "run"
                    ];

                    KillSignal = "SIGINT";
                    Restart = "on-failure";
                    RestartSec = "5s";
                    User = cfg.user;
                    Group = cfg.group;
                    WorkingDirectory = cfg.dataDir;
                    UMask = "0022";
                  };
                };
              };
            });
          default = config.flake.nixosModules.arbi;
        };
      };
    });
}