From f2d46f92e6411aba0c97cb05ca6cdffda1b0ee0a Mon Sep 17 00:00:00 2001 From: chenxu Date: Thu, 24 Sep 2026 18:31:49 +0800 Subject: [PATCH] feat: improve networking, deployment, and roadmap Add managed SSH connection pools and integrate them with agent networking. Improve KCP initialization and repair incomplete UI authentication secrets. Add bilingual roadmap and update architecture and configuration documentation. Signed-off-by: chenxu --- README.md | 17 +- README.zh-CN.md | 17 +- ROADMAP.md | 107 ++ ROADMAP.zh-CN.md | 107 ++ apps/rlark/docs/architecture.md | 24 +- apps/rlark/docs/reference/configuration.md | 8 +- apps/rlark/docs/zh/architecture.md | 24 +- apps/rlark/docs/zh/reference/configuration.md | 8 +- apps/rlark/pkg/agent/container/network.go | 15 +- apps/rlark/pkg/network/nodeserver/metrics.go | 2 +- apps/rlark/pkg/network/nodeserver/server.go | 6 +- apps/rlark/pkg/network/sshdialer/config.go | 105 ++ apps/rlark/pkg/network/sshdialer/dialer.go | 327 ++++++ .../pkg/network/sshdialer/dialer_test.go | 957 ++++++++++++++++++ .../pkg/network/sshdialer/maintenance.go | 127 +++ apps/rlark/pkg/network/sshdialer/pool.go | 333 ++++++ apps/rlark/pkg/network/sshdialer/ssh.go | 116 +++ .../pkg/network/sshdialer/tests/README.md | 27 + .../sshdialer/tests/conformance/main.go | 464 +++++++++ .../rlarkadm/deployer/kubernetes/install.go | 57 +- .../deployer/kubernetes/install_test.go | 25 + 21 files changed, 2784 insertions(+), 89 deletions(-) create mode 100644 ROADMAP.md create mode 100644 ROADMAP.zh-CN.md create mode 100644 apps/rlark/pkg/network/sshdialer/config.go create mode 100644 apps/rlark/pkg/network/sshdialer/dialer.go create mode 100644 apps/rlark/pkg/network/sshdialer/dialer_test.go create mode 100644 apps/rlark/pkg/network/sshdialer/maintenance.go create mode 100644 apps/rlark/pkg/network/sshdialer/pool.go create mode 100644 apps/rlark/pkg/network/sshdialer/ssh.go create mode 100644 apps/rlark/pkg/network/sshdialer/tests/README.md create mode 100644 apps/rlark/pkg/network/sshdialer/tests/conformance/main.go diff --git a/README.md b/README.md index af3dec6..5e04cc0 100644 --- a/README.md +++ b/README.md @@ -36,22 +36,7 @@ Manage cross-cluster embodied intelligence workloads through a unified cloud-nat ## Roadmap -### Available Today - -- Kubernetes-based management of cloud GPU clusters and edge devices -- Declarative Job/Task orchestration for distributed training and Ray workloads -- Cross-cluster resource management through Domain and Node CRDs -- Cross-cluster Pod networking over TUN devices, gVisor netstack, and SSH tunnels -- X.509 and SSH certificate-based authentication -- Prometheus metrics, real-time Pod logs, and a web management UI - -### Planned - -- Docker runtime support for lightweight data planes -- Raw runtime support for hosts and edge devices without a container orchestrator -- More complete account, role, and permission management -- Continued web UI usability and workflow improvements -- Cross-cluster network throughput, latency, and resource-efficiency optimizations +See the [RLark Roadmap](ROADMAP.md) for planned runtime, account management, and custom workload resource improvements. ## Architecture Overview diff --git a/README.zh-CN.md b/README.zh-CN.md index 44e5bca..ca7439f 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -36,22 +36,7 @@ ## Roadmap -### 已实现 - -- 基于 Kubernetes 统一纳管云端 GPU 集群与端侧设备 -- 面向分布式训练和 Ray 工作负载的声明式 Job/Task 编排 -- 通过 Domain 和 Node CRD 管理跨集群资源 -- 基于 TUN 设备、gVisor 协议栈和 SSH 隧道的跨集群 Pod 网络 -- 基于 X.509 和 SSH 证书的身份认证 -- Prometheus 指标、Pod 实时日志和 Web 管理界面 - -### 未来计划 - -- 面向轻量数据面的 Docker 运行时支持 -- 面向无容器编排环境主机和端侧设备的 Raw 运行时支持 -- 更完善的账号、角色与权限管理 -- 持续优化 Web UI 易用性和操作流程 -- 优化跨集群网络吞吐、时延与资源效率 +计划中的运行时、账号管理与自定义任务资源改进，请参阅 [RLark 路线图](ROADMAP.zh-CN.md)。 ## 架构概览 diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 0000000..52d8d3e --- /dev/null +++ b/ROADMAP.md @@ -0,0 +1,107 @@ +# RLark Roadmap + +[English](ROADMAP.md) | [简体中文](ROADMAP.zh-CN.md) + +This roadmap describes how the open-source RLark project will evolve from a Kubernetes-based workload platform into open, extensible infrastructure for cross-cluster embodied intelligence. RLark will provide stable resource models, APIs, and extension mechanisms for integrating external systems and adapting deployments to different environments. Priorities and implementation details may change as the project evolves. + +## Design Principles + +- Keep tenant isolation, identity, authorization, resource ownership, and lifecycle state in the RLark core rather than delegating security boundaries to extensions. +- Integrate extensions as independently deployed services through versioned gRPC or HTTP APIs instead of loading code into RLark processes. +- Define focused extension contracts for each domain instead of one unrestricted plugin API. +- Support capability discovery so the control plane and API consumers do not assume every data plane is Kubernetes or supports the same operations. +- Preserve declarative, idempotent reconciliation so operations can recover from retries, service restarts, and network failures. + +## Phase 1: Runtime and Workload Extensibility + +### Data Plane Runtimes + +The data plane currently supports Kubernetes only. Introduce a versioned Runtime Provider contract and add: + +- **Docker**: run and manage tasks directly with Docker for lightweight environments where Kubernetes is not suitable. +- **Raw**: run and manage tasks directly on hosts or edge devices without a container runtime or orchestrator. +- **Capability discovery**: advertise supported workload types, logs, exec, networking, storage, devices, and other runtime operations so unsupported features can be rejected before scheduling. + +The existing Kubernetes implementation should become the first implementation of the same runtime contract used by Docker and Raw providers. + +### Custom Workload Resources + +The Kubernetes data plane currently supports three built-in workload resource types: Deployment, StatefulSet, and DaemonSet. + +Introduce a service-based Workload Adapter contract for additional Kubernetes resources and external workload systems. The contract should cover the complete lifecycle, including validation, rendering or creation, status observation, updates, and deletion. It should support both: + +- **Declarative translation**: convert an RLark Task into resources such as RayJob, PyTorchJob, or a custom Kubernetes resource, while the Agent retains responsibility for applying and observing them. +- **Delegated execution**: submit and reconcile work through an external service such as Slurm, a cloud training service, or a domain-specific robotics platform. + +The exact API remains under design. Requests must be versioned and idempotent, and generated resources must pass RLark authorization and policy checks before execution. + +### Device Capability Foundation + +RLark should elevate embodied devices from runtime implementation details into reusable platform primitives. This provides a common foundation for discovering, observing, allocating, and controlling cameras, robot arms, sensors, and accelerators without hard-coding every device type into RLark core. + +- **Device model and inventory**: stable identities, ownership and scope, location and topology, device class, vendor metadata, structured capabilities, endpoints, and lifecycle state. +- **Discovery and synchronization**: service-based Device Providers discover physical or virtual devices and continuously reconcile their health, availability, capabilities, and connection state. +- **Observation**: uniform APIs and events for status, telemetry, logs, alarms, and capability changes. +- **Controlled operations**: capability-described commands such as opening a camera, capturing a frame, switching a robot mode, or resetting a device, with validation, timeout, idempotency, and result reporting. +- **Data channels**: authorized access to real-time streams and large device data without forcing video frames or telemetry through the control-plane resource store. +- **Allocation and leases**: exclusive or shared reservations, expiration, renewal, release, and conflict handling so interactive control and scheduled Tasks cannot use a device unsafely at the same time. +- **Safety and authorization**: operation-level permissions, policy checks, emergency-stop and other safety boundaries where supported, complete auditing, and explicit distinction between read-only observation and physical control. + +The existing embodied-runtime robot and camera APIs can serve as initial use cases, while the platform contract should remain extensible to new device classes and vendor-specific operations. Generic metadata and capability schemas should be combined with typed, versioned operation interfaces where safety or interoperability requires stronger semantics. + +## Phase 2: Identity, Authorization, and Resource Scopes + +Improve account management and provide reusable security primitives without prescribing a single tenancy model: + +- Users, service accounts, machine identities, roles, and permissions. +- Explicit ownership and scope for Domains, Nodes, Jobs, Workflows, Tasks, runtimes, credentials, policies, and quotas. +- Delegated, short-lived, renewable, and revocable credentials for Agents, extension services, and external controllers. +- End-to-end propagation of verified identity and scope context across the Gateway, controllers, database, Server, Agents, and extension services. +- Scope-aware authorization for API access, list and watch operations, logs, exec, proxying, events, and secrets. +- Audited privileged access and integration with external identity providers such as OIDC where appropriate. + +These primitives should allow operators to model users, projects, workspaces, organizations, or other security boundaries without requiring RLark to prescribe a single tenancy workflow. Core authorization and resource isolation must still be enforced by RLark rather than delegated entirely to schedulers or policy extensions. + +## Phase 3: Service-Based Extension Framework + +Provide a common operational framework for independently deployed extension services while keeping separate contracts for different extension domains. + +Initial extension types: + +- **Runtime Provider**: execute Tasks on Kubernetes, Docker, Raw hosts, or other runtimes. +- **Workload Adapter**: translate or delegate custom workload lifecycles. +- **Device Provider**: discover devices, publish structured capabilities and status, and expose authorized observation and control operations. +- **Admission and Policy Service**: validate or mutate requests according to deployment-specific policy. +- **Event Consumer**: integrate approvals, notifications, metering, auditing, and automation through webhooks or event streams. + +The framework should provide: + +- Declarative extension registration, health checks, API version negotiation, and capability discovery. +- mTLS service identity, least-privilege authorization, request auditing, and secret-safe configuration. +- Deadlines, retries, idempotency keys, circuit breaking, and reconciliation after lost events or partial failures. +- `ExtensionDefinition` and scoped `ExtensionBinding` resources, separating installation and trust decisions from scope-specific configuration. +- Protocol specifications, SDKs as convenience libraries, conformance tests, and reference implementations. + +Initially, platform administrators should approve extension endpoints. Allowing regular users to register arbitrary endpoints requires additional controls against data exfiltration, credential misuse, and server-side request forgery. + +## Phase 4: Open Integration and Customization + +Provide stable, composable interfaces so operators and developers can add domain capabilities without modifying RLark core: + +- **Extension resources and external controllers**: define resources such as `TrainingPipeline` or `RobotMission`, then reconcile them into RLark Jobs and Tasks through stable APIs. +- **Stable public APIs and events**: prevent external systems and extensions from depending on internal database schemas or implementation details. +- **Device APIs and data-plane access**: support integrations such as device inventories, status views, remote operations, and live camera access through stable discovery, control, event, and streaming contracts. +- **Schema-driven UI contributions**: generate forms and resource views from schemas before considering arbitrary frontend code or micro-frontends. +- **Quotas, policies, and metering events**: provide foundations for approvals, resource governance, usage accounting, and external system integration. + +RLark should expose composable primitives rather than hard-code a particular application workflow. Deployment-specific onboarding, user interfaces, approvals, and domain resources can be implemented outside RLark core. + +## Future Extension Areas + +After the foundational contracts are proven, evaluate focused service interfaces for: + +- Scheduler filtering and scoring based on topology, locality, latency, cost, quota, or domain-specific constraints. +- Dataset, model, artifact, object storage, and secret providers. +- Logs, metrics, traces, and alert integrations based on open observability standards. + +These interfaces should be introduced from concrete use cases rather than through a single generic plugin abstraction. diff --git a/ROADMAP.zh-CN.md b/ROADMAP.zh-CN.md new file mode 100644 index 0000000..36d101d --- /dev/null +++ b/ROADMAP.zh-CN.md @@ -0,0 +1,107 @@ +# RLark 路线图 + +[English](ROADMAP.md) | [简体中文](ROADMAP.zh-CN.md) + +本文档描述开源项目 RLark 如何从基于 Kubernetes 的任务平台演进为开放、可扩展的跨集群具身智能基础设施。RLark 将提供稳定的资源模型、接口与扩展机制，便于在不同部署环境中集成外部系统并按需扩展能力。随着项目演进，优先级和具体实现方式可能会有所调整。 + +## 设计原则 + +- 租户隔离、身份、授权、资源归属和生命周期状态由 RLark 核心负责，不把安全边界委托给扩展服务。 +- 扩展以独立服务部署，通过版本化的 gRPC 或 HTTP API 接入，而不是向 RLark 进程内加载代码。 +- 按领域定义职责清晰的扩展协议，不设计一个无边界的通用插件接口。 +- 提供能力发现机制，使控制面和 API 使用方不再假设所有数据面都是 Kubernetes 或支持相同操作。 +- 保持声明式、幂等的协调机制，使操作能够从重试、服务重启和网络故障中恢复。 + +## 第一阶段：运行时与任务扩展 + +### 数据面运行时 + +目前数据面仅支持 Kubernetes 模式。后续将定义版本化的 Runtime Provider 协议，并增加： + +- **Docker**：面向不适合部署 Kubernetes 的轻量环境，直接通过 Docker 运行和管理任务。 +- **Raw**：面向没有容器运行时或编排系统的主机及端侧设备，直接运行和管理任务。 +- **能力发现**：声明支持的任务类型、日志、Exec、网络、存储、设备和其他运行时操作，使平台能够在调度前拒绝不受支持的功能。 + +现有 Kubernetes 实现应成为 Runtime Provider 协议的首个实现，并与 Docker、Raw Provider 使用相同的核心契约。 + +### 自定义任务资源 + +目前 Kubernetes 数据面仅支持 Deployment、StatefulSet 和 DaemonSet 三种内置任务资源类型。 + +后续将定义服务化的 Workload Adapter 协议，用于接入更多 Kubernetes 资源和外部任务系统。协议应覆盖校验、渲染或创建、状态观察、更新和删除等完整生命周期，并支持两种模式： + +- **声明式转换**：将 RLark Task 转换为 RayJob、PyTorchJob 或自定义 Kubernetes 资源，由 Agent 负责应用资源并观察状态。 +- **委托执行**：通过 Slurm、云训练服务或机器人任务平台等外部服务提交和协调任务。 + +具体 API 仍在设计中。调用必须支持版本化和幂等处理，生成的资源在执行前必须再次经过 RLark 的授权与策略校验。 + +### 设备能力基座 + +RLark 应将具身设备从运行时实现细节提升为可复用的平台原语，为摄像头、机械臂、传感器与加速器提供统一的发现、观察、分配和控制基础，而不需要将每种设备类型硬编码到 RLark 核心。 + +- **设备模型与资产目录**：提供稳定身份、归属与作用域、位置与拓扑、设备类别、厂商信息、结构化能力、访问端点和生命周期状态。 +- **发现与同步**：通过服务化 Device Provider 发现物理或虚拟设备，并持续协调其健康、可用性、能力和连接状态。 +- **观察能力**：通过统一 API 与事件获取状态、遥测、日志、告警和能力变化。 +- **受控操作**：根据设备能力描述开放摄像头、抓取图像、切换机械臂模式或重置设备等命令，并提供参数校验、超时、幂等和结果回报。 +- **数据通道**：在完成授权后访问实时流和大体量设备数据，避免将视频帧或高频遥测数据写入控制面资源存储。 +- **分配与租约**：支持独占或共享预约、过期、续期、释放和冲突处理，避免交互式控制与调度任务同时不安全地使用同一设备。 +- **安全与授权**：提供操作级权限、策略校验、设备支持范围内的急停等安全边界、完整审计，并明确区分只读观察与物理控制权限。 + +现有 embodied-runtime 的机器人和摄像头 API 可以作为首批验证场景，但平台协议应能够扩展到新的设备类别和厂商专有操作。通用元数据与能力 Schema 可以保证开放性；对于涉及安全或互操作性的操作，则应采用强类型、版本化的接口。 + +## 第二阶段：身份、授权与资源作用域 + +完善账号管理并提供可复用的安全原语，但不限定唯一的租户组织方式： + +- 用户、服务账号、机器身份、角色与权限。 +- Domain、Node、Job、Workflow、Task、Runtime、凭证、策略和配额的明确归属与作用域。 +- 面向 Agent、扩展服务和外部 Controller 的可委托、短期、可续期和可吊销凭证。 +- 经过验证的身份与作用域上下文贯穿 Gateway、Controller、数据库、Server、Agent 和扩展服务。 +- API、List/Watch、日志、Exec、代理、事件和 Secret 的作用域级授权。 +- 特权操作审计，并按需对接 OIDC 等外部身份提供方。 + +这些原语应允许部署者按需建立用户、项目、Workspace、组织或其他安全边界，而不要求 RLark 固化唯一的租户流程。核心授权与资源隔离仍必须由 RLark 强制保证，不能完全委托给调度器或策略扩展。 + +## 第三阶段：服务化扩展框架 + +为独立部署的扩展服务提供统一的运行机制，同时为不同扩展领域保留独立协议。 + +首批扩展类型包括： + +- **Runtime Provider**：在 Kubernetes、Docker、Raw 主机或其他运行环境中执行 Task。 +- **Workload Adapter**：转换或委托自定义任务的完整生命周期。 +- **Device Provider**：发现设备、发布结构化能力与状态，并开放经过授权的观察和控制操作。 +- **Admission 与 Policy Service**：根据部署环境定义的策略校验或修改请求。 +- **Event Consumer**：通过 Webhook 或事件流接入审批、通知、计量、审计和自动化流程。 + +扩展框架应提供： + +- 声明式扩展注册、健康检查、API 版本协商和能力发现。 +- mTLS 服务身份、最小权限授权、调用审计和安全的 Secret 配置方式。 +- 超时、重试、幂等键、熔断，以及在事件丢失或部分失败后的状态协调。 +- `ExtensionDefinition` 和限定作用域的 `ExtensionBinding`，将扩展安装及信任决策与具体作用域配置分离。 +- 协议规范、便捷使用的 SDK、一致性测试和参考实现。 + +初期扩展服务地址应由平台管理员审核。允许普通用户注册任意地址前，需要防范数据外泄、凭证滥用和服务端请求伪造等风险。 + +## 第四阶段：开放集成与二次开发能力 + +提供稳定、可组合的开放接口，使部署者和二次开发者无需修改 RLark 核心代码即可扩展领域能力： + +- **扩展资源与外部 Controller**：允许定义 `TrainingPipeline`、`RobotMission` 等领域资源，并通过稳定 API 将其协调为 RLark Job 和 Task。 +- **稳定的公开 API 与事件**：避免外部系统和扩展依赖内部数据库结构或实现细节。 +- **设备 API 与数据面访问**：通过稳定的发现、控制、事件和流式传输协议，支持设备资产目录、状态视图、远程操作和实时摄像头等集成场景。 +- **Schema 驱动的 UI 扩展**：优先通过 Schema 生成表单和资源视图，而不是直接引入任意前端代码或微前端。 +- **配额、策略和计量事件**：为审批、资源治理、用量核算和外部系统集成提供基础能力。 + +RLark 应提供可组合的基础原语，而不是固化某一种应用流程。具体的用户接入、界面、审批和领域资源可根据部署需求在 RLark 之外实现。 + +## 后续扩展方向 + +在基础协议得到实际验证后，再根据具体场景评估以下服务化扩展接口： + +- 基于拓扑、数据位置、网络时延、成本、配额或领域约束的调度过滤与评分。 +- 数据集、模型、制品、对象存储和 Secret Provider。 +- 基于开放可观测标准的日志、指标、链路追踪和告警集成。 + +这些接口应从真实用例出发逐步引入，而不是通过一个通用插件抽象提前覆盖所有场景。 diff --git a/apps/rlark/docs/architecture.md b/apps/rlark/docs/architecture.md index fadf8fd..0f94553 100644 --- a/apps/rlark/docs/architecture.md +++ b/apps/rlark/docs/architecture.md @@ -185,21 +185,27 @@ func (a *containerNetworkAdapter) GetContainerNetworkDial(...) (utils.Dial, erro return dialer.DialContext(ctx, "tcp", net.JoinHostPort(targetPod.LocalIP, "5700")) } // 2. Cross-cluster → SSH tunnel - return a.sshDialer.DialContext(ctx, domainID, sshAddr, cert, key, target) + dialer := a.sshDialerManager.GetDialer(sshdialer.DomainInfo{ + ID: domainID, SSHAddress: sshAddr, Certificate: cert, PrivateKey: key, + }) + return dialer.DialContext(ctx, "tcp", target) } ``` -### 4.5 SSHDialer +### 4.5 SSH Dialer Manager -**Key File**: [container/ssh\_dialer.go](https://github.com/RLinf/RLark/tree/main/apps/rlark/pkg/agent/container/ssh_dialer.go) +**Key Package**: [network/sshdialer](https://github.com/RLinf/RLark/tree/main/apps/rlark/pkg/network/sshdialer) -Per-Domain SSH connection pool. Design highlights: +`DialerManager` owns the connection pools and `GetDialer(DomainInfo)` returns a dialer bound to one Domain. Design highlights: -- Each Domain starts with one SSH connection and grows up to four connections when all existing connections have active channels -- New channels use the least-loaded connection; idle physical connections are reclaimed by background GC -- Auto-reconnect on disconnect; concurrent requests wait during reconnection instead of creating separate connections +- Each Domain establishes one shared, multiplexed SSH transport during recovery; concurrent callers wait for the same bounded handshake instead of creating a reconnect storm +- New channels use the least-loaded healthy connection; idle physical connections are reclaimed by background GC +- Caller cancellation only stops that caller from waiting and does not cancel or penalize the shared recovery attempt +- Definite transport failures are closed immediately; channel-local timeouts stay isolated, while keepalive timeouts drain active channels before forced closure - Exponential backoff on reconnection failure (1s → 2s → 4s → ... → 30s) -- Background GC closes idle connections (default 24 hour timeout) +- SSH handshakes have a hard timeout, and manager shutdown cancels pending reconnects +- Certificate, private-key, or endpoint changes advance a Domain generation and drain transports created from the old generation +- Background GC closes idle or over-age connections and evicts inactive Domain entries (default 24 hour timeout) - Data-path activity timestamps are updated atomically and rate-limited to avoid a mutex on every read and write ### 4.6 Embodied Runtime @@ -238,7 +244,7 @@ sequenceDiagram participant SA as Sidecar A participant NS as NodeServer A participant CNA as ContainerNetworkAdapter - participant SD as SSHDialer + participant SD as SSH Dialer Manager participant SRV as Server participant NB as NodeServer B participant SB as Sidecar Proxy B diff --git a/apps/rlark/docs/reference/configuration.md b/apps/rlark/docs/reference/configuration.md index 2e71a02..c47cd1f 100644 --- a/apps/rlark/docs/reference/configuration.md +++ b/apps/rlark/docs/reference/configuration.md @@ -169,11 +169,9 @@ Data plane agent. Deployed on each cluster or node. Manages node registration, T | `--pod-orphan-sweep-interval` | duration | `5m` | Interval between agent-scoped management Pod orphan sweeps | | `--pod-orphan-sweep-page-size` | int | `200` | Management Pods processed per orphan sweep page | | `--pod-stale-ttl` | duration | `15m` | Time a missing local Pod is retained as `Unknown`/stale before its management Pod is deleted | - -The Pod orphan sweep is a fallback for missed local delete events. It deletes only agent-scoped mirrors whose local Pod UID or verified management Task UID is no longer current. Legacy mirrors are adopted only when the UID-named mirror, live local Pod annotations, management namespace, Task UID, and available domain all agree; ambiguous legacy objects remain untouched and require manual cleanup. A delayed delete intentionally preserves a same-name replacement, so stale mirrors may remain until the next sweep interval. | `--rlark-server-ssh-address` | string | `""` | RLark server SSH address (user@host:port) | | `--rlark-server-ssh-host-key` | string | `""` | RLark server SSH host key | -| `--ssh-max-connections-per-domain` | int | `4` | Maximum adaptive physical SSH connections per Domain | +| `--ssh-max-connections-per-domain` | int | `4` | Upper bound for physical SSH connections per Domain; the pool adds transports as per-connection channel load grows | | `--image` | string | `""` | RLark network sidecar image | | `--enable-same-cluster-direct` | bool | `true` | Enable same-cluster direct Pod access | | `--enable-cross-cluster-direct` | bool | `true` | Enable cross-cluster direct Pod access | @@ -191,6 +189,10 @@ The Pod orphan sweep is a fallback for missed local delete events. It deletes on | `--kube-burst` | int | `8000` | Kubernetes client burst | | `--kube-timeout` | duration | `0` | Kubernetes client request timeout | +The Pod orphan sweep is a fallback for missed local delete events. It deletes only agent-scoped mirrors whose local Pod UID or verified management Task UID is no longer current. Legacy mirrors are adopted only when the UID-named mirror, live local Pod annotations, management namespace, Task UID, and available domain all agree; ambiguous legacy objects remain untouched and require manual cleanup. A delayed delete intentionally preserves a same-name replacement, so stale mirrors may remain until the next sweep interval. + +The SSH pool starts with one transport per active Domain and expands in the background when a transport carries sustained channel load, up to `--ssh-max-connections-per-domain`. Requests continue on existing transports while expansion is in progress, and idle transports are reclaimed automatically. + !!! tip "`--mode` values" - `cluster`: Cluster-level agent only, manages cluster-wide resources - `node`: Node-level agent only, manages Tasks on a single node diff --git a/apps/rlark/docs/zh/architecture.md b/apps/rlark/docs/zh/architecture.md index fe93ff9..594203f 100644 --- a/apps/rlark/docs/zh/architecture.md +++ b/apps/rlark/docs/zh/architecture.md @@ -185,21 +185,27 @@ func (a *containerNetworkAdapter) GetContainerNetworkDial(...) (utils.Dial, erro return dialer.DialContext(ctx, "tcp", net.JoinHostPort(targetPod.LocalIP, "5700")) } // 2. 跨集群 → SSH 隧道 - return a.sshDialer.DialContext(ctx, domainID, sshAddr, cert, key, target) + dialer := a.sshDialerManager.GetDialer(sshdialer.DomainInfo{ + ID: domainID, SSHAddress: sshAddr, Certificate: cert, PrivateKey: key, + }) + return dialer.DialContext(ctx, "tcp", target) } ``` -### 4.5 SSHDialer +### 4.5 SSH Dialer Manager -**关键文件**：[container/ssh\_dialer.go](https://github.com/RLinf/RLark/tree/main/apps/rlark/pkg/agent/container/ssh_dialer.go) +**关键包**：[network/sshdialer](https://github.com/RLinf/RLark/tree/main/apps/rlark/pkg/network/sshdialer) -按 Domain 维护的 SSH 连接池，设计要点： +`DialerManager` 维护所有 Domain 的连接池，`GetDialer(DomainInfo)` 返回绑定单个 Domain 的拨号器。设计要点： -- 每个 Domain 从一条 SSH 连接开始；现有连接均承载活跃 channel 时按需扩展，默认最多四条 -- 新 channel 选择负载最低的连接，空闲物理连接由后台 GC 回收 -- 连接断开时自动重连，重连期间并发请求等待而非各自新建 +- 每个 Domain 在故障恢复时只建立一条共享的多路复用 SSH transport；并发调用者等待同一个有超时上限的握手，避免形成重连风暴 +- 新 channel 选择负载最低的健康连接，空闲物理连接由后台 GC 回收 +- 调用者取消只停止自身等待，不会取消共享恢复，也不会增加 Domain 的重连退避 +- 明确的传输层故障立即关闭；channel 自身超时不会影响其他 channel，keepalive 超时则先等待活跃 channel 排空，再按宽限期强制关闭 - 重连失败指数退避（1s → 2s → 4s → ... → 30s） -- 后台 GC 关闭空闲超时连接（默认 24 小时） +- SSH 握手有硬超时，Manager 关闭时会取消仍在进行的重连 +- 证书、私钥或端点发生变化时推进 Domain generation，并排空旧 generation 的 transport +- 后台 GC 关闭空闲或超龄连接并回收不活跃的 Domain entry（默认 24 小时） - 数据路径上的活跃时间使用原子、限频更新，避免每次读写获取互斥锁 ### 4.6 Embodied Runtime @@ -238,7 +244,7 @@ sequenceDiagram participant SA as Sidecar A participant NS as NodeServer A participant CNA as ContainerNetworkAdapter - participant SD as SSHDialer + participant SD as SSH Dialer Manager participant SRV as Server participant NB as NodeServer B participant SB as Sidecar Proxy B diff --git a/apps/rlark/docs/zh/reference/configuration.md b/apps/rlark/docs/zh/reference/configuration.md index 187d0d8..f3be2c9 100644 --- a/apps/rlark/docs/zh/reference/configuration.md +++ b/apps/rlark/docs/zh/reference/configuration.md @@ -169,11 +169,9 @@ rlark-controller-manager \ | `--pod-orphan-sweep-interval` | duration | `5m` | Agent 范围内管理 Pod 孤儿扫描间隔 | | `--pod-orphan-sweep-page-size` | int | `200` | 每页扫描的管理 Pod 数量 | | `--pod-stale-ttl` | duration | `15m` | 本地 Pod 缺失后以 `Unknown`/陈旧状态保留、再删除管理 Pod 的时长 | - -Pod 孤儿扫描用于兜底处理遗漏的本地删除事件。它只删除 Agent 作用域内、本地 Pod UID 或已验证管理 Task UID 已失效的镜像。旧版镜像仅在以 UID 命名的镜像、存活本地 Pod 注解、管理命名空间、Task UID 以及可用的 Domain 全部一致时才会被接管；身份不明确的旧对象保持不变，需要手动清理。延迟删除会有意保留同名替代 Pod，因此旧镜像可能持续到下一次扫描。 | `--rlark-server-ssh-address` | string | `""` | RLark Server SSH 地址（user@host:port） | | `--rlark-server-ssh-host-key` | string | `""` | RLark Server SSH Host Key | -| `--ssh-max-connections-per-domain` | int | `4` | 每个 Domain 按负载自适应扩展的物理 SSH 连接上限 | +| `--ssh-max-connections-per-domain` | int | `4` | 每个 Domain 的物理 SSH 连接上限；单连接 channel 负载增长时连接池会弹性增加 transport | | `--image` | string | `""` | RLark 网络 Sidecar 镜像 | | `--enable-same-cluster-direct` | bool | `true` | 启用同集群 Pod 直接访问 | | `--enable-cross-cluster-direct` | bool | `true` | 启用跨集群 Pod 直接访问 | @@ -191,6 +189,10 @@ Pod 孤儿扫描用于兜底处理遗漏的本地删除事件。它只删除 Age | `--kube-burst` | int | `8000` | Kubernetes 客户端 Burst | | `--kube-timeout` | duration | `0` | Kubernetes 客户端请求超时 | +Pod 孤儿扫描用于兜底处理遗漏的本地删除事件。它只删除 Agent 作用域内、本地 Pod UID 或已验证管理 Task UID 已失效的镜像。旧版镜像仅在以 UID 命名的镜像、存活本地 Pod 注解、管理命名空间、Task UID 以及可用的 Domain 全部一致时才会被接管；身份不明确的旧对象保持不变，需要手动清理。延迟删除会有意保留同名替代 Pod，因此旧镜像可能持续到下一次扫描。 + +SSH 连接池为每个活跃 Domain 从一条 transport 开始；单条 transport 持续承载较高 channel 负载时，连接池会在后台扩容，最多不超过 `--ssh-max-connections-per-domain`。扩容期间请求继续使用已有 transport，空闲 transport 会自动回收。 + !!! tip "`--mode` 参数说明" - `cluster`：仅运行集群级 Agent，管理集群范围内的资源 - `node`：仅运行节点级 Agent，管理单个节点的 Task diff --git a/apps/rlark/pkg/agent/container/network.go b/apps/rlark/pkg/agent/container/network.go index 29d7a10..c94ab5d 100644 --- a/apps/rlark/pkg/agent/container/network.go +++ b/apps/rlark/pkg/agent/container/network.go @@ -19,6 +19,7 @@ import ( "github.com/rlinf/rlark/apps/rlark/pkg/common" "github.com/rlinf/rlark/apps/rlark/pkg/log" nodeservermetrics "github.com/rlinf/rlark/apps/rlark/pkg/network/nodeserver" + "github.com/rlinf/rlark/apps/rlark/pkg/network/sshdialer" "github.com/rlinf/rlark/apps/rlark/pkg/utils" ) @@ -72,7 +73,7 @@ type containerNetworkAdapter struct { domainPeerLister listerv1alpha1.DomainPeerLister managementPodLister listerv1alpha1.PodLister sshAddr string - sshDialer *SSHDialer + sshDialerManager *sshdialer.DialerManager enableSameClusterDirect bool enableCrossClusterDirect bool @@ -99,7 +100,7 @@ func NewContainerNetworkAdapter( domainPeerLister: domainPeerLister, managementPodLister: managementPodLister, sshAddr: sshAddr, - sshDialer: NewSSHDialer(SSHDialerConfig{ + sshDialerManager: sshdialer.NewDialerManager(sshdialer.Config{ HostKeyCallback: hostKeyCallback, OnReconnect: nodeservermetrics.OnReconnect(), MaxConnectionsPerDomain: sshMaxConnectionsPerDomain, @@ -112,7 +113,7 @@ func NewContainerNetworkAdapter( } func (a *containerNetworkAdapter) Close() error { - return a.sshDialer.Close() + return a.sshDialerManager.Close() } // makeHostKeyCallback 解析 SSH 主机公钥字符串，返回对应的 HostKeyCallback。 @@ -243,7 +244,13 @@ func (a *containerNetworkAdapter) GetContainerNetworkDial(ctx context.Context, c target := fmt.Sprintf("%s.%s.%s.agent-node:5700", targetPod.LocalIP, targetPod.Node, agentID) logger.V(1).Info("Target pod is in a different cluster, using control plane proxy", "targetPod", target) return func(ctx context.Context) (net.Conn, error) { - conn, err := a.sshDialer.DialContext(ctx, cred.DomainID, a.sshAddr, dpeer.Spec.Cert, dpeer.Spec.Key, target) + dialer := a.sshDialerManager.GetDialer(sshdialer.DomainInfo{ + ID: cred.DomainID, + SSHAddress: a.sshAddr, + Certificate: dpeer.Spec.Cert, + PrivateKey: dpeer.Spec.Key, + }) + conn, err := dialer.DialContext(ctx, "tcp", target) status := "success" if err != nil { status = "error" diff --git a/apps/rlark/pkg/network/nodeserver/metrics.go b/apps/rlark/pkg/network/nodeserver/metrics.go index cd60df2..6ccc9af 100644 --- a/apps/rlark/pkg/network/nodeserver/metrics.go +++ b/apps/rlark/pkg/network/nodeserver/metrics.go @@ -94,7 +94,7 @@ func (m *nodeServerMetrics) IncSSHReconnect(domain string) { m.sshReconnectTotal.WithLabelValues(domain).Inc() } -// OnReconnect returns a callback suitable for SSHDialerConfig.OnReconnect. +// OnReconnect returns a callback suitable for sshdialer.Config.OnReconnect. func OnReconnect() func(domainID string) { return func(domainID string) { metrics.IncSSHReconnect(domainID) diff --git a/apps/rlark/pkg/network/nodeserver/server.go b/apps/rlark/pkg/network/nodeserver/server.go index 6bd19da..61acde2 100644 --- a/apps/rlark/pkg/network/nodeserver/server.go +++ b/apps/rlark/pkg/network/nodeserver/server.go @@ -217,8 +217,6 @@ func (s *NodeServer[C]) closeConnections() { // handleConnection 处理来自本地进程的连接请求，读取目标地址并通过 dialer 连接到目标。 func (s *NodeServer[C]) handleConnection(ctx context.Context, conn *utils.WrapConn, cred C) { logger := log.FromContext(ctx) - ctx, cancel := context.WithTimeout(ctx, time.Second*30) - defer cancel() metrics.IncActive() defer metrics.DecActive() @@ -257,7 +255,9 @@ func (s *NodeServer[C]) handleConnection(ctx context.Context, conn *utils.WrapCo logger.Error(nil, "Failed to get target", "host", host, "err", err) return } - conn2, err = dial(ctx) + dialCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + conn2, err = dial(dialCtx) + cancel() if err != nil { logger.Error(nil, "Failed to connect to target", "host", host, "port", port, "err", err) return diff --git a/apps/rlark/pkg/network/sshdialer/config.go b/apps/rlark/pkg/network/sshdialer/config.go new file mode 100644 index 0000000..02fba26 --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/config.go @@ -0,0 +1,105 @@ +package sshdialer + +import ( + "errors" + "time" + + "golang.org/x/crypto/ssh" +) + +const ( + defaultIdleTimeout = 24 * time.Hour + defaultCleanupInterval = time.Minute + defaultSSHUser = "root" + defaultSSHTimeout = 10 * time.Second + defaultKeepaliveInterval = 30 * time.Second + defaultKeepaliveTimeout = 10 * time.Second + defaultKeepaliveDrainGrace = 30 * time.Second + defaultMaxConnections = 4 + defaultMaxConnectionAge = 24 * time.Hour + defaultMaxChannels = 4096 + defaultMaxChannelsPerConnection = 256 + maxReconnectBackoff = 30 * time.Second + initialReconnectBackoff = time.Second + activityUpdateInterval = time.Second +) + +const keepaliveRequest = "keepalive@openssh.com" + +var ( + ErrClosed = errors.New("ssh dialer: closed") + ErrDomainRemoved = errors.New("ssh dialer: domain removed") + ErrOverloaded = errors.New("ssh dialer: channel limit reached") +) + +// Config configures the SSH connection pool. +type Config struct { + IdleTimeout time.Duration `json:"idleTimeout,omitempty" yaml:"idleTimeout,omitempty"` + CleanupInterval time.Duration `json:"cleanupInterval,omitempty" yaml:"cleanupInterval,omitempty"` + SSHUser string `json:"sshUser,omitempty" yaml:"sshUser,omitempty"` + SSHTimeout time.Duration `json:"sshTimeout,omitempty" yaml:"sshTimeout,omitempty"` + InitialReconnectBackoff time.Duration `json:"initialReconnectBackoff,omitempty" yaml:"initialReconnectBackoff,omitempty"` + MaxReconnectBackoff time.Duration `json:"maxReconnectBackoff,omitempty" yaml:"maxReconnectBackoff,omitempty"` + KeepaliveInterval time.Duration `json:"keepaliveInterval,omitempty" yaml:"keepaliveInterval,omitempty"` + KeepaliveTimeout time.Duration `json:"keepaliveTimeout,omitempty" yaml:"keepaliveTimeout,omitempty"` + KeepaliveDrainGrace time.Duration `json:"keepaliveDrainGrace,omitempty" yaml:"keepaliveDrainGrace,omitempty"` + MaxConnectionsPerDomain int `json:"maxConnectionsPerDomain,omitempty" yaml:"maxConnectionsPerDomain,omitempty"` + MaxConnectionAge time.Duration `json:"maxConnectionAge,omitempty" yaml:"maxConnectionAge,omitempty"` + MaxChannelsPerDomain int `json:"maxChannelsPerDomain,omitempty" yaml:"maxChannelsPerDomain,omitempty"` + MaxChannelsPerConnection int `json:"maxChannelsPerConnection,omitempty" yaml:"maxChannelsPerConnection,omitempty"` + OnReconnect func(domainID string) `json:"-" yaml:"-"` + HostKeyCallback ssh.HostKeyCallback `json:"-" yaml:"-"` +} + +// DomainInfo identifies a remote SSH endpoint and its credentials. +type DomainInfo struct { + ID string + SSHAddress string + Certificate string + PrivateKey string +} + +func (c *Config) setDefaults() { + if c.IdleTimeout <= 0 { + c.IdleTimeout = defaultIdleTimeout + } + if c.CleanupInterval <= 0 { + c.CleanupInterval = defaultCleanupInterval + } + if c.SSHUser == "" { + c.SSHUser = defaultSSHUser + } + if c.SSHTimeout <= 0 { + c.SSHTimeout = defaultSSHTimeout + } + if c.InitialReconnectBackoff <= 0 { + c.InitialReconnectBackoff = initialReconnectBackoff + } + if c.MaxReconnectBackoff <= 0 { + c.MaxReconnectBackoff = maxReconnectBackoff + } + if c.KeepaliveInterval <= 0 { + c.KeepaliveInterval = defaultKeepaliveInterval + } + if c.KeepaliveTimeout <= 0 { + c.KeepaliveTimeout = defaultKeepaliveTimeout + } + if c.KeepaliveDrainGrace <= 0 { + c.KeepaliveDrainGrace = defaultKeepaliveDrainGrace + } + if c.MaxConnectionsPerDomain <= 0 { + c.MaxConnectionsPerDomain = defaultMaxConnections + } + if c.MaxConnectionAge <= 0 { + c.MaxConnectionAge = defaultMaxConnectionAge + } + if c.MaxChannelsPerDomain <= 0 { + c.MaxChannelsPerDomain = defaultMaxChannels + } + if c.MaxChannelsPerConnection <= 0 { + c.MaxChannelsPerConnection = defaultMaxChannelsPerConnection + } + if c.HostKeyCallback == nil { + c.HostKeyCallback = ssh.InsecureIgnoreHostKey() + } +} diff --git a/apps/rlark/pkg/network/sshdialer/dialer.go b/apps/rlark/pkg/network/sshdialer/dialer.go new file mode 100644 index 0000000..88d1aba --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/dialer.go @@ -0,0 +1,327 @@ +package sshdialer + +import ( + "context" + "fmt" + "net" + "sync" + "sync/atomic" + "time" + + "github.com/rlinf/rlark/apps/rlark/pkg/log" + "golang.org/x/crypto/ssh" +) + +// DialerManager owns independent SSH transport pools keyed by domain. +// +// Design principles: +// - GetDialer binds domain identity and credentials; Dialer only handles target dialing. +// - Each domain reuses multiplexed transports and grows the pool by load up to the configured limit. +// - Concurrent transport creation is coordinated, and failed reconnects use exponential backoff. +// - Definite transport failures are closed immediately; keepalive timeouts drain active channels first. +// - Idle transports are reclaimed in the background and shared state is concurrency-safe. + +// pooledSSHClient tracks one multiplexed physical SSH transport. +type pooledSSHClient struct { + client *ssh.Client + generation uint64 + active int + draining bool + lastUsedNanos atomic.Int64 + createdAt time.Time + keepaliveDone chan struct{} + keepaliveExit chan struct{} + drainTimer *time.Timer + closeOnce sync.Once +} + +func newPooledSSHClient(client *ssh.Client) *pooledSSHClient { + p := &pooledSSHClient{ + client: client, + createdAt: time.Now(), + keepaliveDone: make(chan struct{}), + keepaliveExit: make(chan struct{}), + } + p.touch() + return p +} + +func (p *pooledSSHClient) touch() { + now := time.Now().UnixNano() + last := p.lastUsedNanos.Load() + if now-last >= int64(activityUpdateInterval) { + p.lastUsedNanos.CompareAndSwap(last, now) + } +} + +func (p *pooledSSHClient) lastUsed() time.Time { + return time.Unix(0, p.lastUsedNanos.Load()) +} + +func (p *pooledSSHClient) close() { + p.closeOnce.Do(func() { + if p.drainTimer != nil { + p.drainTimer.Stop() + } + close(p.keepaliveDone) + _ = p.client.Close() + }) +} + +type domainEntry struct { + id string + domainInfo DomainInfo + generation uint64 + lastUsedNanos atomic.Int64 + + mu sync.Mutex + clients []*pooledSSHClient + next int + + // Coordinates one manager-owned recovery attempt for this domain. + reconnecting bool + reconnectCh chan struct{} + reconnectCancel context.CancelFunc + lastReconnectErr error + reconnectBackoff time.Duration + initialBackoff time.Duration + maxBackoff time.Duration + removed bool +} + +func (entry *domainEntry) touch() { + entry.lastUsedNanos.Store(time.Now().UnixNano()) +} + +func (entry *domainEntry) lastUsed() time.Time { + return time.Unix(0, entry.lastUsedNanos.Load()) +} + +// DialerManager owns SSH transports for all domains. +type DialerManager struct { + cfg Config + closed atomic.Bool + + mu sync.RWMutex + domains map[string]*domainEntry + + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup + closeOnce sync.Once + closeDone chan struct{} +} + +// Dialer opens connections through one domain's SSH transports. +type Dialer struct { + manager *DialerManager + entry *domainEntry +} + +// NewDialerManager creates and starts an SSH connection manager. +func NewDialerManager(cfg Config) *DialerManager { + cfg.setDefaults() + ctx, cancel := context.WithCancel(context.Background()) + d := &DialerManager{ + cfg: cfg, + domains: make(map[string]*domainEntry), + ctx: ctx, + cancel: cancel, + closeDone: make(chan struct{}), + } + d.wg.Add(1) + go d.cleanupLoop() + return d +} + +type activityConn struct { + net.Conn + onActivity func() + onRelease func() + onError func(error) + releaseOnce sync.Once + errorOnce sync.Once +} + +func (c *activityConn) release() { + c.releaseOnce.Do(c.onRelease) +} + +func (c *activityConn) Read(b []byte) (int, error) { + n, err := c.Conn.Read(b) + if n > 0 { + c.onActivity() + } + if err != nil && isSSHChannelTransportError(err) { + c.errorOnce.Do(func() { c.onError(err) }) + } + if err != nil { + c.release() + } + return n, err +} + +func (c *activityConn) Write(b []byte) (int, error) { + n, err := c.Conn.Write(b) + if n > 0 { + c.onActivity() + } + if err != nil && isSSHChannelTransportError(err) { + c.errorOnce.Do(func() { c.onError(err) }) + } + if err != nil { + c.release() + } + return n, err +} + +func (c *activityConn) Close() error { + err := c.Conn.Close() + c.release() + return err +} + +// GetDialer returns a dialer bound to a domain. The latest DomainInfo is used +// when a new physical SSH transport is established. +func (d *DialerManager) GetDialer(info DomainInfo) *Dialer { + entry := d.getOrCreate(info) + return &Dialer{manager: d, entry: entry} +} + +// RemoveDomain removes a domain, closes its transports, and clears credentials. +func (d *DialerManager) RemoveDomain(id string) { + d.mu.Lock() + entry := d.domains[id] + delete(d.domains, id) + d.mu.Unlock() + if entry == nil { + return + } + entry.mu.Lock() + entry.removed = true + entry.domainInfo = DomainInfo{ID: id} + entry.mu.Unlock() + entry.close() +} + +// DialContext opens a connection through the domain's SSH transport pool. +func (d *Dialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) { + if network != "tcp" { + return nil, fmt.Errorf("ssh dialer: unsupported network %q", network) + } + manager := d.manager + if manager.closed.Load() { + return nil, ErrClosed + } + + entry := d.entry + pooled, err := entry.borrow(ctx, manager) + if err != nil { + return nil, fmt.Errorf("ssh dialer: %w", err) + } + + conn, err := pooled.client.DialContext(ctx, "tcp", addr) + if err != nil { + entry.release(pooled) + if isSSHTransportError(err) { + log.GetLogger().Info("SSH channel dial failed with transport error, aborting transport", + "domain", entry.id, + "target", addr, + "err", err, + "errType", fmt.Sprintf("%T", err), + ) + entry.abort(pooled, "channel-dial-error") + } + return nil, fmt.Errorf("ssh proxy to %s: %w", addr, err) + } + + return &activityConn{ + Conn: conn, + onActivity: pooled.touch, + onRelease: func() { entry.release(pooled) }, + onError: func(err error) { + log.GetLogger().Info("SSH channel I/O failed with transport error, aborting transport", + "domain", entry.id, + "target", addr, + "err", err, + "errType", fmt.Sprintf("%T", err), + ) + entry.abort(pooled, "channel-io-error") + }, + }, nil +} + +// Dial opens a connection through the domain's SSH transport pool. +func (d *Dialer) Dial(network, addr string) (net.Conn, error) { + return d.DialContext(context.Background(), network, addr) +} + +// Close closes every managed transport and stops background maintenance. +func (d *DialerManager) Close() error { + d.closeOnce.Do(func() { + d.closed.Store(true) + d.cancel() + d.mu.RLock() + entries := make([]*domainEntry, 0, len(d.domains)) + for _, entry := range d.domains { + entries = append(entries, entry) + } + d.mu.RUnlock() + for _, entry := range entries { + entry.close() + } + d.wg.Wait() + close(d.closeDone) + }) + <-d.closeDone + return nil +} + +// Stats returns the number of transports available for new channels. +func (d *DialerManager) Stats() (open int) { + d.mu.RLock() + defer d.mu.RUnlock() + for _, entry := range d.domains { + entry.mu.Lock() + for _, client := range entry.clients { + if !client.draining { + open++ + } + } + entry.mu.Unlock() + } + return +} + +// PoolStats describes current physical transports and multiplexed channels. +type PoolStats struct { + Domains int + Open int + Draining int + Reconnecting int + Channels int +} + +// PoolStats returns a concurrency-safe pool snapshot. +func (d *DialerManager) PoolStats() PoolStats { + var stats PoolStats + d.mu.RLock() + defer d.mu.RUnlock() + stats.Domains = len(d.domains) + for _, entry := range d.domains { + entry.mu.Lock() + if entry.reconnecting { + stats.Reconnecting++ + } + for _, client := range entry.clients { + stats.Channels += client.active + if client.draining { + stats.Draining++ + } else { + stats.Open++ + } + } + entry.mu.Unlock() + } + return stats +} diff --git a/apps/rlark/pkg/network/sshdialer/dialer_test.go b/apps/rlark/pkg/network/sshdialer/dialer_test.go new file mode 100644 index 0000000..28aa357 --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/dialer_test.go @@ -0,0 +1,957 @@ +package sshdialer + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "sync" + "syscall" + "testing" + "time" + + "golang.org/x/crypto/ssh" +) + +// --------------------------------------------------------------------------- +// DialerManager and domain Dialer tests. +// --------------------------------------------------------------------------- + +const testPrivateKeyPEM = `-----BEGIN OPENSSH PRIVATE KEY----- +b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW +QyNTUxOQAAACDYgEohV8cyTPhXqw3J4KJZ814GmHJAVqXy5IkEH6RBBgAAAKCK3Czsitws +7AAAAAtzc2gtZWQyNTUxOQAAACDYgEohV8cyTPhXqw3J4KJZ814GmHJAVqXy5IkEH6RBBg +AAAEDun/wMJd+XLqbF/nKfrayvmXeLhHjzLd4L+yQ/yFAgD9iASiFXxzJM+FerDcngolnz +XgaYckBWpfLkiQQfpEEGAAAAGmxpZ2h0bmluZ0BDaGVueHVNYWNib29rQWlyAQID +-----END OPENSSH PRIVATE KEY-----` + +const testPublicKeyPEM = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINiASiFXxzJM+FerDcngolnzXgaYckBWpfLkiQQfpEEG" + +// newSSHClient creates a *ssh.Client with a fully initialized transport +// by going through a real SSH handshake with a local test server. +func newSSHClient(t *testing.T) *ssh.Client { + t.Helper() + + signer, err := ssh.ParsePrivateKey([]byte(testPrivateKeyPEM)) + if err != nil { + t.Fatalf("parse server key: %v", err) + } + + serverConfig := &ssh.ServerConfig{ + PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) { + return &ssh.Permissions{}, nil + }, + } + serverConfig.AddHostKey(signer) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + t.Cleanup(func() { _ = ln.Close() }) + + go func() { + tcpConn, err := ln.Accept() + if err != nil { + return + } + _, _, _, err = ssh.NewServerConn(tcpConn, serverConfig) + if err != nil { + _ = tcpConn.Close() + } + }() + + tcpConn, err := net.DialTimeout("tcp", ln.Addr().String(), 5*time.Second) + if err != nil { + t.Fatalf("client dial: %v", err) + } + + clientConfig := &ssh.ClientConfig{ + User: "test", + Auth: []ssh.AuthMethod{ + ssh.PublicKeys(signer), + }, + HostKeyCallback: ssh.InsecureIgnoreHostKey(), + Timeout: 5 * time.Second, + } + + c, _, _, err := ssh.NewClientConn(tcpConn, ln.Addr().String(), clientConfig) + if err != nil { + t.Fatalf("client handshake: %v", err) + } + client := ssh.NewClient(c, nil, nil) + t.Cleanup(func() { _ = client.Close() }) + return client +} + +// testDialer returns a minimal dialer for tests that need to call entry.borrow() directly. +func testDialer(t *testing.T) *DialerManager { + t.Helper() + d := NewDialerManager(Config{ + IdleTimeout: 100 * time.Millisecond, + CleanupInterval: 50 * time.Millisecond, + }) + t.Cleanup(func() { _ = d.Close() }) + return d +} + +func newDomainEntry(id string) *domainEntry { + return &domainEntry{id: id, domainInfo: DomainInfo{ID: id}, generation: 1} +} + +func TestConfig_KeepaliveTimeoutDefault(t *testing.T) { + cfg := Config{} + cfg.setDefaults() + if cfg.KeepaliveTimeout != defaultKeepaliveTimeout { + t.Fatalf("KeepaliveTimeout = %v, want %v", cfg.KeepaliveTimeout, defaultKeepaliveTimeout) + } + if cfg.KeepaliveDrainGrace != defaultKeepaliveDrainGrace { + t.Fatalf("KeepaliveDrainGrace = %v, want %v", cfg.KeepaliveDrainGrace, defaultKeepaliveDrainGrace) + } +} + +func TestDomainEntry_KeepaliveTimeoutDrains(t *testing.T) { + entry := newDomainEntry("test") + pooled := newPooledSSHClient(newSSHClient(t)) + pooled.active = 1 + entry.clients = []*pooledSSHClient{pooled} + + done := make(chan struct{}) + go func() { + entry.keepaliveLoop(pooled, time.Millisecond, 10*time.Millisecond, time.Second) + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("keepalive loop did not time out") + } + + entry.mu.Lock() + if len(entry.clients) != 1 || !pooled.draining { + entry.mu.Unlock() + t.Fatal("timed-out keepalive should drain the active SSH transport") + } + entry.mu.Unlock() + entry.release(pooled) + entry.mu.Lock() + defer entry.mu.Unlock() + if len(entry.clients) != 0 { + t.Fatal("draining transport should close after its last channel exits") + } +} + +func TestDomainEntry_DrainGraceForcesClose(t *testing.T) { + entry := newDomainEntry("test") + pooled := newPooledSSHClient(newSSHClient(t)) + pooled.active = 1 + entry.clients = []*pooledSSHClient{pooled} + + entry.drain(pooled, "test", 10*time.Millisecond) + deadline := time.Now().Add(time.Second) + for { + entry.mu.Lock() + removed := len(entry.clients) == 0 + entry.mu.Unlock() + if removed { + break + } + if time.Now().After(deadline) { + t.Fatal("drain grace did not force-close the transport") + } + time.Sleep(time.Millisecond) + } + + entry.release(pooled) + if pooled.active != 0 { + t.Fatalf("expected delayed release to remain safe, got active=%d", pooled.active) + } +} + +// TestDomainEntry_Borrow 测试正常路径：借用健康的连接。 +func TestDomainEntry_Borrow(t *testing.T) { + entry := newDomainEntry("test") + client := newSSHClient(t) + pooled := newPooledSSHClient(client) + entry.clients = []*pooledSSHClient{pooled} + + // 健康的连接走 fast path，d 不会被使用 + got, err := entry.borrow(context.Background(), nil) + if err != nil { + t.Fatalf("borrow: %v", err) + } + if got != pooled { + t.Fatal("borrow returned wrong client") + } + if pooled.lastUsed().Equal(time.Time{}) { + t.Fatal("expected lastUsed to be updated") + } + entry.release(got) +} + +func TestDomainEntry_AdaptivePoolSelection(t *testing.T) { + d := NewDialerManager(Config{MaxConnectionsPerDomain: 2}) + t.Cleanup(func() { _ = d.Close() }) + entry := newDomainEntry("test") + first := newPooledSSHClient(newSSHClient(t)) + second := newPooledSSHClient(newSSHClient(t)) + first.active = 2 + second.active = 1 + entry.clients = []*pooledSSHClient{first, second} + + got, err := entry.borrow(context.Background(), d) + if err != nil { + t.Fatalf("borrow: %v", err) + } + if got != second { + t.Fatal("expected least-loaded SSH connection") + } + if second.active != 2 { + t.Fatalf("expected active=2, got %d", second.active) + } + entry.release(got) + if second.active != 1 { + t.Fatalf("expected active=1 after release, got %d", second.active) + } +} + +func TestDomainEntry_RoundRobinsEquallyLoadedClients(t *testing.T) { + entry := newDomainEntry("test") + first := newPooledSSHClient(newSSHClient(t)) + second := newPooledSSHClient(newSSHClient(t)) + first.generation = entry.generation + second.generation = entry.generation + entry.clients = []*pooledSSHClient{first, second} + + got, err := entry.borrow(context.Background(), nil) + if err != nil { + t.Fatalf("first borrow: %v", err) + } + if got != first { + t.Fatal("expected first borrow to use first transport") + } + entry.release(got) + + got, err = entry.borrow(context.Background(), nil) + if err != nil { + t.Fatalf("second borrow: %v", err) + } + if got != second { + t.Fatal("expected retry to rotate to equally loaded second transport") + } + entry.release(got) +} + +func TestDomainEntry_LoadTriggersBackgroundExpansion(t *testing.T) { + d := NewDialerManager(Config{ + MaxConnectionsPerDomain: 2, + MaxChannelsPerConnection: 2, + }) + t.Cleanup(func() { _ = d.Close() }) + entry := newDomainEntry("test") + entry.domainInfo = DomainInfo{ID: "test", SSHAddress: "127.0.0.1:1", PrivateKey: testPrivateKeyPEM} + client := newPooledSSHClient(newSSHClient(t)) + client.generation = entry.generation + client.active = 1 + entry.clients = []*pooledSSHClient{client} + + started := time.Now() + got, err := entry.borrow(context.Background(), d) + if err != nil { + t.Fatalf("borrow: %v", err) + } + if got != client { + t.Fatal("expected the existing transport while expansion runs") + } + if elapsed := time.Since(started); elapsed > 100*time.Millisecond { + t.Fatalf("borrow waited for background expansion: %v", elapsed) + } + entry.mu.Lock() + reconnecting := entry.reconnecting + entry.mu.Unlock() + if !reconnecting { + t.Fatal("expected channel load to start background expansion") + } + entry.release(got) +} + +func TestDomainEntry_LoadExpansionRespectsConnectionLimit(t *testing.T) { + d := NewDialerManager(Config{ + MaxConnectionsPerDomain: 1, + MaxChannelsPerConnection: 1, + }) + t.Cleanup(func() { _ = d.Close() }) + entry := newDomainEntry("test") + client := newPooledSSHClient(newSSHClient(t)) + client.generation = entry.generation + entry.clients = []*pooledSSHClient{client} + + got, err := entry.borrow(context.Background(), d) + if err != nil { + t.Fatalf("borrow: %v", err) + } + entry.mu.Lock() + reconnecting := entry.reconnecting + entry.mu.Unlock() + if reconnecting { + t.Fatal("connection limit should prevent background expansion") + } + entry.release(got) +} + +func TestDomainEntry_LowLoadDoesNotExpand(t *testing.T) { + d := NewDialerManager(Config{ + MaxConnectionsPerDomain: 2, + MaxChannelsPerConnection: 2, + }) + t.Cleanup(func() { _ = d.Close() }) + entry := newDomainEntry("test") + client := newPooledSSHClient(newSSHClient(t)) + client.generation = entry.generation + entry.clients = []*pooledSSHClient{client} + + got, err := entry.borrow(context.Background(), d) + if err != nil { + t.Fatalf("borrow: %v", err) + } + entry.mu.Lock() + reconnecting := entry.reconnecting + entry.mu.Unlock() + if reconnecting { + t.Fatal("load below the soft limit should not expand the pool") + } + entry.release(got) +} + +func TestActivityConn_CloseReleasesOnce(t *testing.T) { + left, right := net.Pipe() + t.Cleanup(func() { _ = right.Close() }) + releases := 0 + conn := &activityConn{ + Conn: left, + onActivity: func() {}, + onRelease: func() { releases++ }, + onError: func(error) {}, + } + + _ = conn.Close() + _ = conn.Close() + if releases != 1 { + t.Fatalf("expected one release, got %d", releases) + } +} + +func TestActivityConn_ChannelTimeoutReleasesWithoutTransportError(t *testing.T) { + transportErr := &net.OpError{Op: "write", Err: syscall.ETIMEDOUT} + errorsReported := 0 + releases := 0 + conn := &activityConn{ + Conn: &errorConn{err: transportErr}, + onActivity: func() {}, + onRelease: func() { releases++ }, + onError: func(err error) { + if !errors.Is(err, transportErr) { + t.Errorf("unexpected transport error: %v", err) + } + errorsReported++ + }, + } + + _, _ = conn.Write(nil) + _, _ = conn.Read(nil) + if errorsReported != 0 { + t.Fatalf("expected channel timeout not to abort transport, got %d reports", errorsReported) + } + if releases != 1 { + t.Fatalf("expected terminal channel error to release once, got %d", releases) + } +} + +func TestActivityConn_EOFDoesNotReportTransportError(t *testing.T) { + errorsReported := 0 + conn := &activityConn{ + Conn: &errorConn{err: io.EOF}, + onActivity: func() {}, + onRelease: func() {}, + onError: func(error) { errorsReported++ }, + } + + _, _ = conn.Read(nil) + _, _ = conn.Write(nil) + if errorsReported != 0 { + t.Fatalf("expected channel EOF not to report a transport error, got %d", errorsReported) + } +} + +func TestActivityConn_UnexpectedEOFReportsTransportError(t *testing.T) { + errorsReported := 0 + conn := &activityConn{ + Conn: &errorConn{err: io.ErrUnexpectedEOF}, + onActivity: func() {}, + onRelease: func() {}, + onError: func(error) { errorsReported++ }, + } + + _, _ = conn.Read(nil) + if errorsReported != 1 { + t.Fatalf("expected unexpected EOF to report a transport error, got %d", errorsReported) + } +} + +type errorConn struct { + net.Conn + err error +} + +func (c *errorConn) Read([]byte) (int, error) { return 0, c.err } +func (c *errorConn) Write([]byte) (int, error) { return 0, c.err } + +// TestDomainEntry_BorrowBroken 测试连接损坏后触发重连。 +func TestDomainEntry_BorrowBroken(t *testing.T) { + d := testDialer(t) + entry := newDomainEntry("test") + + // 没有可用连接 → 尝试重连 → 失败 + entry.domainInfo = DomainInfo{ID: "test", SSHAddress: "127.0.0.1:1", PrivateKey: testPrivateKeyPEM} + _, err := entry.borrow(context.Background(), d) + if err == nil { + t.Fatal("expected error when no SSH server") + } + + // 重连失败后不应有 client + entry.mu.Lock() + if len(entry.clients) != 0 { + t.Fatal("expected nil client after failed reconnect") + } + entry.mu.Unlock() +} + +// TestDomainEntry_MarkBroken 测试标记为损坏并关闭连接。 +func TestDomainEntry_MarkBroken(t *testing.T) { + entry := newDomainEntry("test") + client := newSSHClient(t) + pooled := newPooledSSHClient(client) + entry.clients = []*pooledSSHClient{pooled} + + entry.abort(pooled, "test") + if len(entry.clients) != 0 { + t.Fatal("expected transport to be removed after abort") + } +} + +func TestDomainEntry_MarkBrokenClosesActiveTransport(t *testing.T) { + entry := newDomainEntry("test") + pooled := newPooledSSHClient(newSSHClient(t)) + pooled.active = 2 + entry.clients = []*pooledSSHClient{pooled} + + entry.abort(pooled, "test") + if len(entry.clients) != 0 || !pooled.draining { + t.Fatal("broken client should be removed even with active channels") + } + if got := entry.leastLoadedLocked(); got != nil { + t.Fatal("draining client must not accept new channels") + } + + entry.release(pooled) + entry.release(pooled) + if pooled.active != 0 { + t.Fatalf("expected active references to drain, got %d", pooled.active) + } +} + +// TestDialerManager_ConcurrentSafety 高并发下不 panic 不死锁。 +func TestDialerManager_ConcurrentSafety(t *testing.T) { + d := NewDialerManager(Config{ + IdleTimeout: 100 * time.Millisecond, + CleanupInterval: 50 * time.Millisecond, + InitialReconnectBackoff: 1 * time.Millisecond, + MaxReconnectBackoff: 10 * time.Millisecond, + }) + defer func() { _ = d.Close() }() + + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + dialer := d.GetDialer(DomainInfo{ID: "test-domain", SSHAddress: "127.0.0.1:1", PrivateKey: testPrivateKeyPEM}) + _, _ = dialer.DialContext(context.Background(), "tcp", "127.0.0.1:80") + }() + } + wg.Wait() + + time.Sleep(150 * time.Millisecond) + t.Log("50 concurrent dials completed without panic") +} + +// TestDialerManager_ConcurrentReconnect 50 个并发请求失败后都应及时返回。 +func TestDialerManager_ConcurrentReconnect(t *testing.T) { + d := NewDialerManager(Config{ + InitialReconnectBackoff: 1 * time.Millisecond, + MaxReconnectBackoff: 10 * time.Millisecond, + }) + defer func() { _ = d.Close() }() + entry := d.getOrCreate(DomainInfo{ID: "test", SSHAddress: "127.0.0.1:1", PrivateKey: testPrivateKeyPEM}) + + // 模拟连接断开 + entry.close() + + // 50 个并发请求，最多并行建立连接池容量个连接。 + var wg sync.WaitGroup + errCh := make(chan error, 50) + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, err := entry.borrow(context.Background(), d) + errCh <- err + }() + } + wg.Wait() + close(errCh) + + count := 0 + for err := range errCh { + if err == nil { + t.Fatal("all should fail - no SSH server") + } + count++ + } + if count != 50 { + t.Fatalf("expected 50 results, got %d", count) + } + t.Logf("50 concurrent reconnects: all completed, none deadlocked") +} + +// TestDialerManager_ReconnectCoordination verifies concurrent requests share +// one recovery handshake instead of creating a connection storm. +func TestDialerManager_ReconnectCoordination(t *testing.T) { + d := NewDialerManager(Config{ + InitialReconnectBackoff: 1 * time.Millisecond, + MaxReconnectBackoff: 10 * time.Millisecond, + }) + defer func() { _ = d.Close() }() + + // 启动一个真实的 SSH 服务器，慢速握手 + signer, err := ssh.ParsePrivateKey([]byte(testPrivateKeyPEM)) + if err != nil { + t.Fatalf("parse key: %v", err) + } + serverConfig := &ssh.ServerConfig{ + PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) { + return &ssh.Permissions{}, nil + }, + } + serverConfig.AddHostKey(signer) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + t.Cleanup(func() { _ = ln.Close() }) + entry := d.getOrCreate(DomainInfo{ID: "test", SSHAddress: ln.Addr().String(), PrivateKey: testPrivateKeyPEM}) + + go func() { + for { + tcpConn, err := ln.Accept() + if err != nil { + return + } + go func() { + time.Sleep(200 * time.Millisecond) + _, _, _, err := ssh.NewServerConn(tcpConn, serverConfig) + if err != nil { + _ = tcpConn.Close() + } + }() + } + }() + + const requestCount = defaultMaxConnections * 10 + var wg sync.WaitGroup + start := time.Now() + for i := 0; i < requestCount; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, err := entry.borrow(context.Background(), d) + if err != nil { + t.Errorf("borrow failed: %v", err) + } + }() + } + wg.Wait() + + elapsed := time.Since(start) + t.Logf("%d concurrent borrows with 200ms handshake: took %v", requestCount, elapsed) + if elapsed > time.Second { + t.Fatalf("expected parallel connection setup, but took %v", elapsed) + } + if got := d.Stats(); got != 1 { + t.Fatalf("expected one shared recovery connection, got %d", got) + } +} + +// TestDialerManager_Close 测试正确关闭。 +func TestDialerManager_Close(t *testing.T) { + d := NewDialerManager(Config{}) + entry := d.getOrCreate(DomainInfo{ID: "test-domain"}) + entry.clients = []*pooledSSHClient{newPooledSSHClient(newSSHClient(t))} + _ = d.Close() +} + +// TestDialerManager_CloseRacesReconnect 测试 Close 与重连的竞态： +// 重连中调用 Close，不应产生泄漏。 +func TestDialerManager_CloseRacesReconnect(t *testing.T) { + d := NewDialerManager(Config{}) + entry := d.getOrCreate(DomainInfo{ID: "test", SSHAddress: "127.0.0.1:1", PrivateKey: testPrivateKeyPEM}) + + // 在一个 goroutine 中启动重连（但阻塞在退避或拨号中） + done := make(chan struct{}) + go func() { + _, _ = entry.borrow(context.Background(), d) + close(done) + }() + + // 立即 Close + time.Sleep(5 * time.Millisecond) + _ = d.Close() + + // 等待重连返回 + <-done + + // Close 后不应有 client + entry.mu.Lock() + hasClient := len(entry.clients) != 0 + entry.mu.Unlock() + if hasClient { + t.Fatal("expected no client after Close, got leaked connection") + } + + // 再次调用 Close 应安全（幂等） + _ = d.Close() +} + +// TestDialerManager_Stats 测试统计。 +func TestDialerManager_Stats(t *testing.T) { + d := NewDialerManager(Config{}) + defer func() { _ = d.Close() }() + + entry1 := d.getOrCreate(DomainInfo{ID: "a"}) + entry1.clients = []*pooledSSHClient{newPooledSSHClient(newSSHClient(t))} + entry2 := d.getOrCreate(DomainInfo{ID: "b"}) + entry2.clients = []*pooledSSHClient{newPooledSSHClient(newSSHClient(t))} + + if stats := d.Stats(); stats != 2 { + t.Fatalf("expected 2 open, got %d", stats) + } + + entry1.abort(entry1.clients[0], "test") + if stats := d.Stats(); stats != 1 { + t.Fatalf("expected 1 open after broken, got %d", stats) + } +} + +func TestDialerManager_ChannelLimit(t *testing.T) { + d := NewDialerManager(Config{MaxChannelsPerDomain: 1}) + defer func() { _ = d.Close() }() + entry := d.getOrCreate(DomainInfo{ID: "test"}) + client := newPooledSSHClient(newSSHClient(t)) + client.generation = entry.generation + client.active = 1 + entry.clients = []*pooledSSHClient{client} + _, err := entry.borrow(context.Background(), d) + if !errors.Is(err, ErrOverloaded) { + t.Fatalf("borrow error = %v, want ErrOverloaded", err) + } +} + +func TestDialerManager_RemoveDomain(t *testing.T) { + d := NewDialerManager(Config{}) + defer func() { _ = d.Close() }() + dialer := d.GetDialer(DomainInfo{ID: "test", PrivateKey: "secret"}) + d.RemoveDomain("test") + if _, err := dialer.DialContext(context.Background(), "tcp", "localhost:1"); !errors.Is(err, ErrDomainRemoved) { + t.Fatalf("dial error = %v, want ErrDomainRemoved", err) + } + if stats := d.PoolStats(); stats.Domains != 0 { + t.Fatalf("expected domain removal, got %+v", stats) + } +} + +// TestDialerManager_GC 测试空闲连接回收。 +func TestDialerManager_GC(t *testing.T) { + d := NewDialerManager(Config{ + IdleTimeout: 50 * time.Millisecond, + CleanupInterval: 20 * time.Millisecond, + }) + defer func() { _ = d.Close() }() + + entry := d.getOrCreate(DomainInfo{ID: "test-domain"}) + client := newSSHClient(t) + entry.mu.Lock() + pooled := newPooledSSHClient(client) + pooled.lastUsedNanos.Store(time.Now().Add(-1 * time.Hour).UnixNano()) + entry.clients = []*pooledSSHClient{pooled} + entry.mu.Unlock() + + time.Sleep(100 * time.Millisecond) + + entry.mu.Lock() + hasClient := len(entry.clients) != 0 + entry.mu.Unlock() + + if hasClient { + t.Fatal("expected idle connection to be GC'd") + } +} + +// TestDialerManager_GetOrCreate 测试 domain 创建和复用。 +func TestDialerManager_GetOrCreate(t *testing.T) { + d := NewDialerManager(Config{}) + defer func() { _ = d.Close() }() + + e1 := d.getOrCreate(DomainInfo{ID: "test", SSHAddress: "first"}) + e2 := d.getOrCreate(DomainInfo{ID: "test", SSHAddress: "second"}) + e3 := d.getOrCreate(DomainInfo{ID: "other"}) + + if e1 != e2 { + t.Fatal("expected same instance for same ID") + } + if e1 == e3 { + t.Fatal("expected different instance for different ID") + } + if e1.domainInfo.SSHAddress != "second" { + t.Fatal("expected the latest domain info to be used") + } + if e1.generation != 2 { + t.Fatalf("expected changed domain info to advance generation, got %d", e1.generation) + } +} + +func TestDialerManager_ConfigChangeDrainsOldTransport(t *testing.T) { + d := NewDialerManager(Config{KeepaliveDrainGrace: time.Second}) + defer func() { _ = d.Close() }() + entry := d.getOrCreate(DomainInfo{ID: "test", SSHAddress: "first"}) + client := newPooledSSHClient(newSSHClient(t)) + client.generation = entry.generation + client.active = 1 + entry.clients = []*pooledSSHClient{client} + + d.getOrCreate(DomainInfo{ID: "test", SSHAddress: "second"}) + entry.mu.Lock() + defer entry.mu.Unlock() + if !client.draining { + t.Fatal("old-generation transport should drain after endpoint change") + } + if entry.leastLoadedLocked() != nil { + t.Fatal("old-generation transport must not accept new channels") + } +} + +func TestDialerManager_UnchangedConfigKeepsGeneration(t *testing.T) { + d := NewDialerManager(Config{}) + defer func() { _ = d.Close() }() + info := DomainInfo{ID: "test", SSHAddress: "same", PrivateKey: "key"} + entry := d.getOrCreate(info) + generation := entry.generation + d.getOrCreate(info) + if entry.generation != generation { + t.Fatalf("unchanged config advanced generation from %d to %d", generation, entry.generation) + } +} + +func TestBorrowCancellationDoesNotChangeBackoff(t *testing.T) { + d := NewDialerManager(Config{InitialReconnectBackoff: time.Second}) + defer func() { _ = d.Close() }() + entry := d.getOrCreate(DomainInfo{ID: "test", SSHAddress: "127.0.0.1:1", PrivateKey: testPrivateKeyPEM}) + entry.mu.Lock() + entry.reconnectBackoff = time.Second + entry.mu.Unlock() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := entry.borrow(ctx, d) + if !errors.Is(err, context.Canceled) { + t.Fatalf("borrow error = %v, want context canceled", err) + } + entry.mu.Lock() + defer entry.mu.Unlock() + if entry.reconnectBackoff != time.Second { + t.Fatalf("caller cancellation changed shared backoff to %v", entry.reconnectBackoff) + } +} + +func TestDialerManager_GetDialerBindsDomain(t *testing.T) { + manager := NewDialerManager(Config{}) + defer func() { _ = manager.Close() }() + + first := manager.GetDialer(DomainInfo{ID: "test", SSHAddress: "first"}) + second := manager.GetDialer(DomainInfo{ID: "test", SSHAddress: "second"}) + if first.entry != second.entry { + t.Fatal("dialers for the same domain should share a connection pool") + } + if second.entry.domainInfo.SSHAddress != "second" { + t.Fatal("domain dialer should use the latest domain info") + } + if _, err := second.DialContext(context.Background(), "udp", "localhost:1"); err == nil { + t.Fatal("expected unsupported network error") + } +} + +// TestDialSSH_ParseKey 测试密钥解析。 +func TestDialSSH_ParseKey(t *testing.T) { + signer, err := ssh.ParsePrivateKey([]byte(testPrivateKeyPEM)) + if err != nil { + t.Fatalf("parse private key: %v", err) + } + if signer.PublicKey().Type() != ssh.KeyAlgoED25519 { + t.Fatalf("expected ed25519, got %s", signer.PublicKey().Type()) + } +} + +// TestDialSSH_ParsePubkey 测试公钥解析。 +func TestDialSSH_ParsePubkey(t *testing.T) { + _, _, _, _, err := ssh.ParseAuthorizedKey([]byte(testPublicKeyPEM)) + if err != nil { + t.Fatalf("parse authorized key: %v", err) + } +} + +func TestDialSSH_HandshakeTimeout(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + t.Cleanup(func() { _ = ln.Close() }) + accepted := make(chan net.Conn, 1) + go func() { + conn, acceptErr := ln.Accept() + if acceptErr == nil { + accepted <- conn + } + }() + + cfg := Config{SSHTimeout: 25 * time.Millisecond, HostKeyCallback: ssh.InsecureIgnoreHostKey()} + started := time.Now() + _, err = dialSSH(context.Background(), ln.Addr().String(), "", testPrivateKeyPEM, cfg) + if err == nil { + t.Fatal("expected stalled SSH handshake to time out") + } + if elapsed := time.Since(started); elapsed > time.Second { + t.Fatalf("SSH handshake exceeded configured timeout: %v", elapsed) + } + select { + case conn := <-accepted: + _ = conn.Close() + default: + } +} + +// TestNextBackoff 测试退避计算。 +func TestNextBackoff(t *testing.T) { + tests := []struct { + current time.Duration + max time.Duration + expected time.Duration + }{ + {0, maxReconnectBackoff, initialReconnectBackoff}, + {initialReconnectBackoff, maxReconnectBackoff, initialReconnectBackoff * 2}, + {maxReconnectBackoff / 2, maxReconnectBackoff, maxReconnectBackoff}, + {maxReconnectBackoff, maxReconnectBackoff, maxReconnectBackoff}, + {maxReconnectBackoff * 2, maxReconnectBackoff, maxReconnectBackoff}, + } + for _, tt := range tests { + got := nextBackoff(tt.current, initialReconnectBackoff, tt.max) + if got != tt.expected { + t.Errorf("nextBackoff(%v, %v, %v) = %v, want %v", tt.current, initialReconnectBackoff, tt.max, got, tt.expected) + } + } +} + +func TestNextBackoffUsesConfiguredInitial(t *testing.T) { + initial := 25 * time.Millisecond + if got := nextBackoff(0, initial, time.Second); got != initial { + t.Fatalf("nextBackoff initial = %v, want %v", got, initial) + } +} + +// TestIsSSHTransportError 覆盖各类传输错误与误判场景。 +func TestIsSSHTransportError(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"io.EOF", io.EOF, true}, + {"io.ErrUnexpectedEOF", io.ErrUnexpectedEOF, true}, + {"ECONNRESET", syscall.ECONNRESET, true}, + {"EPIPE", syscall.EPIPE, true}, + {"ETIMEDOUT", syscall.ETIMEDOUT, false}, + {"wrapped ECONNRESET", fmt.Errorf("dial: %w", syscall.ECONNRESET), true}, + {"net.OpError reset", &net.OpError{Op: "read", Err: syscall.ECONNRESET}, true}, + {"net.OpError timeout", &net.OpError{Op: "read", Err: syscall.ETIMEDOUT}, false}, + {"ssh transport closed", errors.New("ssh: tcp transport closed"), true}, + {"context.Canceled", context.Canceled, false}, + {"context.DeadlineExceeded", context.DeadlineExceeded, false}, + {"wrapped context.Canceled", fmt.Errorf("dial: %w", context.Canceled), false}, + {"wrapped context.DeadlineExceeded", fmt.Errorf("dial: %w", context.DeadlineExceeded), false}, + {"generic error", errors.New("connection refused"), false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isSSHTransportError(tt.err); got != tt.want { + t.Errorf("isSSHTransportError(%v) = %v, want %v", tt.err, got, tt.want) + } + }) + } +} + +// TestDomainEntry_AbortIgnoresRemovedTransport verifies stale maintenance work +// cannot abort a replacement transport. +func TestDomainEntry_AbortIgnoresRemovedTransport(t *testing.T) { + entry := newDomainEntry("test") + + old := newPooledSSHClient(newSSHClient(t)) + entry.clients = []*pooledSSHClient{old} + + // Replace the old transport as if reconnect had succeeded. + newClient := newPooledSSHClient(newSSHClient(t)) + entry.clients = []*pooledSSHClient{newClient} + + entry.abort(old, "test") + if len(entry.clients) != 1 || entry.clients[0] != newClient { + t.Fatal("replacement transport should remain untouched") + } + + entry.abort(newClient, "test") + if len(entry.clients) != 0 { + t.Fatal("expected the current transport to be aborted") + } +} + +// TestParseSSHAddr 测试 user@host:port 解析。 +func TestParseSSHAddr(t *testing.T) { + tests := []struct { + addr string + defaultUser string + wantUser string + wantHost string + }{ + {"root@192.168.1.1:22", "admin", "root", "192.168.1.1:22"}, + {"app@10.0.0.5:2222", "root", "app", "10.0.0.5:2222"}, + {"192.168.1.1:22", "root", "root", "192.168.1.1:22"}, + {"10.0.0.5:2222", "admin", "admin", "10.0.0.5:2222"}, + {"user@host:0", "", "user", "host:0"}, + {"@host:22", "root", "", "host:22"}, + } + for _, tt := range tests { + gotUser, gotHost := parseSSHAddr(tt.addr, tt.defaultUser) + if gotUser != tt.wantUser || gotHost != tt.wantHost { + t.Errorf("parseSSHAddr(%q, %q) = (%q, %q), want (%q, %q)", + tt.addr, tt.defaultUser, gotUser, gotHost, tt.wantUser, tt.wantHost) + } + } +} diff --git a/apps/rlark/pkg/network/sshdialer/maintenance.go b/apps/rlark/pkg/network/sshdialer/maintenance.go new file mode 100644 index 0000000..7e88505 --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/maintenance.go @@ -0,0 +1,127 @@ +package sshdialer + +import ( + "fmt" + "time" + + "github.com/rlinf/rlark/apps/rlark/pkg/log" +) + +func (entry *domainEntry) keepaliveLoop(client *pooledSSHClient, interval, timeout, drainGrace time.Duration) { + logger := log.GetLogger() + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + result := make(chan error, 1) + go func() { + _, _, err := client.client.SendRequest(keepaliveRequest, true, nil) + result <- err + }() + timer := time.NewTimer(timeout) + select { + case err := <-result: + if !timer.Stop() { + <-timer.C + } + if err == nil { + continue + } + logger.Info("SSH keepalive failed, aborting transport", + "domain", entry.id, + "err", err, + "errType", fmt.Sprintf("%T", err), + ) + entry.abort(client, "keepalive-failed") + return + case <-timer.C: + logger.Info("SSH keepalive timed out, draining connection", + "domain", entry.id, + "timeout", timeout, + ) + entry.drain(client, "keepalive-timeout", drainGrace) + return + case <-client.keepaliveDone: + if !timer.Stop() { + <-timer.C + } + return + } + case <-client.keepaliveDone: + return + } + } +} + +func (entry *domainEntry) close() { + entry.mu.Lock() + clients := entry.clients + entry.clients = nil + if entry.reconnectCancel != nil { + entry.reconnectCancel() + } + for _, client := range clients { + client.close() + } + entry.mu.Unlock() + for _, client := range clients { + select { + case <-client.keepaliveExit: + default: + // Test-created clients may not have a maintenance goroutine. + } + } +} + +func (d *DialerManager) cleanupLoop() { + defer d.wg.Done() + ticker := time.NewTicker(d.cfg.CleanupInterval) + defer ticker.Stop() + for { + select { + case <-d.ctx.Done(): + return + case <-ticker.C: + d.cleanup() + } + } +} + +func (d *DialerManager) cleanup() { + cutoff := time.Now().Add(-d.cfg.IdleTimeout) + d.mu.RLock() + entries := make([]*domainEntry, 0, len(d.domains)) + for _, entry := range d.domains { + entries = append(entries, entry) + } + d.mu.RUnlock() + + for _, entry := range entries { + entry.mu.Lock() + kept := entry.clients[:0] + for _, client := range entry.clients { + if client.active == 0 && (client.lastUsed().Before(cutoff) || time.Since(client.createdAt) >= d.cfg.MaxConnectionAge) { + client.close() + continue + } + if !client.draining && time.Since(client.createdAt) >= d.cfg.MaxConnectionAge { + client.draining = true + } + kept = append(kept, client) + } + entry.clients = kept + removeEntry := len(entry.clients) == 0 && !entry.reconnecting && entry.lastUsed().Before(cutoff) + entry.mu.Unlock() + if removeEntry { + d.mu.Lock() + if d.domains[entry.id] == entry { + delete(d.domains, entry.id) + entry.mu.Lock() + entry.domainInfo = DomainInfo{ID: entry.id} + entry.mu.Unlock() + } + d.mu.Unlock() + } + } +} diff --git a/apps/rlark/pkg/network/sshdialer/pool.go b/apps/rlark/pkg/network/sshdialer/pool.go new file mode 100644 index 0000000..9c0c47e --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/pool.go @@ -0,0 +1,333 @@ +package sshdialer + +import ( + "context" + "errors" + "fmt" + "math/rand" + "time" + + "github.com/rlinf/rlark/apps/rlark/pkg/log" +) + +func (d *DialerManager) getOrCreate(info DomainInfo) *domainEntry { + d.mu.RLock() + entry, ok := d.domains[info.ID] + d.mu.RUnlock() + if !ok { + d.mu.Lock() + entry, ok = d.domains[info.ID] + if !ok { + entry = &domainEntry{ + id: info.ID, + domainInfo: info, + generation: 1, + initialBackoff: d.cfg.InitialReconnectBackoff, + maxBackoff: d.cfg.MaxReconnectBackoff, + } + entry.touch() + d.domains[info.ID] = entry + } + d.mu.Unlock() + } + + entry.mu.Lock() + entry.touch() + if entry.domainInfo != info { + entry.domainInfo = info + entry.generation++ + entry.reconnectBackoff = 0 + if entry.reconnectCancel != nil { + entry.reconnectCancel() + } + clients := append([]*pooledSSHClient(nil), entry.clients...) + entry.mu.Unlock() + for _, client := range clients { + entry.drain(client, "domain-config-changed", d.cfg.KeepaliveDrainGrace) + } + return entry + } + entry.mu.Unlock() + return entry +} + +func (entry *domainEntry) borrow(ctx context.Context, d *DialerManager) (*pooledSSHClient, error) { + for { + entry.mu.Lock() + if entry.removed { + entry.mu.Unlock() + return nil, ErrDomainRemoved + } + if d != nil && entry.activeChannelsLocked() >= d.cfg.MaxChannelsPerDomain { + entry.mu.Unlock() + return nil, ErrOverloaded + } + client := entry.borrowCandidateLocked() + if client != nil { + client.active++ + client.touch() + entry.touch() + if d != nil && client.active >= d.cfg.MaxChannelsPerConnection && len(entry.clients) < d.cfg.MaxConnectionsPerDomain && !entry.reconnecting { + entry.startReconnectLocked(d) + } + entry.mu.Unlock() + return client, nil + } + if d == nil { + entry.mu.Unlock() + return nil, fmt.Errorf("no SSH transport available") + } + if !entry.reconnecting { + entry.startReconnectLocked(d) + ch := entry.reconnectCh + entry.mu.Unlock() + if err := entry.waitForReconnect(ctx, d, ch); err != nil { + return nil, err + } + } else { + ch := entry.reconnectCh + entry.mu.Unlock() + if err := entry.waitForReconnect(ctx, d, ch); err != nil { + return nil, err + } + } + } +} + +func (entry *domainEntry) startReconnectLocked(d *DialerManager) { + entry.reconnecting = true + entry.reconnectCh = make(chan struct{}) + reconnectCtx, reconnectCancel := context.WithCancel(d.ctx) + entry.reconnectCancel = reconnectCancel + info := entry.domainInfo + generation := entry.generation + backoff := entry.reconnectBackoff + d.wg.Add(1) + go entry.reconnect(reconnectCtx, d, info, generation, backoff) +} + +func (entry *domainEntry) waitForReconnect(ctx context.Context, d *DialerManager, ch <-chan struct{}) error { + select { + case <-ch: + entry.mu.Lock() + err := entry.lastReconnectErr + entry.mu.Unlock() + if err != nil { + return fmt.Errorf("ssh reconnect: %w", err) + } + return nil + case <-ctx.Done(): + return ctx.Err() + case <-d.ctx.Done(): + return ErrClosed + } +} + +func (entry *domainEntry) activeChannelsLocked() int { + active := 0 + for _, client := range entry.clients { + active += client.active + } + return active +} + +func (entry *domainEntry) leastLoadedLocked() *pooledSSHClient { + return entry.selectLeastLoadedLocked(false) +} + +func (entry *domainEntry) borrowCandidateLocked() *pooledSSHClient { + return entry.selectLeastLoadedLocked(true) +} + +func (entry *domainEntry) selectLeastLoadedLocked(advance bool) *pooledSSHClient { + if len(entry.clients) == 0 { + return nil + } + + var selected *pooledSSHClient + selectedIndex := -1 + for offset := range entry.clients { + index := (entry.next + offset) % len(entry.clients) + client := entry.clients[index] + if client.draining || (client.generation != 0 && client.generation != entry.generation) { + continue + } + if selected == nil || client.active < selected.active { + selected = client + selectedIndex = index + } + } + if advance && selectedIndex >= 0 { + entry.next = (selectedIndex + 1) % len(entry.clients) + } + return selected +} + +func (entry *domainEntry) release(client *pooledSSHClient) { + entry.mu.Lock() + if client.active > 0 { + client.active-- + } + client.touch() + entry.touch() + if client.draining && client.active == 0 { + entry.removeLocked(client) + client.close() + } + entry.mu.Unlock() +} + +func (entry *domainEntry) reconnect(ctx context.Context, d *DialerManager, info DomainInfo, generation uint64, backoff time.Duration) { + defer d.wg.Done() + var err error + if backoff > 0 { + wait := backoff + time.Duration(rand.Int63n(max(1, int64(backoff/2)))) + timer := time.NewTimer(wait) + select { + case <-timer.C: + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + err = ErrClosed + } + } + + var client *pooledSSHClient + if err == nil { + var sshClientErr error + sshClient, dialErr := d.dialSSHWithMergedCtx(ctx, info) + sshClientErr = dialErr + err = sshClientErr + if err == nil { + client = newPooledSSHClient(sshClient) + client.generation = generation + } + } + + entry.mu.Lock() + stale := generation != entry.generation + if client != nil && !stale && !d.closed.Load() { + entry.clients = append(entry.clients, client) + entry.reconnectBackoff = 0 + entry.lastReconnectErr = nil + } else { + if client != nil { + client.close() + } + if stale { + err = nil + } else if !errors.Is(err, ErrClosed) { + entry.reconnectBackoff = nextBackoff(entry.reconnectBackoff, entry.initialBackoff, entry.maxBackoff) + } + entry.lastReconnectErr = err + } + entry.reconnecting = false + entry.reconnectCancel = nil + ch := entry.reconnectCh + entry.reconnectCh = nil + entry.mu.Unlock() + if ch != nil { + close(ch) + } + + if client != nil && !stale && !d.closed.Load() { + go func() { + defer close(client.keepaliveExit) + entry.keepaliveLoop(client, d.cfg.KeepaliveInterval, d.cfg.KeepaliveTimeout, d.cfg.KeepaliveDrainGrace) + }() + if d.cfg.OnReconnect != nil { + d.cfg.OnReconnect(entry.id) + } + } +} + +func nextBackoff(current, initial, maxBackoff time.Duration) time.Duration { + if initial <= 0 { + initial = initialReconnectBackoff + } + if maxBackoff <= 0 { + maxBackoff = maxReconnectBackoff + } + if current <= 0 { + return min(initial, maxBackoff) + } + return min(current*2, maxBackoff) +} + +func (entry *domainEntry) abort(client *pooledSSHClient, reason string) { + entry.mu.Lock() + defer entry.mu.Unlock() + if !entry.containsLocked(client) || client.draining { + return + } + client.draining = true + log.GetLogger().Info("SSH transport aborted", + "domain", entry.id, + "reason", reason, + "activeChannels", client.active, + "lastUsed", client.lastUsed(), + "idleFor", time.Since(client.lastUsed()).Round(time.Second), + ) + entry.removeLocked(client) + client.close() +} + +func (entry *domainEntry) drain(client *pooledSSHClient, reason string, grace time.Duration) { + entry.mu.Lock() + if !entry.containsLocked(client) || client.draining { + entry.mu.Unlock() + return + } + client.draining = true + log.GetLogger().Info("SSH connection draining", + "domain", entry.id, + "reason", reason, + "activeChannels", client.active, + "grace", grace, + ) + if client.active == 0 { + entry.removeLocked(client) + client.close() + entry.mu.Unlock() + return + } + entry.mu.Unlock() + + entry.mu.Lock() + client.drainTimer = time.AfterFunc(grace, func() { + entry.mu.Lock() + defer entry.mu.Unlock() + if entry.containsLocked(client) { + entry.removeLocked(client) + client.close() + } + }) + entry.mu.Unlock() +} + +func (entry *domainEntry) containsLocked(client *pooledSSHClient) bool { + for _, candidate := range entry.clients { + if candidate == client { + return true + } + } + return false +} + +func (entry *domainEntry) removeLocked(client *pooledSSHClient) { + for i, candidate := range entry.clients { + if candidate == client { + entry.clients = append(entry.clients[:i], entry.clients[i+1:]...) + if len(entry.clients) == 0 { + entry.next = 0 + } else { + if i < entry.next { + entry.next-- + } + entry.next %= len(entry.clients) + } + return + } + } +} diff --git a/apps/rlark/pkg/network/sshdialer/ssh.go b/apps/rlark/pkg/network/sshdialer/ssh.go new file mode 100644 index 0000000..32aed26 --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/ssh.go @@ -0,0 +1,116 @@ +package sshdialer + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "strings" + "syscall" + "time" + + "github.com/rlinf/rlark/apps/rlark/pkg/auth/cert" + "golang.org/x/crypto/ssh" +) + +func dialSSH(ctx context.Context, sshAddr, certPEM, keyPEM string, cfg Config) (*ssh.Client, error) { + user, address := parseSSHAddr(sshAddr, cfg.SSHUser) + signer, err := ssh.ParsePrivateKey([]byte(keyPEM)) + if err != nil { + return nil, fmt.Errorf("parse ssh key: %w", err) + } + + var auth ssh.AuthMethod + if certPEM != "" { + certificate, err := cert.DecodeSSHCertificateFromPEM([]byte(certPEM)) + if err != nil { + return nil, fmt.Errorf("parse ssh cert: %w", err) + } + certSigner, err := ssh.NewCertSigner(certificate, signer) + if err != nil { + return nil, fmt.Errorf("new cert signer: %w", err) + } + auth = ssh.PublicKeys(certSigner) + } else { + auth = ssh.PublicKeys(signer) + } + + config := &ssh.ClientConfig{ + User: user, + Auth: []ssh.AuthMethod{auth}, + HostKeyCallback: cfg.HostKeyCallback, + Timeout: cfg.SSHTimeout, + } + dialer := &net.Dialer{Timeout: cfg.SSHTimeout, KeepAlive: 30 * time.Second} + conn, err := dialer.DialContext(ctx, "tcp", address) + if err != nil { + return nil, fmt.Errorf("dial ssh server %s: %w", address, err) + } + handshakeDeadline := time.Now().Add(cfg.SSHTimeout) + if deadline, ok := ctx.Deadline(); ok && deadline.Before(handshakeDeadline) { + handshakeDeadline = deadline + } + if err := conn.SetDeadline(handshakeDeadline); err != nil { + _ = conn.Close() + return nil, fmt.Errorf("set SSH handshake deadline: %w", err) + } + handshakeDone := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + _ = conn.Close() + case <-handshakeDone: + } + }() + + c, chans, reqs, err := ssh.NewClientConn(conn, address, config) + close(handshakeDone) + if err != nil { + _ = conn.Close() + return nil, fmt.Errorf("ssh handshake with %s: %w", address, err) + } + if err := conn.SetDeadline(time.Time{}); err != nil { + _ = c.Close() + return nil, fmt.Errorf("clear SSH handshake deadline: %w", err) + } + return ssh.NewClient(c, chans, reqs), nil +} + +func parseSSHAddr(addr, defaultUser string) (string, string) { + user, hostPort, ok := strings.Cut(addr, "@") + if ok { + return user, hostPort + } + return defaultUser, addr +} + +func isSSHTransportError(err error) bool { + if err == nil || errors.Is(err, context.Canceled) { + return false + } + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return true + } + if errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.EPIPE) { + return true + } + return strings.Contains(err.Error(), "ssh: tcp transport closed") +} + +func isSSHChannelTransportError(err error) bool { + return !errors.Is(err, io.EOF) && isSSHTransportError(err) +} + +func (d *DialerManager) dialSSHWithMergedCtx(ctx context.Context, info DomainInfo) (*ssh.Client, error) { + dialCtx, cancel := context.WithCancel(ctx) + defer cancel() + go func() { + select { + case <-d.ctx.Done(): + cancel() + case <-dialCtx.Done(): + } + }() + return dialSSH(dialCtx, info.SSHAddress, info.Certificate, info.PrivateKey, d.cfg) +} diff --git a/apps/rlark/pkg/network/sshdialer/tests/README.md b/apps/rlark/pkg/network/sshdialer/tests/README.md new file mode 100644 index 0000000..665c1b7 --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/tests/README.md @@ -0,0 +1,27 @@ +# SSH Dialer Conformance Tests + +These long-running integration scenarios exercise the SSH dialer with a real +SSH server and TCP echo backend. They are intentionally excluded from regular +`go test` runs. + +```bash +go run ./pkg/network/sshdialer/tests/conformance +``` + +The suite covers: + +- concurrent multiplexed channels with sustained traffic; +- forced SSH transport interruption and automatic reconnection; +- SSH service shutdown, same-address restart, and recovery without restarting + the dialer manager. + +Use flags to increase or reduce the workload: + +```bash +go run ./pkg/network/sshdialer/tests/conformance \ + -connections 64 \ + -payload-size 8388608 \ + -rounds 5 \ + -duration 1m \ + -timeout 60s +``` diff --git a/apps/rlark/pkg/network/sshdialer/tests/conformance/main.go b/apps/rlark/pkg/network/sshdialer/tests/conformance/main.go new file mode 100644 index 0000000..6e58d48 --- /dev/null +++ b/apps/rlark/pkg/network/sshdialer/tests/conformance/main.go @@ -0,0 +1,464 @@ +// Run with: +// +// go run ./pkg/network/sshdialer/tests/conformance +package main + +import ( + "context" + "crypto/rand" + "flag" + "fmt" + "io" + "log" + "net" + "os" + "sync" + "sync/atomic" + "time" + + "github.com/rlinf/rlark/apps/rlark/pkg/network/sshdialer" + "golang.org/x/crypto/ssh" +) + +const domainID = "conformance-domain" + +type options struct { + connections int + payloadSize int + rounds int + duration time.Duration + timeout time.Duration +} + +func main() { + var opts options + flag.IntVar(&opts.connections, "connections", 32, "number of concurrent tunneled connections") + flag.IntVar(&opts.payloadSize, "payload-size", 4<<20, "bytes transferred by each connection") + flag.IntVar(&opts.rounds, "rounds", 3, "concurrent rounds per scenario") + flag.DurationVar(&opts.duration, "duration", 10*time.Second, "minimum duration of the sustained traffic scenario") + flag.DurationVar(&opts.timeout, "timeout", 30*time.Second, "timeout for each connection") + flag.Parse() + + scenarios := []struct { + name string + run func(context.Context, options) error + }{ + {"high concurrency and sustained traffic", scenarioConcurrentTraffic}, + {"transport interruption and automatic reconnect", scenarioTransportInterruption}, + {"SSH service restart and recovery", scenarioServerRestart}, + } + + failures := 0 + for _, scenario := range scenarios { + fmt.Printf("== %s\n", scenario.name) + started := time.Now() + if err := scenario.run(context.Background(), opts); err != nil { + failures++ + fmt.Printf(" FAIL (%s): %v\n", time.Since(started).Round(time.Millisecond), err) + continue + } + fmt.Printf(" ok (%s)\n", time.Since(started).Round(time.Millisecond)) + } + if failures > 0 { + fmt.Printf("\n%d scenario(s) failed\n", failures) + os.Exit(1) + } + fmt.Println("\nall scenarios passed") +} + +func scenarioConcurrentTraffic(ctx context.Context, opts options) error { + env, err := startEnvironment(opts) + if err != nil { + return err + } + defer env.close() + + deadline := time.Now().Add(opts.duration) + for round := 1; round <= opts.rounds || time.Now().Before(deadline); round++ { + if err := runConcurrentRound(ctx, env.dialer, env.backend.Addr().String(), opts); err != nil { + return fmt.Errorf("round %d: %w", round, err) + } + } + if open := env.manager.Stats(); open < 1 || open > 4 { + return fmt.Errorf("unexpected available transport count: %d", open) + } + return nil +} + +func scenarioTransportInterruption(ctx context.Context, opts options) error { + env, err := startEnvironment(opts) + if err != nil { + return err + } + defer env.close() + + if err := roundTrip(ctx, env.dialer, env.backend.Addr().String(), opts.payloadSize, opts.timeout); err != nil { + return fmt.Errorf("warm-up: %w", err) + } + before := env.server.accepted.Load() + env.server.dropTransports() + + if err := retry(opts.timeout, func() error { + return roundTrip(ctx, env.dialer, env.backend.Addr().String(), opts.payloadSize, opts.timeout) + }); err != nil { + return fmt.Errorf("recover after transport interruption: %w", err) + } + if env.server.accepted.Load() <= before { + return fmt.Errorf("dialer recovered without establishing a replacement transport") + } + return runConcurrentRound(ctx, env.dialer, env.backend.Addr().String(), opts) +} + +func scenarioServerRestart(ctx context.Context, opts options) error { + backend, err := startEchoServer() + if err != nil { + return err + } + defer func() { _ = backend.Close() }() + + server, err := newSSHTestServer() + if err != nil { + return err + } + address, err := server.start("127.0.0.1:0") + if err != nil { + return err + } + defer server.close() + + manager, dialer := newManager(address, opts) + defer func() { _ = manager.Close() }() + if err := roundTrip(ctx, dialer, backend.Addr().String(), opts.payloadSize, opts.timeout); err != nil { + return fmt.Errorf("before restart: %w", err) + } + + server.stop() + failCtx, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + conn, dialErr := dialer.DialContext(failCtx, "tcp", backend.Addr().String()) + cancel() + if dialErr == nil { + _ = conn.Close() + return fmt.Errorf("expected dial to fail while SSH service is stopped") + } + + if _, err := server.start(address); err != nil { + return fmt.Errorf("restart SSH service: %w", err) + } + if err := retry(opts.timeout, func() error { + return roundTrip(ctx, dialer, backend.Addr().String(), opts.payloadSize, opts.timeout) + }); err != nil { + return fmt.Errorf("recover after SSH service restart: %w", err) + } + return runConcurrentRound(ctx, dialer, backend.Addr().String(), opts) +} + +type testEnvironment struct { + backend net.Listener + server *sshTestServer + manager *sshdialer.DialerManager + dialer *sshdialer.Dialer +} + +func startEnvironment(opts options) (*testEnvironment, error) { + backend, err := startEchoServer() + if err != nil { + return nil, err + } + server, err := newSSHTestServer() + if err != nil { + _ = backend.Close() + return nil, err + } + address, err := server.start("127.0.0.1:0") + if err != nil { + _ = backend.Close() + return nil, err + } + manager, dialer := newManager(address, opts) + return &testEnvironment{backend: backend, server: server, manager: manager, dialer: dialer}, nil +} + +func (e *testEnvironment) close() { + _ = e.manager.Close() + e.server.close() + _ = e.backend.Close() +} + +func newManager(address string, opts options) (*sshdialer.DialerManager, *sshdialer.Dialer) { + manager := sshdialer.NewDialerManager(sshdialer.Config{ + SSHTimeout: opts.timeout, + InitialReconnectBackoff: 10 * time.Millisecond, + MaxReconnectBackoff: 250 * time.Millisecond, + KeepaliveInterval: 100 * time.Millisecond, + KeepaliveTimeout: 100 * time.Millisecond, + KeepaliveDrainGrace: 500 * time.Millisecond, + MaxConnectionsPerDomain: 4, + }) + dialer := manager.GetDialer(sshdialer.DomainInfo{ + ID: domainID, + SSHAddress: address, + PrivateKey: testPrivateKeyPEM, + }) + return manager, dialer +} + +func runConcurrentRound(ctx context.Context, dialer *sshdialer.Dialer, address string, opts options) error { + errs := make(chan error, opts.connections) + var wg sync.WaitGroup + for i := 0; i < opts.connections; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + if err := roundTrip(ctx, dialer, address, opts.payloadSize, opts.timeout); err != nil { + errs <- fmt.Errorf("connection %d: %w", id, err) + } + }(i) + } + wg.Wait() + close(errs) + for err := range errs { + return err + } + return nil +} + +func roundTrip(parent context.Context, dialer *sshdialer.Dialer, address string, size int, timeout time.Duration) error { + ctx, cancel := context.WithTimeout(parent, timeout) + defer cancel() + conn, err := dialer.DialContext(ctx, "tcp", address) + if err != nil { + return fmt.Errorf("dial: %w", err) + } + defer func() { _ = conn.Close() }() + _ = conn.SetDeadline(time.Now().Add(timeout)) + + payload := make([]byte, size) + if _, err := rand.Read(payload); err != nil { + return err + } + writeDone := make(chan error, 1) + go func() { + _, err := conn.Write(payload) + writeDone <- err + }() + response := make([]byte, len(payload)) + if _, err := io.ReadFull(conn, response); err != nil { + return fmt.Errorf("read: %w", err) + } + if err := <-writeDone; err != nil { + return fmt.Errorf("write: %w", err) + } + for i := range payload { + if payload[i] != response[i] { + return fmt.Errorf("payload mismatch at byte %d", i) + } + } + return nil +} + +func retry(timeout time.Duration, fn func() error) error { + deadline := time.Now().Add(timeout) + var lastErr error + for time.Now().Before(deadline) { + if err := fn(); err == nil { + return nil + } else { + lastErr = err + } + time.Sleep(50 * time.Millisecond) + } + return lastErr +} + +type sshTestServer struct { + config *ssh.ServerConfig + + mu sync.Mutex + listener net.Listener + conns map[*ssh.ServerConn]struct{} + accepted atomic.Int64 + wg sync.WaitGroup +} + +func newSSHTestServer() (*sshTestServer, error) { + signer, err := ssh.ParsePrivateKey([]byte(testPrivateKeyPEM)) + if err != nil { + return nil, err + } + config := &ssh.ServerConfig{PublicKeyCallback: func(ssh.ConnMetadata, ssh.PublicKey) (*ssh.Permissions, error) { + return &ssh.Permissions{}, nil + }} + config.AddHostKey(signer) + return &sshTestServer{config: config, conns: make(map[*ssh.ServerConn]struct{})}, nil +} + +func (s *sshTestServer) start(address string) (string, error) { + ln, err := net.Listen("tcp", address) + if err != nil { + return "", err + } + s.mu.Lock() + s.listener = ln + s.mu.Unlock() + s.wg.Add(1) + go s.serve(ln) + return ln.Addr().String(), nil +} + +func (s *sshTestServer) serve(listener net.Listener) { + defer s.wg.Done() + for { + conn, err := listener.Accept() + if err != nil { + return + } + s.wg.Add(1) + go s.serveConn(conn) + } +} + +func (s *sshTestServer) serveConn(conn net.Conn) { + defer s.wg.Done() + serverConn, channels, requests, err := ssh.NewServerConn(conn, s.config) + if err != nil { + _ = conn.Close() + return + } + s.mu.Lock() + s.conns[serverConn] = struct{}{} + s.mu.Unlock() + s.accepted.Add(1) + defer func() { + s.mu.Lock() + delete(s.conns, serverConn) + s.mu.Unlock() + _ = serverConn.Close() + }() + go ssh.DiscardRequests(requests) + for channel := range channels { + if channel.ChannelType() != "direct-tcpip" { + _ = channel.Reject(ssh.UnknownChannelType, "unsupported channel type") + continue + } + var request directTCPIPRequest + if err := ssh.Unmarshal(channel.ExtraData(), &request); err != nil { + _ = channel.Reject(ssh.ConnectionFailed, "invalid direct-tcpip request") + continue + } + upstream, err := net.DialTimeout("tcp", net.JoinHostPort(request.Host, fmt.Sprint(request.Port)), 5*time.Second) + if err != nil { + _ = channel.Reject(ssh.ConnectionFailed, err.Error()) + continue + } + sshChannel, channelRequests, err := channel.Accept() + if err != nil { + _ = upstream.Close() + continue + } + go ssh.DiscardRequests(channelRequests) + go proxy(sshChannel, upstream) + } +} + +type directTCPIPRequest struct { + Host string + Port uint32 + OriginHost string + OriginPort uint32 +} + +func proxy(left io.ReadWriteCloser, right net.Conn) { + defer func() { _ = left.Close() }() + defer func() { _ = right.Close() }() + done := make(chan struct{}, 2) + go func() { + _, _ = io.Copy(left, right) + done <- struct{}{} + }() + go func() { + _, _ = io.Copy(right, left) + done <- struct{}{} + }() + <-done +} + +func (s *sshTestServer) dropTransports() { + s.mu.Lock() + conns := make([]*ssh.ServerConn, 0, len(s.conns)) + for conn := range s.conns { + conns = append(conns, conn) + } + s.mu.Unlock() + for _, conn := range conns { + _ = conn.Close() + } +} + +func (s *sshTestServer) stop() { + s.mu.Lock() + listener := s.listener + s.listener = nil + s.mu.Unlock() + if listener != nil { + _ = listener.Close() + } + s.dropTransports() + s.wg.Wait() +} + +func (s *sshTestServer) close() { + s.stop() +} + +func startEchoServer() (net.Listener, error) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, err + } + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go echo(conn) + } + }() + return listener, nil +} + +func echo(conn net.Conn) { + defer func() { _ = conn.Close() }() + pending := make(chan []byte, 128) + go func() { + for data := range pending { + if _, err := conn.Write(data); err != nil { + return + } + } + }() + buffer := make([]byte, 64<<10) + for { + n, err := conn.Read(buffer) + if n > 0 { + data := append([]byte(nil), buffer[:n]...) + pending <- data + } + if err != nil { + close(pending) + return + } + } +} + +func init() { + log.SetFlags(log.LstdFlags | log.Lmicroseconds) +} + +const testPrivateKeyPEM = `-----BEGIN OPENSSH PRIVATE KEY----- +b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW +QyNTUxOQAAACDYgEohV8cyTPhXqw3J4KJZ814GmHJAVqXy5IkEH6RBBgAAAKCK3Czsitws +7AAAAAtzc2gtZWQyNTUxOQAAACDYgEohV8cyTPhXqw3J4KJZ814GmHJAVqXy5IkEH6RBBg +AAAEDun/wMJd+XLqbF/nKfrayvmXeLhHjzLd4L+yQ/yFAgD9iASiFXxzJM+FerDcngolnz +XgaYckBWpfLkiQQfpEEGAAAAGmxpZ2h0bmluZ0BDaGVueHVNYWNib29rQWlyAQID +-----END OPENSSH PRIVATE KEY-----` diff --git a/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install.go b/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install.go index 1d6c385..f0cadb0 100644 --- a/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install.go +++ b/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install.go @@ -89,7 +89,7 @@ func (d *Installer) Install(cfg *types.DeployConfig, certBundle *cert.Bundle) er for _, c := range component.ComponentsForPlane(cfg) { c.HealthCheckFn = health.K8sWorkloadHealthCheck(clientset, c) if c.Name == constants.ComponentKCP { - c.PostDeployFn = extractKCPKubeconfigFn(ctx, clientset, kubeconfig) + c.PostDeployFn = initializeKCPManagementAPIFn(ctx, clientset, kubeconfig) } if err := ensureRBAC(ctx, clientset, cfg, &c); err != nil { return err @@ -119,16 +119,6 @@ func (d *Installer) Install(cfg *types.DeployConfig, certBundle *cert.Bundle) er } } - // 始终 apply CRD，无论 KCP 是否新部署 - if cfg.Plane == types.PlaneControl && !cfg.UsesKubernetesManagementAPI() { - if err := createUIAuthSecretInKCP(ctx, clientset, kubeconfig); err != nil { - return err - } - if err := applyKCP(ctx, clientset, kubeconfig); err != nil { - return err - } - } - logger.Info("plane deployed", "plane", cfg.Plane, "namespace", constants.Namespace) d.summary = d.buildSummary(ctx, clientset, cfg) @@ -226,17 +216,35 @@ func installCRDsToKubernetes(ctx context.Context, client apiextensionsclient.Int func ensureUIAuthSecretInKubernetes(ctx context.Context, clientset kubernetes.Interface, namespace string) error { if secret, err := clientset.CoreV1().Secrets(namespace).Get(ctx, common.UIAuthSecretName, metav1.GetOptions{}); err == nil { - if len(secret.Data[common.UIAuthJWTSigningKey]) >= 32 { + if len(secret.Data[common.UIAuthAdminPasswordKey]) > 0 && + len(secret.Data[common.UIAuthUserPasswordKey]) > 0 && + len(secret.Data[common.UIAuthJWTSigningKey]) >= 32 { return nil } - signingKey, err := generateSecret(32) - if err != nil { - return fmt.Errorf("generate JWT signing key: %w", err) - } if secret.Data == nil { secret.Data = make(map[string][]byte) } - secret.Data[common.UIAuthJWTSigningKey] = signingKey + if len(secret.Data[common.UIAuthAdminPasswordKey]) == 0 { + password, err := generatePassword(16) + if err != nil { + return fmt.Errorf("generate admin password: %w", err) + } + secret.Data[common.UIAuthAdminPasswordKey] = []byte(password) + } + if len(secret.Data[common.UIAuthUserPasswordKey]) == 0 { + password, err := generatePassword(16) + if err != nil { + return fmt.Errorf("generate user password: %w", err) + } + secret.Data[common.UIAuthUserPasswordKey] = []byte(password) + } + if len(secret.Data[common.UIAuthJWTSigningKey]) < 32 { + signingKey, err := generateSecret(32) + if err != nil { + return fmt.Errorf("generate JWT signing key: %w", err) + } + secret.Data[common.UIAuthJWTSigningKey] = signingKey + } if _, err := clientset.CoreV1().Secrets(namespace).Update(ctx, secret, metav1.UpdateOptions{}); err != nil { return fmt.Errorf("update ui auth secret: %w", err) } @@ -359,18 +367,17 @@ func ensureRBAC(ctx context.Context, clientset *kubernetes.Clientset, cfg *types return nil } -// extractKCPKubeconfigFn returns a PostDeployFn that extracts admin.kubeconfig -// from the KCP pod and creates a ConfigMap for other components to mount. -func extractKCPKubeconfigFn(ctx context.Context, clientset *kubernetes.Clientset, kubeconfig string) func(cfg *types.DeployConfig) error { +// initializeKCPManagementAPIFn initializes resources required by components +// immediately after kcp is healthy and before those components are deployed. +func initializeKCPManagementAPIFn(ctx context.Context, clientset *kubernetes.Clientset, kubeconfig string) func(cfg *types.DeployConfig) error { return func(cfg *types.DeployConfig) error { - return extractAndApplyKCP(ctx, clientset, kubeconfig) + if err := extractAndApplyKCP(ctx, clientset, kubeconfig); err != nil { + return err + } + return createUIAuthSecretInKCP(ctx, clientset, kubeconfig) } } -func applyKCP(ctx context.Context, clientset *kubernetes.Clientset, kubeconfig string) error { - return extractAndApplyKCP(ctx, clientset, kubeconfig) -} - func extractAndApplyKCP(ctx context.Context, clientset *kubernetes.Clientset, kubeconfig string) error { logger := log.GetLogger() pods, err := clientset.CoreV1().Pods(constants.Namespace).List(ctx, metav1.ListOptions{ diff --git a/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install_test.go b/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install_test.go index 3fe0c31..04e1de2 100644 --- a/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install_test.go +++ b/apps/rlark/pkg/rlarkadm/deployer/kubernetes/install_test.go @@ -90,3 +90,28 @@ func TestEnsureUIAuthSecretInKubernetes(t *testing.T) { t.Fatal("existing UI auth secret was changed") } } + +func TestEnsureUIAuthSecretInKubernetesRepairsMissingFields(t *testing.T) { + ctx := context.Background() + signingKey := []byte("01234567890123456789012345678901") + client := kubernetesfake.NewSimpleClientset(&corev1.Secret{ + ObjectMeta: metav1.ObjectMeta{Name: common.UIAuthSecretName, Namespace: constants.Namespace}, + Data: map[string][]byte{ + common.UIAuthJWTSigningKey: signingKey, + }, + }) + + if err := ensureUIAuthSecretInKubernetes(ctx, client, constants.Namespace); err != nil { + t.Fatalf("ensureUIAuthSecretInKubernetes() error = %v", err) + } + secret, err := client.CoreV1().Secrets(constants.Namespace).Get(ctx, common.UIAuthSecretName, metav1.GetOptions{}) + if err != nil { + t.Fatalf("get repaired UI auth secret: %v", err) + } + if len(secret.Data[common.UIAuthAdminPasswordKey]) != 16 || len(secret.Data[common.UIAuthUserPasswordKey]) != 16 { + t.Fatalf("missing passwords were not repaired: %#v", secret.Data) + } + if string(secret.Data[common.UIAuthJWTSigningKey]) != string(signingKey) { + t.Fatal("valid JWT signing key was changed") + } +}