labmcp/nix/nixpkgs-search-module.nix

{ config, lib, pkgs, ... }:

let
  cfg = config.services.nixpkgs-search;

  # Determine database URL based on configuration
  # For postgres with connectionStringFile, the URL is set at runtime via script
  useConnectionStringFile = cfg.database.type == "postgres" && cfg.database.connectionStringFile != null;

  databaseUrl = if cfg.database.type == "sqlite"
    then "sqlite://${cfg.dataDir}/${cfg.database.name}"
    else if useConnectionStringFile
    then "" # Will be set at runtime from file
    else cfg.database.connectionString;

  # Build HTTP transport flags for a service
  mkHttpFlags = httpCfg: lib.concatStringsSep " " ([
    "--transport http"
    "--http-address '${httpCfg.address}'"
    "--http-endpoint '${httpCfg.endpoint}'"
    "--session-ttl '${httpCfg.sessionTTL}'"
  ] ++ lib.optionals (httpCfg.allowedOrigins != []) (
    map (origin: "--allowed-origins '${origin}'") httpCfg.allowedOrigins
  ) ++ lib.optionals httpCfg.tls.enable [
    "--tls-cert '${httpCfg.tls.certFile}'"
    "--tls-key '${httpCfg.tls.keyFile}'"
  ]);

  # Common HTTP options
  mkHttpOptions = defaultPort: {
    address = lib.mkOption {
      type = lib.types.str;
      default = "127.0.0.1:${toString defaultPort}";
      description = "HTTP listen address for the MCP server.";
    };

    endpoint = lib.mkOption {
      type = lib.types.str;
      default = "/mcp";
      description = "HTTP endpoint path for MCP requests.";
    };

    allowedOrigins = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      example = [ "http://localhost:3000" "https://example.com" ];
      description = ''
        Allowed Origin headers for CORS.
        Empty list means only localhost origins are allowed.
      '';
    };

    sessionTTL = lib.mkOption {
      type = lib.types.str;
      default = "30m";
      description = "Session TTL for HTTP transport (Go duration format).";
    };

    tls = {
      enable = lib.mkEnableOption "TLS for HTTP transport";

      certFile = lib.mkOption {
        type = lib.types.nullOr lib.types.path;
        default = null;
        description = "Path to TLS certificate file.";
      };

      keyFile = lib.mkOption {
        type = lib.types.nullOr lib.types.path;
        default = null;
        description = "Path to TLS private key file.";
      };
    };
  };

  # Service configuration factory
  mkServiceConfig = serviceName: subcommand: httpCfg: {
    description = "Nixpkgs Search ${serviceName} MCP Server";
    wantedBy = [ "multi-user.target" ];
    after = [ "network.target" ]
      ++ lib.optional (cfg.database.type == "postgres") "postgresql.service";

    environment = lib.mkIf (!useConnectionStringFile) {
      NIXPKGS_SEARCH_DATABASE = databaseUrl;
    };

    path = [ cfg.package ];

    script = let
      httpFlags = mkHttpFlags httpCfg;
    in
      if useConnectionStringFile then ''
        # Read database connection string from file
        if [ ! -f "${cfg.database.connectionStringFile}" ]; then
          echo "Error: connectionStringFile not found: ${cfg.database.connectionStringFile}" >&2
          exit 1
        fi
        export NIXPKGS_SEARCH_DATABASE="$(cat "${cfg.database.connectionStringFile}")"

        exec nixpkgs-search ${subcommand} serve ${httpFlags}
      '' else ''
        exec nixpkgs-search ${subcommand} serve ${httpFlags}
      '';

    serviceConfig = {
      Type = "simple";
      User = cfg.user;
      Group = cfg.group;
      Restart = "on-failure";
      RestartSec = "5s";

      # Hardening
      NoNewPrivileges = true;
      ProtectSystem = "strict";
      ProtectHome = true;
      PrivateTmp = true;
      PrivateDevices = true;
      ProtectKernelTunables = true;
      ProtectKernelModules = true;
      ProtectControlGroups = true;
      RestrictNamespaces = true;
      RestrictRealtime = true;
      RestrictSUIDSGID = true;
      MemoryDenyWriteExecute = true;
      LockPersonality = true;

      ReadWritePaths = [ cfg.dataDir ];
      WorkingDirectory = cfg.dataDir;
      StateDirectory = lib.mkIf (cfg.dataDir == "/var/lib/nixpkgs-search") "nixpkgs-search";
    };
  };
in
{
  options.services.nixpkgs-search = {
    enable = lib.mkEnableOption "Nixpkgs Search MCP server(s)";

    package = lib.mkPackageOption pkgs "nixpkgs-search" { };

    user = lib.mkOption {
      type = lib.types.str;
      default = "nixpkgs-search";
      description = "User account under which the service runs.";
    };

    group = lib.mkOption {
      type = lib.types.str;
      default = "nixpkgs-search";
      description = "Group under which the service runs.";
    };

    dataDir = lib.mkOption {
      type = lib.types.path;
      default = "/var/lib/nixpkgs-search";
      description = "Directory to store data files.";
    };

    database = {
      type = lib.mkOption {
        type = lib.types.enum [ "sqlite" "postgres" ];
        default = "sqlite";
        description = "Database backend to use.";
      };

      name = lib.mkOption {
        type = lib.types.str;
        default = "nixpkgs-search.db";
        description = "SQLite database filename (when using sqlite backend).";
      };

      connectionString = lib.mkOption {
        type = lib.types.str;
        default = "";
        description = ''
          PostgreSQL connection string (when using postgres backend).
          Example: "postgres://user:password@localhost/nixpkgs_search?sslmode=disable"

          WARNING: This value will be stored in the Nix store, which is world-readable.
          For production use with sensitive credentials, use connectionStringFile instead.
        '';
      };

      connectionStringFile = lib.mkOption {
        type = lib.types.nullOr lib.types.path;
        default = null;
        description = ''
          Path to a file containing the PostgreSQL connection string.
          The file should contain just the connection string, e.g.:
          postgres://user:password@localhost/nixpkgs_search?sslmode=disable

          This is the recommended way to configure PostgreSQL credentials
          as the file is not stored in the world-readable Nix store.
          The file must be readable by the service user.
        '';
        example = "/run/secrets/nixpkgs-search-db";
      };
    };

    indexOnStart = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      example = [ "nixos-unstable" "nixos-24.11" ];
      description = ''
        List of nixpkgs revisions to index on service start.
        Can be channel names (nixos-unstable) or git hashes.
        Indexing is skipped if the revision is already indexed.
      '';
    };

    indexFlags = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      example = [ "--no-packages" "--no-files" ];
      description = ''
        Additional flags to pass to the index command.
        Useful for skipping packages (--no-packages), options (--no-options),
        or files (--no-files) during indexing.
      '';
    };

    options = {
      enable = lib.mkOption {
        type = lib.types.bool;
        default = true;
        description = "Enable the NixOS options MCP server.";
      };

      http = mkHttpOptions 8082;

      openFirewall = lib.mkOption {
        type = lib.types.bool;
        default = false;
        description = "Whether to open the firewall for the options MCP HTTP server.";
      };
    };

    packages = {
      enable = lib.mkOption {
        type = lib.types.bool;
        default = true;
        description = "Enable the Nix packages MCP server.";
      };

      http = mkHttpOptions 8083;

      openFirewall = lib.mkOption {
        type = lib.types.bool;
        default = false;
        description = "Whether to open the firewall for the packages MCP HTTP server.";
      };
    };
  };

  config = lib.mkIf cfg.enable {
    assertions = [
      {
        assertion = cfg.database.type == "sqlite"
          || cfg.database.connectionString != ""
          || cfg.database.connectionStringFile != null;
        message = "services.nixpkgs-search.database: when using postgres backend, either connectionString or connectionStringFile must be set";
      }
      {
        assertion = cfg.database.connectionString == "" || cfg.database.connectionStringFile == null;
        message = "services.nixpkgs-search.database: connectionString and connectionStringFile are mutually exclusive";
      }
      {
        assertion = !cfg.options.http.tls.enable || (cfg.options.http.tls.certFile != null && cfg.options.http.tls.keyFile != null);
        message = "services.nixpkgs-search.options.http.tls: both certFile and keyFile must be set when TLS is enabled";
      }
      {
        assertion = !cfg.packages.http.tls.enable || (cfg.packages.http.tls.certFile != null && cfg.packages.http.tls.keyFile != null);
        message = "services.nixpkgs-search.packages.http.tls: both certFile and keyFile must be set when TLS is enabled";
      }
      {
        assertion = cfg.options.enable || cfg.packages.enable;
        message = "services.nixpkgs-search: at least one of options.enable or packages.enable must be true";
      }
    ];

    users.users.${cfg.user} = lib.mkIf (cfg.user == "nixpkgs-search") {
      isSystemUser = true;
      group = cfg.group;
      home = cfg.dataDir;
      description = "Nixpkgs Search MCP server user";
    };

    users.groups.${cfg.group} = lib.mkIf (cfg.group == "nixpkgs-search") { };

    systemd.tmpfiles.rules = [
      "d ${cfg.dataDir} 0750 ${cfg.user} ${cfg.group} -"
    ];

    # Indexing service (runs once on startup if indexOnStart is set)
    systemd.services.nixpkgs-search-index = lib.mkIf (cfg.indexOnStart != []) {
      description = "Nixpkgs Search Indexer";
      wantedBy = [ "multi-user.target" ];
      after = [ "network.target" ]
        ++ lib.optional (cfg.database.type == "postgres") "postgresql.service";
      before = lib.optionals cfg.options.enable [ "nixpkgs-search-options.service" ]
        ++ lib.optionals cfg.packages.enable [ "nixpkgs-search-packages.service" ];

      environment = lib.mkIf (!useConnectionStringFile) {
        NIXPKGS_SEARCH_DATABASE = databaseUrl;
      };

      path = [ cfg.package ];

      script = let
        indexFlags = lib.concatStringsSep " " cfg.indexFlags;
        indexCommands = lib.concatMapStringsSep "\n" (rev: ''
          echo "Indexing revision: ${rev}"
          nixpkgs-search index ${indexFlags} "${rev}" || true
        '') cfg.indexOnStart;
      in
        if useConnectionStringFile then ''
          # Read database connection string from file
          if [ ! -f "${cfg.database.connectionStringFile}" ]; then
            echo "Error: connectionStringFile not found: ${cfg.database.connectionStringFile}" >&2
            exit 1
          fi
          export NIXPKGS_SEARCH_DATABASE="$(cat "${cfg.database.connectionStringFile}")"

          ${indexCommands}
        '' else ''
          ${indexCommands}
        '';

      serviceConfig = {
        Type = "oneshot";
        User = cfg.user;
        Group = cfg.group;
        RemainAfterExit = true;

        # Hardening
        NoNewPrivileges = true;
        ProtectSystem = "strict";
        ProtectHome = true;
        PrivateTmp = true;
        PrivateDevices = true;
        ProtectKernelTunables = true;
        ProtectKernelModules = true;
        ProtectControlGroups = true;
        RestrictNamespaces = true;
        RestrictRealtime = true;
        RestrictSUIDSGID = true;
        LockPersonality = true;

        ReadWritePaths = [ cfg.dataDir ];
        WorkingDirectory = cfg.dataDir;
        StateDirectory = lib.mkIf (cfg.dataDir == "/var/lib/nixpkgs-search") "nixpkgs-search";
      };
    };

    # Options MCP server
    systemd.services.nixpkgs-search-options = lib.mkIf cfg.options.enable
      (mkServiceConfig "Options" "options" cfg.options.http // {
        after = (mkServiceConfig "Options" "options" cfg.options.http).after
          ++ lib.optionals (cfg.indexOnStart != []) [ "nixpkgs-search-index.service" ];
      });

    # Packages MCP server
    systemd.services.nixpkgs-search-packages = lib.mkIf cfg.packages.enable
      (mkServiceConfig "Packages" "packages" cfg.packages.http // {
        after = (mkServiceConfig "Packages" "packages" cfg.packages.http).after
          ++ lib.optionals (cfg.indexOnStart != []) [ "nixpkgs-search-index.service" ];
      });

    # Open firewall ports if configured
    networking.firewall = lib.mkMerge [
      (lib.mkIf cfg.options.openFirewall (let
        addressParts = lib.splitString ":" cfg.options.http.address;
        port = lib.toInt (lib.last addressParts);
      in {
        allowedTCPPorts = [ port ];
      }))
      (lib.mkIf cfg.packages.openFirewall (let
        addressParts = lib.splitString ":" cfg.packages.http.address;
        port = lib.toInt (lib.last addressParts);
      in {
        allowedTCPPorts = [ port ];
      }))
    ];
  };
}