From 421848e3c0950286c75ac9aa61c8d15ae8d237e5 Mon Sep 17 00:00:00 2001 From: pionxe Date: Thu, 16 Apr 2026 00:04:07 +0800 Subject: [PATCH 1/4] =?UTF-8?q?feat:=20[EPIC-GW-04]=20=E6=90=AD=E5=BB=BA?= =?UTF-8?q?=E7=BD=91=E7=BB=9C=E8=AE=BF=E9=97=AE=E9=9D=A2=E9=AA=A8=E6=9E=B6?= =?UTF-8?q?=20(HTTP/WS/SSE)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 12 + internal/cli/gateway_commands.go | 58 ++- internal/cli/root_test.go | 79 ++- internal/gateway/network_server.go | 626 +++++++++++++++++++++++ internal/gateway/network_server_test.go | 314 ++++++++++++ internal/gateway/rpc_dispatch.go | 59 +++ internal/gateway/server.go | 46 +- internal/tui/infra/clipboard_common.go | 8 +- internal/tui/infra/clipboard_fallback.go | 4 - internal/tui/infra/clipboard_native.go | 2 +- internal/tui/infra/image.go | 18 +- 11 files changed, 1161 insertions(+), 65 deletions(-) create mode 100644 internal/gateway/network_server.go create mode 100644 internal/gateway/network_server_test.go create mode 100644 internal/gateway/rpc_dispatch.go diff --git a/README.md b/README.md index f1089cd4..18a52d44 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,18 @@ Gateway 子命令（Step 1 骨架）： go run ./cmd/neocode gateway ``` +指定网络访问面监听地址（默认 `127.0.0.1:8080`，仅允许 Loopback）： + +```bash +go run ./cmd/neocode gateway --http-listen 127.0.0.1:8080 +``` + +网络访问面骨架端点（EPIC-GW-04）： + +- `POST /rpc`：单次 JSON-RPC 请求入口 +- `GET /ws`：WebSocket 流式入口（含心跳） +- `GET /sse`：SSE 流式入口（MVP 默认触发 `gateway.ping`，含心跳） + URL Scheme 派发骨架命令（EPIC-GW-02A）： ```bash diff --git a/internal/cli/gateway_commands.go b/internal/cli/gateway_commands.go index 3b777e45..1003bc81 100644 --- a/internal/cli/gateway_commands.go +++ b/internal/cli/gateway_commands.go @@ -27,6 +27,7 @@ var ( runGatewayCommand = defaultGatewayCommandRunner runURLDispatchCommand = defaultURLDispatchCommandRunner newGatewayServer = defaultNewGatewayServer + newGatewayNetwork = defaultNewGatewayNetworkServer dispatchURLThroughIPC = urlscheme.Dispatch exitProcess = os.Exit writeDispatchError = writeURLDispatchErrorOutput @@ -35,6 +36,7 @@ var ( type gatewayCommandOptions struct { ListenAddress string + HTTPAddress string LogLevel string } @@ -63,6 +65,12 @@ type gatewayServer interface { Close(ctx context.Context) error } +type gatewayNetworkServer interface { + ListenAddress() string + Serve(ctx context.Context, runtimePort gateway.RuntimePort) error + Close(ctx context.Context) error +} + // newGatewayCommand 创建并返回网关子命令，负责启动本地 Gateway 进程。 func newGatewayCommand() *cobra.Command { options := &gatewayCommandOptions{} @@ -80,12 +88,19 @@ func newGatewayCommand() *cobra.Command { return runGatewayCommand(cmd.Context(), gatewayCommandOptions{ ListenAddress: strings.TrimSpace(options.ListenAddress), + HTTPAddress: strings.TrimSpace(options.HTTPAddress), LogLevel: normalizedLogLevel, }) }, } cmd.Flags().StringVar(&options.ListenAddress, "listen", "", "gateway listen address (optional override)") + cmd.Flags().StringVar( + &options.HTTPAddress, + "http-listen", + gateway.DefaultNetworkListenAddress, + "gateway network listen address (loopback only)", + ) cmd.Flags().StringVar(&options.LogLevel, "log-level", defaultGatewayLogLevel, "gateway log level: debug|info|warn|error") return cmd @@ -110,19 +125,49 @@ func defaultGatewayCommandRunner(ctx context.Context, options gatewayCommandOpti signalContext, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM) defer stop() - server, err := newGatewayServer(gateway.ServerOptions{ + ipcServer, err := newGatewayServer(gateway.ServerOptions{ ListenAddress: options.ListenAddress, Logger: logger, }) if err != nil { return err } + networkServer, err := newGatewayNetwork(gateway.NetworkServerOptions{ + ListenAddress: options.HTTPAddress, + Logger: logger, + }) + if err != nil { + _ = ipcServer.Close(context.Background()) + return err + } defer func() { - _ = server.Close(context.Background()) + _ = networkServer.Close(context.Background()) + _ = ipcServer.Close(context.Background()) }() - logger.Printf("gateway listen address: %s", server.ListenAddress()) - return server.Serve(signalContext, nil) + logger.Printf("gateway ipc listen address: %s", ipcServer.ListenAddress()) + logger.Printf("gateway network listen address: %s", networkServer.ListenAddress()) + + serveContext, cancel := context.WithCancel(signalContext) + defer cancel() + + serveErrorChannel := make(chan error, 2) + go func() { + serveErrorChannel <- ipcServer.Serve(serveContext, nil) + }() + go func() { + serveErrorChannel <- networkServer.Serve(serveContext, nil) + }() + + var firstError error + for index := 0; index < 2; index++ { + serveError := <-serveErrorChannel + if serveError != nil && firstError == nil { + firstError = serveError + cancel() + } + } + return firstError } // defaultNewGatewayServer 创建默认网关服务实例，供命令层启动流程调用。 @@ -130,6 +175,11 @@ func defaultNewGatewayServer(options gateway.ServerOptions) (gatewayServer, erro return gateway.NewServer(options) } +// defaultNewGatewayNetworkServer 创建默认网关网络访问面服务实例，供命令层启动流程调用。 +func defaultNewGatewayNetworkServer(options gateway.NetworkServerOptions) (gatewayNetworkServer, error) { + return gateway.NewNetworkServer(options) +} + // newURLDispatchCommand 创建 URL Scheme 派发子命令骨架，仅做参数收敛与调用转发。 func newURLDispatchCommand() *cobra.Command { options := &urlDispatchCommandOptions{} diff --git a/internal/cli/root_test.go b/internal/cli/root_test.go index 50dbc562..dc65c18a 100644 --- a/internal/cli/root_test.go +++ b/internal/cli/root_test.go @@ -189,7 +189,12 @@ func TestGatewaySubcommandPassesFlagsToRunner(t *testing.T) { } command := NewRootCommand() - command.SetArgs([]string{"gateway", "--listen", " /tmp/gateway.sock ", "--log-level", " WARN "}) + command.SetArgs([]string{ + "gateway", + "--listen", " /tmp/gateway.sock ", + "--http-listen", " 127.0.0.1:19080 ", + "--log-level", " WARN ", + }) if err := command.ExecuteContext(context.Background()); err != nil { t.Fatalf("ExecuteContext() error = %v", err) } @@ -200,6 +205,9 @@ func TestGatewaySubcommandPassesFlagsToRunner(t *testing.T) { if captured.LogLevel != "warn" { t.Fatalf("log level = %q, want %q", captured.LogLevel, "warn") } + if captured.HTTPAddress != "127.0.0.1:19080" { + t.Fatalf("http address = %q, want %q", captured.HTTPAddress, "127.0.0.1:19080") + } } func TestGatewaySubcommandRejectsInvalidLogLevel(t *testing.T) { @@ -220,15 +228,22 @@ func TestGatewaySubcommandRejectsInvalidLogLevel(t *testing.T) { func TestDefaultGatewayCommandRunnerSuccess(t *testing.T) { originalNewGatewayServer := newGatewayServer + originalNewGatewayNetwork := newGatewayNetwork t.Cleanup(func() { newGatewayServer = originalNewGatewayServer }) + t.Cleanup(func() { newGatewayNetwork = originalNewGatewayNetwork }) server := &stubGatewayServer{listenAddress: "stub://gateway"} newGatewayServer = func(options gateway.ServerOptions) (gatewayServer, error) { return server, nil } + networkServer := &stubGatewayServer{listenAddress: "127.0.0.1:8080"} + newGatewayNetwork = func(options gateway.NetworkServerOptions) (gatewayNetworkServer, error) { + return networkServer, nil + } err := defaultGatewayCommandRunner(context.Background(), gatewayCommandOptions{ ListenAddress: "stub://gateway", + HTTPAddress: "127.0.0.1:8080", LogLevel: "info", }) if err != nil { @@ -240,19 +255,31 @@ func TestDefaultGatewayCommandRunnerSuccess(t *testing.T) { if !server.closeCalled { t.Fatal("expected server Close to be called") } + if !networkServer.serveCalled { + t.Fatal("expected network server Serve to be called") + } + if !networkServer.closeCalled { + t.Fatal("expected network server Close to be called") + } } func TestDefaultGatewayCommandRunnerReturnsConstructorError(t *testing.T) { originalNewGatewayServer := newGatewayServer + originalNewGatewayNetwork := newGatewayNetwork t.Cleanup(func() { newGatewayServer = originalNewGatewayServer }) + t.Cleanup(func() { newGatewayNetwork = originalNewGatewayNetwork }) expected := errors.New("new gateway server failed") newGatewayServer = func(options gateway.ServerOptions) (gatewayServer, error) { return nil, expected } + newGatewayNetwork = func(options gateway.NetworkServerOptions) (gatewayNetworkServer, error) { + return &stubGatewayServer{listenAddress: "127.0.0.1:8080"}, nil + } err := defaultGatewayCommandRunner(context.Background(), gatewayCommandOptions{ ListenAddress: "stub://gateway", + HTTPAddress: "127.0.0.1:8080", LogLevel: "info", }) if !errors.Is(err, expected) { @@ -262,7 +289,9 @@ func TestDefaultGatewayCommandRunnerReturnsConstructorError(t *testing.T) { func TestDefaultGatewayCommandRunnerReturnsServeError(t *testing.T) { originalNewGatewayServer := newGatewayServer + originalNewGatewayNetwork := newGatewayNetwork t.Cleanup(func() { newGatewayServer = originalNewGatewayServer }) + t.Cleanup(func() { newGatewayNetwork = originalNewGatewayNetwork }) expected := errors.New("serve failed") server := &stubGatewayServer{ @@ -272,9 +301,14 @@ func TestDefaultGatewayCommandRunnerReturnsServeError(t *testing.T) { newGatewayServer = func(options gateway.ServerOptions) (gatewayServer, error) { return server, nil } + networkServer := &stubGatewayServer{listenAddress: "127.0.0.1:8080"} + newGatewayNetwork = func(options gateway.NetworkServerOptions) (gatewayNetworkServer, error) { + return networkServer, nil + } err := defaultGatewayCommandRunner(context.Background(), gatewayCommandOptions{ ListenAddress: "stub://gateway", + HTTPAddress: "127.0.0.1:8080", LogLevel: "info", }) if !errors.Is(err, expected) { @@ -283,6 +317,37 @@ func TestDefaultGatewayCommandRunnerReturnsServeError(t *testing.T) { if !server.closeCalled { t.Fatal("expected server Close to be called") } + if !networkServer.closeCalled { + t.Fatal("expected network server Close to be called") + } +} + +func TestDefaultGatewayCommandRunnerReturnsNetworkConstructorError(t *testing.T) { + originalNewGatewayServer := newGatewayServer + originalNewGatewayNetwork := newGatewayNetwork + t.Cleanup(func() { newGatewayServer = originalNewGatewayServer }) + t.Cleanup(func() { newGatewayNetwork = originalNewGatewayNetwork }) + + networkErr := errors.New("new network server failed") + ipcServer := &stubGatewayServer{listenAddress: "stub://gateway"} + newGatewayServer = func(options gateway.ServerOptions) (gatewayServer, error) { + return ipcServer, nil + } + newGatewayNetwork = func(options gateway.NetworkServerOptions) (gatewayNetworkServer, error) { + return nil, networkErr + } + + err := defaultGatewayCommandRunner(context.Background(), gatewayCommandOptions{ + ListenAddress: "stub://gateway", + HTTPAddress: "127.0.0.1:8080", + LogLevel: "info", + }) + if !errors.Is(err, networkErr) { + t.Fatalf("expected network constructor error %v, got %v", networkErr, err) + } + if !ipcServer.closeCalled { + t.Fatal("expected ipc server Close to be called when network constructor fails") + } } func TestDefaultNewGatewayServer(t *testing.T) { @@ -297,6 +362,18 @@ func TestDefaultNewGatewayServer(t *testing.T) { } } +func TestDefaultNewGatewayNetworkServer(t *testing.T) { + server, err := defaultNewGatewayNetworkServer(gateway.NetworkServerOptions{ + ListenAddress: "127.0.0.1:8080", + }) + if err != nil { + t.Fatalf("defaultNewGatewayNetworkServer() error = %v", err) + } + if server == nil { + t.Fatal("defaultNewGatewayNetworkServer() returned nil server") + } +} + func TestURLDispatchSubcommandUsesURLFlag(t *testing.T) { originalRunner := runURLDispatchCommand originalPreload := runGlobalPreload diff --git a/internal/gateway/network_server.go b/internal/gateway/network_server.go new file mode 100644 index 00000000..410ad1db --- /dev/null +++ b/internal/gateway/network_server.go @@ -0,0 +1,626 @@ +package gateway + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net" + "net/http" + "os" + "strconv" + "strings" + "sync" + "time" + + "golang.org/x/net/websocket" + + "neo-code/internal/gateway/protocol" +) + +const ( + // DefaultNetworkListenAddress 定义网关网络访问面的默认监听地址，仅允许本机环回访问。 + DefaultNetworkListenAddress = "127.0.0.1:8080" + // DefaultNetworkReadTimeout 定义网络入口单次读取超时时间，防止慢连接长期占用资源。 + DefaultNetworkReadTimeout = 15 * time.Second + // DefaultNetworkWriteTimeout 定义网络入口单次写入超时时间，避免写阻塞导致协程泄漏。 + DefaultNetworkWriteTimeout = 15 * time.Second + // DefaultNetworkShutdownTimeout 定义网络入口优雅关闭的最大等待时间。 + DefaultNetworkShutdownTimeout = 2 * time.Second + // DefaultNetworkHeartbeatInterval 定义 WS/SSE 长连接的保活心跳周期。 + DefaultNetworkHeartbeatInterval = 3 * time.Second + // DefaultNetworkMaxRequestBytes 定义 HTTP/WS 单次请求体的最大字节数。 + DefaultNetworkMaxRequestBytes int64 = MaxFrameSize + // DefaultNetworkMaxStreamConnections 定义 WS/SSE 长连接总上限。 + DefaultNetworkMaxStreamConnections = 128 +) + +var ( + resolveNetworkListenAddressFn = ResolveNetworkListenAddress +) + +// NetworkServerOptions 描述网关网络访问面服务启动所需的可选配置。 +type NetworkServerOptions struct { + ListenAddress string + Logger *log.Logger + ReadTimeout time.Duration + WriteTimeout time.Duration + ShutdownTimeout time.Duration + HeartbeatInterval time.Duration + MaxRequestBytes int64 + MaxStreamConnections int + listenFn func(network, address string) (net.Listener, error) +} + +// NetworkServer 提供 HTTP/WebSocket/SSE 网络访问面的统一入口服务。 +type NetworkServer struct { + listenAddress string + logger *log.Logger + readTimeout time.Duration + writeTimeout time.Duration + shutdownTimeout time.Duration + heartbeatInterval time.Duration + maxRequestBytes int64 + maxStreamConnections int + listenFn func(network, address string) (net.Listener, error) + + mu sync.Mutex + server *http.Server + listener net.Listener + wsConns map[*websocket.Conn]struct{} + sseCancels map[int]context.CancelFunc + nextSSEID int +} + +// NewNetworkServer 创建网关网络访问面服务实例，并执行监听地址合法性校验。 +func NewNetworkServer(options NetworkServerOptions) (*NetworkServer, error) { + listenAddress, err := resolveNetworkListenAddressFn(options.ListenAddress) + if err != nil { + return nil, err + } + + logger := options.Logger + if logger == nil { + logger = log.New(os.Stderr, "gateway-network: ", log.LstdFlags) + } + + listenFn := options.listenFn + if listenFn == nil { + listenFn = net.Listen + } + + readTimeout := options.ReadTimeout + if readTimeout <= 0 { + readTimeout = DefaultNetworkReadTimeout + } + + writeTimeout := options.WriteTimeout + if writeTimeout <= 0 { + writeTimeout = DefaultNetworkWriteTimeout + } + + shutdownTimeout := options.ShutdownTimeout + if shutdownTimeout <= 0 { + shutdownTimeout = DefaultNetworkShutdownTimeout + } + + heartbeatInterval := options.HeartbeatInterval + if heartbeatInterval <= 0 { + heartbeatInterval = DefaultNetworkHeartbeatInterval + } + + maxRequestBytes := options.MaxRequestBytes + if maxRequestBytes <= 0 { + maxRequestBytes = DefaultNetworkMaxRequestBytes + } + + maxStreamConnections := options.MaxStreamConnections + if maxStreamConnections <= 0 { + maxStreamConnections = DefaultNetworkMaxStreamConnections + } + + return &NetworkServer{ + listenAddress: listenAddress, + logger: logger, + readTimeout: readTimeout, + writeTimeout: writeTimeout, + shutdownTimeout: shutdownTimeout, + heartbeatInterval: heartbeatInterval, + maxRequestBytes: maxRequestBytes, + maxStreamConnections: maxStreamConnections, + listenFn: listenFn, + wsConns: make(map[*websocket.Conn]struct{}), + sseCancels: make(map[int]context.CancelFunc), + }, nil +} + +// ResolveNetworkListenAddress 解析网络访问面监听地址，默认值固定为本机环回地址。 +func ResolveNetworkListenAddress(override string) (string, error) { + address := strings.TrimSpace(override) + if address == "" { + address = DefaultNetworkListenAddress + } + if err := validateLoopbackListenAddress(address); err != nil { + return "", err + } + return address, nil +} + +// validateLoopbackListenAddress 校验网络监听地址只能绑定到环回接口，避免开放到外网。 +func validateLoopbackListenAddress(address string) error { + host, _, err := net.SplitHostPort(strings.TrimSpace(address)) + if err != nil { + return fmt.Errorf("invalid --http-listen %q: %w", address, err) + } + normalizedHost := strings.TrimSpace(host) + if normalizedHost == "" { + return fmt.Errorf("invalid --http-listen %q: host must be loopback", address) + } + if strings.EqualFold(normalizedHost, "localhost") { + return nil + } + + ip := net.ParseIP(normalizedHost) + if ip == nil || !ip.IsLoopback() { + return fmt.Errorf("invalid --http-listen %q: host must be loopback", address) + } + return nil +} + +// ListenAddress 返回网络访问面当前绑定的监听地址。 +func (s *NetworkServer) ListenAddress() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.listenAddress +} + +// Serve 启动网络访问面服务，并注册 HTTP/WebSocket/SSE 三类入口。 +func (s *NetworkServer) Serve(ctx context.Context, runtimePort RuntimePort) error { + listener, err := s.listenFn("tcp", s.listenAddress) + if err != nil { + return fmt.Errorf("gateway network listen failed: %w", err) + } + + httpServer := &http.Server{ + Handler: s.withCORS(s.buildHandler(runtimePort)), + ReadTimeout: s.readTimeout, + WriteTimeout: s.writeTimeout, + } + + s.mu.Lock() + if s.server != nil { + s.mu.Unlock() + _ = listener.Close() + return fmt.Errorf("gateway: network server is already serving") + } + s.server = httpServer + s.listener = listener + s.listenAddress = listener.Addr().String() + s.mu.Unlock() + + s.logger.Printf("network listening on %s", listener.Addr().String()) + + go func() { + <-ctx.Done() + _ = s.Close(context.Background()) + }() + + if err := httpServer.Serve(listener); err != nil { + if errors.Is(err, http.ErrServerClosed) || ctx.Err() != nil || s.isClosed() { + return nil + } + return fmt.Errorf("gateway: serve network: %w", err) + } + return nil +} + +// Close 关闭网络访问面并主动中断 WS/SSE 长连接，避免进程退出被长连接阻塞。 +func (s *NetworkServer) Close(ctx context.Context) error { + s.mu.Lock() + httpServer := s.server + listener := s.listener + s.server = nil + s.listener = nil + s.mu.Unlock() + + if httpServer == nil && listener == nil { + return nil + } + + s.forceCloseStreamConnections() + + var closeErr error + if httpServer != nil { + shutdownCtx := context.Background() + if ctx != nil { + shutdownCtx = ctx + } + if s.shutdownTimeout > 0 { + var cancel context.CancelFunc + shutdownCtx, cancel = context.WithTimeout(shutdownCtx, s.shutdownTimeout) + defer cancel() + } + + if err := httpServer.Shutdown(shutdownCtx); err != nil { + closeErr = errors.Join(closeErr, err) + closeErr = errors.Join(closeErr, httpServer.Close()) + } + } + + if listener != nil { + if err := listener.Close(); err != nil && !errors.Is(err, net.ErrClosed) { + closeErr = errors.Join(closeErr, err) + } + } + + return closeErr +} + +// isClosed 判断网络服务是否已经处于关闭状态。 +func (s *NetworkServer) isClosed() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.server == nil +} + +// buildHandler 构建网络访问面的路由入口，并将请求统一转入网关分发链路。 +func (s *NetworkServer) buildHandler(runtimePort RuntimePort) http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("/rpc", func(writer http.ResponseWriter, request *http.Request) { + s.handleRPCRequest(writer, request, runtimePort) + }) + mux.Handle("/ws", websocket.Handler(func(conn *websocket.Conn) { + s.handleWebSocket(conn, runtimePort) + })) + mux.HandleFunc("/sse", func(writer http.ResponseWriter, request *http.Request) { + s.handleSSERequest(writer, request, runtimePort) + }) + return mux +} + +// withCORS 为网络入口统一注入基础 CORS 响应头，并快速处理 OPTIONS 预检请求。 +func (s *NetworkServer) withCORS(next http.Handler) http.Handler { + return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Access-Control-Allow-Origin", "*") + writer.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS") + writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization") + if request.Method == http.MethodOptions { + writer.WriteHeader(http.StatusNoContent) + return + } + next.ServeHTTP(writer, request) + }) +} + +// handleRPCRequest 处理 POST /rpc 请求并返回单次 JSON-RPC 响应。 +func (s *NetworkServer) handleRPCRequest(writer http.ResponseWriter, request *http.Request, runtimePort RuntimePort) { + if request.Method != http.MethodPost { + http.Error(writer, "method not allowed", http.StatusMethodNotAllowed) + return + } + + request.Body = http.MaxBytesReader(writer, request.Body, s.maxRequestBytes) + rpcRequest, rpcErr := decodeJSONRPCRequestFromReader(request.Body) + if rpcErr != nil { + writeJSONRPCHTTPResponse(writer, protocol.NewJSONRPCErrorResponse(nil, rpcErr)) + return + } + + response := dispatchRPCRequest(request.Context(), rpcRequest, runtimePort) + writeJSONRPCHTTPResponse(writer, response) +} + +// handleWebSocket 处理 WS 入口请求，并将每条文本消息作为 JSON-RPC 请求进行分发。 +func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort RuntimePort) { + if !s.registerWSConnection(conn) { + _ = conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)) + _ = websocket.Message.Send(conn, `{"status":"error","code":"too_many_connections","message":"stream connection limit exceeded"}`) + _ = conn.Close() + return + } + defer func() { + s.unregisterWSConnection(conn) + _ = conn.Close() + }() + + maxPayloadBytes := int(s.maxRequestBytes) + if maxPayloadBytes > 0 { + conn.MaxPayloadBytes = maxPayloadBytes + } + + var writeMu sync.Mutex + stopHeartbeat := make(chan struct{}) + defer close(stopHeartbeat) + + go s.runWSHeartbeatLoop(conn, &writeMu, stopHeartbeat) + + for { + if s.readTimeout > 0 { + _ = conn.SetReadDeadline(time.Now().Add(s.readTimeout)) + } + + var rawMessage string + if err := websocket.Message.Receive(conn, &rawMessage); err != nil { + if isConnectionClosedError(err) { + return + } + s.logger.Printf("websocket read failed: %v", err) + return + } + + rpcRequest, rpcErr := decodeJSONRPCRequestFromBytes([]byte(rawMessage)) + var rpcResponse protocol.JSONRPCResponse + if rpcErr != nil { + rpcResponse = protocol.NewJSONRPCErrorResponse(nil, rpcErr) + } else { + rpcResponse = dispatchRPCRequest(context.Background(), rpcRequest, runtimePort) + } + + if err := s.writeWebSocketMessage(conn, &writeMu, rpcResponse); err != nil { + if !isConnectionClosedError(err) { + s.logger.Printf("websocket write failed: %v", err) + } + return + } + } +} + +// runWSHeartbeatLoop 周期性推送 WebSocket 心跳帧，保证长连接可观测与保活。 +func (s *NetworkServer) runWSHeartbeatLoop(conn *websocket.Conn, writeMu *sync.Mutex, stop <-chan struct{}) { + ticker := time.NewTicker(s.heartbeatInterval) + defer ticker.Stop() + + for { + select { + case <-stop: + return + case <-ticker.C: + if err := s.writeWebSocketMessage(conn, writeMu, map[string]any{ + "type": "heartbeat", + "timestamp": time.Now().UTC().Format(time.RFC3339Nano), + }); err != nil { + return + } + } + } +} + +// writeWebSocketMessage 将任意结构序列化后写入 WS 连接，并施加写超时约束。 +func (s *NetworkServer) writeWebSocketMessage(conn *websocket.Conn, writeMu *sync.Mutex, payload any) error { + if s.writeTimeout > 0 { + if err := conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)); err != nil { + return err + } + } + rawPayload, err := json.Marshal(payload) + if err != nil { + return err + } + + writeMu.Lock() + defer writeMu.Unlock() + return websocket.Message.Send(conn, string(rawPayload)) +} + +// handleSSERequest 处理 SSE 入口请求，先返回一次结果事件，再持续发送心跳事件。 +func (s *NetworkServer) handleSSERequest(writer http.ResponseWriter, request *http.Request, runtimePort RuntimePort) { + if request.Method != http.MethodGet { + http.Error(writer, "method not allowed", http.StatusMethodNotAllowed) + return + } + + flusher, ok := writer.(http.Flusher) + if !ok { + http.Error(writer, "streaming not supported", http.StatusInternalServerError) + return + } + + streamCtx, cancel := context.WithCancel(request.Context()) + connectionID, registered := s.registerSSEConnection(cancel) + if !registered { + cancel() + http.Error(writer, "stream connection limit exceeded", http.StatusServiceUnavailable) + return + } + defer func() { + cancel() + s.unregisterSSEConnection(connectionID) + }() + + writer.Header().Set("Content-Type", "text/event-stream") + writer.Header().Set("Cache-Control", "no-cache") + writer.Header().Set("Connection", "keep-alive") + flusher.Flush() + + rpcRequest := buildSSETriggerRequest(request) + rpcResponse := dispatchRPCRequest(streamCtx, rpcRequest, runtimePort) + if err := s.writeSSEEvent(writer, flusher, "result", rpcResponse); err != nil { + return + } + + ticker := time.NewTicker(s.heartbeatInterval) + defer ticker.Stop() + + for { + select { + case <-streamCtx.Done(): + return + case <-ticker.C: + if err := s.writeSSEEvent(writer, flusher, "heartbeat", map[string]string{ + "timestamp": time.Now().UTC().Format(time.RFC3339Nano), + }); err != nil { + return + } + } + } +} + +// writeSSEEvent 将结构化数据写入 SSE 事件通道，并在每次发送后立即刷新。 +func (s *NetworkServer) writeSSEEvent(writer http.ResponseWriter, flusher http.Flusher, eventName string, payload any) error { + if s.writeTimeout > 0 { + _ = http.NewResponseController(writer).SetWriteDeadline(time.Now().Add(s.writeTimeout)) + } + rawPayload, err := json.Marshal(payload) + if err != nil { + return err + } + if _, err := fmt.Fprintf(writer, "event: %s\n", strings.TrimSpace(eventName)); err != nil { + return err + } + if _, err := fmt.Fprintf(writer, "data: %s\n\n", string(rawPayload)); err != nil { + return err + } + flusher.Flush() + return nil +} + +// buildSSETriggerRequest 从 SSE 查询参数构建一次 JSON-RPC 触发请求，默认方法为 gateway.ping。 +func buildSSETriggerRequest(request *http.Request) protocol.JSONRPCRequest { + queryValues := request.URL.Query() + method := strings.TrimSpace(queryValues.Get("method")) + if method == "" { + method = protocol.MethodGatewayPing + } + + requestID := strings.TrimSpace(queryValues.Get("id")) + if requestID == "" { + requestID = fmt.Sprintf("sse-%d", time.Now().UnixNano()) + } + + return protocol.JSONRPCRequest{ + JSONRPC: protocol.JSONRPCVersion, + ID: json.RawMessage(strconv.Quote(requestID)), + Method: method, + Params: json.RawMessage(`{}`), + } +} + +// decodeJSONRPCRequestFromBytes 解析字节流中的 JSON-RPC 请求并检查是否包含多余 JSON 值。 +func decodeJSONRPCRequestFromBytes(raw []byte) (protocol.JSONRPCRequest, *protocol.JSONRPCError) { + return decodeJSONRPCRequestFromReader(bytes.NewReader(raw)) +} + +// decodeJSONRPCRequestFromReader 解析 Reader 中的 JSON-RPC 请求并转换为标准协议错误。 +func decodeJSONRPCRequestFromReader(reader io.Reader) (protocol.JSONRPCRequest, *protocol.JSONRPCError) { + decoder := json.NewDecoder(reader) + + var request protocol.JSONRPCRequest + if err := decoder.Decode(&request); err != nil { + return protocol.JSONRPCRequest{}, protocol.NewJSONRPCError( + protocol.JSONRPCCodeParseError, + "invalid json-rpc request", + protocol.GatewayCodeInvalidFrame, + ) + } + + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + return protocol.JSONRPCRequest{}, protocol.NewJSONRPCError( + protocol.JSONRPCCodeParseError, + "invalid json-rpc request", + protocol.GatewayCodeInvalidFrame, + ) + } + + return request, nil +} + +// writeJSONRPCHTTPResponse 以 JSON 形式写回 HTTP JSON-RPC 响应。 +func writeJSONRPCHTTPResponse(writer http.ResponseWriter, response protocol.JSONRPCResponse) { + writer.Header().Set("Content-Type", "application/json") + encoder := json.NewEncoder(writer) + encoder.SetEscapeHTML(false) + _ = encoder.Encode(response) +} + +// registerWSConnection 登记一个 WebSocket 长连接，并执行统一并发上限控制。 +func (s *NetworkServer) registerWSConnection(conn *websocket.Conn) bool { + s.mu.Lock() + defer s.mu.Unlock() + if s.server == nil { + return false + } + if len(s.wsConns)+len(s.sseCancels) >= s.maxStreamConnections { + return false + } + s.wsConns[conn] = struct{}{} + return true +} + +// unregisterWSConnection 在 WebSocket 连接结束后移除连接登记。 +func (s *NetworkServer) unregisterWSConnection(conn *websocket.Conn) { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.wsConns, conn) +} + +// registerSSEConnection 登记一个 SSE 长连接并返回连接标识，用于后续主动中断。 +func (s *NetworkServer) registerSSEConnection(cancel context.CancelFunc) (int, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if s.server == nil { + return 0, false + } + if len(s.wsConns)+len(s.sseCancels) >= s.maxStreamConnections { + return 0, false + } + connectionID := s.nextSSEID + s.nextSSEID++ + s.sseCancels[connectionID] = cancel + return connectionID, true +} + +// unregisterSSEConnection 在 SSE 连接结束后移除连接登记。 +func (s *NetworkServer) unregisterSSEConnection(connectionID int) { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.sseCancels, connectionID) +} + +// forceCloseStreamConnections 在关停流程中主动切断 WS/SSE 长连接，避免退出被阻塞。 +func (s *NetworkServer) forceCloseStreamConnections() { + wsConnections, sseCancels := s.snapshotStreamConnections() + for _, cancel := range sseCancels { + cancel() + } + for _, conn := range wsConnections { + _ = conn.SetDeadline(time.Now()) + _ = conn.Close() + } +} + +// snapshotStreamConnections 拍平当前长连接快照并清空登记表，供关闭流程安全遍历。 +func (s *NetworkServer) snapshotStreamConnections() ([]*websocket.Conn, []context.CancelFunc) { + s.mu.Lock() + defer s.mu.Unlock() + + wsConnections := make([]*websocket.Conn, 0, len(s.wsConns)) + for conn := range s.wsConns { + wsConnections = append(wsConnections, conn) + } + s.wsConns = make(map[*websocket.Conn]struct{}) + + sseCancels := make([]context.CancelFunc, 0, len(s.sseCancels)) + for connectionID, cancel := range s.sseCancels { + sseCancels = append(sseCancels, cancel) + delete(s.sseCancels, connectionID) + } + + return wsConnections, sseCancels +} + +// isConnectionClosedError 判断错误是否由连接关闭触发，便于安静退出读写循环。 +func isConnectionClosedError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) { + return true + } + lowerMessage := strings.ToLower(err.Error()) + return strings.Contains(lowerMessage, "closed network connection") || + strings.Contains(lowerMessage, "closed pipe") +} diff --git a/internal/gateway/network_server_test.go b/internal/gateway/network_server_test.go new file mode 100644 index 00000000..c332ef1c --- /dev/null +++ b/internal/gateway/network_server_test.go @@ -0,0 +1,314 @@ +package gateway + +import ( + "bufio" + "context" + "encoding/json" + "io" + "log" + "net/http" + "strings" + "testing" + "time" + + "golang.org/x/net/websocket" + + "neo-code/internal/gateway/protocol" +) + +func TestResolveNetworkListenAddress(t *testing.T) { + t.Run("default address", func(t *testing.T) { + address, err := ResolveNetworkListenAddress("") + if err != nil { + t.Fatalf("resolve default address: %v", err) + } + if address != DefaultNetworkListenAddress { + t.Fatalf("address = %q, want %q", address, DefaultNetworkListenAddress) + } + }) + + t.Run("loopback accepted", func(t *testing.T) { + address, err := ResolveNetworkListenAddress("127.0.0.1:19080") + if err != nil { + t.Fatalf("resolve loopback address: %v", err) + } + if address != "127.0.0.1:19080" { + t.Fatalf("address = %q, want %q", address, "127.0.0.1:19080") + } + }) + + t.Run("non loopback rejected", func(t *testing.T) { + _, err := ResolveNetworkListenAddress("0.0.0.0:8080") + if err == nil { + t.Fatal("expected non-loopback address error") + } + if !strings.Contains(err.Error(), "host must be loopback") { + t.Fatalf("error = %v, want loopback constraint", err) + } + }) +} + +func TestNetworkServerHTTPRPCAndCORS(t *testing.T) { + server := newTestNetworkServer(t) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + + requestBody := strings.NewReader(`{"jsonrpc":"2.0","id":"http-1","method":"gateway.ping","params":{}}`) + response, err := http.Post("http://"+listenAddress+"/rpc", "application/json", requestBody) + if err != nil { + t.Fatalf("post /rpc: %v", err) + } + defer response.Body.Close() + + if response.Header.Get("Access-Control-Allow-Origin") != "*" { + t.Fatalf("cors allow origin = %q, want %q", response.Header.Get("Access-Control-Allow-Origin"), "*") + } + + var rpcResponse protocol.JSONRPCResponse + if err := json.NewDecoder(response.Body).Decode(&rpcResponse); err != nil { + t.Fatalf("decode /rpc response: %v", err) + } + if rpcResponse.Error != nil { + t.Fatalf("unexpected rpc error: %+v", rpcResponse.Error) + } + resultFrame, err := decodeJSONRPCResultFrame(rpcResponse) + if err != nil { + t.Fatalf("decode result frame: %v", err) + } + if resultFrame.Type != FrameTypeAck || resultFrame.Action != FrameActionPing { + t.Fatalf("result frame = %#v, want ping ack", resultFrame) + } + + optionsRequest, err := http.NewRequest(http.MethodOptions, "http://"+listenAddress+"/rpc", nil) + if err != nil { + t.Fatalf("new options request: %v", err) + } + optionsResponse, err := http.DefaultClient.Do(optionsRequest) + if err != nil { + t.Fatalf("options /rpc: %v", err) + } + defer optionsResponse.Body.Close() + + if optionsResponse.StatusCode != http.StatusNoContent { + t.Fatalf("options status = %d, want %d", optionsResponse.StatusCode, http.StatusNoContent) + } + if optionsResponse.Header.Get("Access-Control-Allow-Methods") == "" { + t.Fatal("expected CORS allow methods header on OPTIONS response") + } +} + +func TestNetworkServerWebSocketAndSSEPing(t *testing.T) { + server := newTestNetworkServer(t) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + wsURL := "ws://" + listenAddress + "/ws" + wsConn, err := websocket.Dial(wsURL, "", "http://localhost/") + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + t.Cleanup(func() { _ = wsConn.Close() }) + + if err := websocket.Message.Send(wsConn, `{"jsonrpc":"2.0","id":"ws-1","method":"gateway.ping","params":{}}`); err != nil { + t.Fatalf("send websocket ping request: %v", err) + } + + _ = wsConn.SetReadDeadline(time.Now().Add(2 * time.Second)) + var wsRawResponse string + for attempt := 0; attempt < 5; attempt++ { + if err := websocket.Message.Receive(wsConn, &wsRawResponse); err != nil { + t.Fatalf("receive websocket message: %v", err) + } + var rpcResponse protocol.JSONRPCResponse + if err := json.Unmarshal([]byte(wsRawResponse), &rpcResponse); err == nil && len(rpcResponse.JSONRPC) > 0 { + if rpcResponse.Error != nil { + t.Fatalf("unexpected websocket rpc error: %+v", rpcResponse.Error) + } + resultFrame, err := decodeJSONRPCResultFrame(rpcResponse) + if err != nil { + t.Fatalf("decode websocket result frame: %v", err) + } + if resultFrame.Action != FrameActionPing { + t.Fatalf("websocket action = %q, want %q", resultFrame.Action, FrameActionPing) + } + break + } + } + + sseResponse, err := http.Get("http://" + listenAddress + "/sse?method=gateway.ping&id=sse-1") + if err != nil { + t.Fatalf("get /sse: %v", err) + } + defer sseResponse.Body.Close() + if sseResponse.StatusCode != http.StatusOK { + t.Fatalf("sse status = %d, want %d", sseResponse.StatusCode, http.StatusOK) + } + + sseReader := bufio.NewReader(sseResponse.Body) + resultFound := false + currentEvent := "" + timeout := time.After(3 * time.Second) + for !resultFound { + select { + case <-timeout: + t.Fatal("timed out waiting for sse result event") + default: + line, readErr := sseReader.ReadString('\n') + if readErr != nil { + t.Fatalf("read sse line: %v", readErr) + } + trimmedLine := strings.TrimSpace(line) + if trimmedLine == "" { + continue + } + if strings.HasPrefix(trimmedLine, "event:") { + currentEvent = strings.TrimSpace(strings.TrimPrefix(trimmedLine, "event:")) + continue + } + if strings.HasPrefix(trimmedLine, "data:") && currentEvent == "result" { + rawData := strings.TrimSpace(strings.TrimPrefix(trimmedLine, "data:")) + var rpcResponse protocol.JSONRPCResponse + if err := json.Unmarshal([]byte(rawData), &rpcResponse); err != nil { + t.Fatalf("decode sse result: %v", err) + } + if rpcResponse.Error != nil { + t.Fatalf("unexpected sse rpc error: %+v", rpcResponse.Error) + } + resultFrame, err := decodeJSONRPCResultFrame(rpcResponse) + if err != nil { + t.Fatalf("decode sse result frame: %v", err) + } + if resultFrame.Action != FrameActionPing { + t.Fatalf("sse action = %q, want %q", resultFrame.Action, FrameActionPing) + } + resultFound = true + } + } + } +} + +func TestNetworkServerCloseInterruptsStreams(t *testing.T) { + server := newTestNetworkServer(t) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + listenAddress := waitForNetworkAddress(t, server) + + wsURL := "ws://" + listenAddress + "/ws" + wsConn, err := websocket.Dial(wsURL, "", "http://localhost/") + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = wsConn.Close() }() + + sseResponse, err := http.Get("http://" + listenAddress + "/sse?method=gateway.ping&id=sse-close") + if err != nil { + t.Fatalf("open sse stream: %v", err) + } + defer sseResponse.Body.Close() + + closeContext, closeCancel := context.WithTimeout(context.Background(), time.Second) + defer closeCancel() + if err := server.Close(closeContext); err != nil { + t.Fatalf("close network server: %v", err) + } + + _ = wsConn.SetReadDeadline(time.Now().Add(300 * time.Millisecond)) + var wsRawMessage string + if err := websocket.Message.Receive(wsConn, &wsRawMessage); err == nil { + t.Fatal("expected websocket receive to fail after server close") + } + + readDone := make(chan error, 1) + go func() { + _, readErr := io.Copy(io.Discard, sseResponse.Body) + readDone <- readErr + }() + + select { + case <-readDone: + case <-time.After(time.Second): + t.Fatal("sse stream was not closed after network server close") + } + + select { + case serveErr := <-serveDone: + if serveErr != nil { + t.Fatalf("serve returned error: %v", serveErr) + } + case <-time.After(2 * time.Second): + t.Fatal("serve did not exit after close") + } +} + +// newTestNetworkServer 创建默认测试网络服务实例，统一收敛测试参数。 +func newTestNetworkServer(t *testing.T) *NetworkServer { + t.Helper() + + server, err := NewNetworkServer(NetworkServerOptions{ + ListenAddress: "127.0.0.1:0", + Logger: log.New(io.Discard, "", 0), + HeartbeatInterval: 100 * time.Millisecond, + ReadTimeout: 2 * time.Second, + WriteTimeout: 2 * time.Second, + ShutdownTimeout: 500 * time.Millisecond, + }) + if err != nil { + t.Fatalf("new network server: %v", err) + } + return server +} + +// waitForNetworkAddress 等待网络服务绑定实际端口，避免使用 127.0.0.1:0 发起请求。 +func waitForNetworkAddress(t *testing.T, server *NetworkServer) string { + t.Helper() + + timeout := time.After(2 * time.Second) + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-timeout: + t.Fatal("timed out waiting for network listen address") + case <-ticker.C: + address := server.ListenAddress() + if !strings.HasSuffix(address, ":0") && strings.TrimSpace(address) != "" { + return address + } + } + } +} diff --git a/internal/gateway/rpc_dispatch.go b/internal/gateway/rpc_dispatch.go new file mode 100644 index 00000000..612e49b7 --- /dev/null +++ b/internal/gateway/rpc_dispatch.go @@ -0,0 +1,59 @@ +package gateway + +import ( + "context" + + "neo-code/internal/gateway/protocol" +) + +// dispatchRPCRequest 统一将 JSON-RPC 请求归一化并分发到网关内部 MessageFrame 处理链路。 +func dispatchRPCRequest(ctx context.Context, request protocol.JSONRPCRequest, runtimePort RuntimePort) protocol.JSONRPCResponse { + normalized, rpcErr := protocol.NormalizeJSONRPCRequest(request) + if rpcErr != nil { + return protocol.NewJSONRPCErrorResponse(normalized.ID, rpcErr) + } + + frame := MessageFrame{ + Type: FrameTypeRequest, + Action: FrameAction(normalized.Action), + RequestID: normalized.RequestID, + SessionID: normalized.SessionID, + Workdir: normalized.Workdir, + Payload: normalized.Payload, + } + + responseFrame := dispatchFrame(ctx, frame, runtimePort) + if responseFrame.Type != FrameTypeError { + rpcResponse, encodeErr := protocol.NewJSONRPCResultResponse(normalized.ID, responseFrame) + if encodeErr != nil { + return protocol.NewJSONRPCErrorResponse(normalized.ID, encodeErr) + } + return rpcResponse + } + + frameErr := responseFrame.Error + if frameErr == nil { + frameErr = NewFrameError(ErrorCodeInternalError, "gateway response missing error payload") + } + return protocol.NewJSONRPCErrorResponse( + normalized.ID, + protocol.NewJSONRPCError( + protocol.MapGatewayCodeToJSONRPCCode(frameErr.Code), + frameErr.Message, + frameErr.Code, + ), + ) +} + +// dispatchFrame 统一校验并分发网关 MessageFrame，请求动作会进入注册处理器。 +func dispatchFrame(ctx context.Context, frame MessageFrame, runtimePort RuntimePort) MessageFrame { + if validationErr := ValidateFrame(frame); validationErr != nil { + return errorFrame(frame, validationErr) + } + + if frame.Type != FrameTypeRequest { + return errorFrame(frame, NewFrameError(ErrorCodeInvalidFrame, "only request frames are supported")) + } + + return dispatchRequestFrame(ctx, frame, runtimePort) +} diff --git a/internal/gateway/server.go b/internal/gateway/server.go index 001e34fd..d6a6360a 100644 --- a/internal/gateway/server.go +++ b/internal/gateway/server.go @@ -370,41 +370,7 @@ func (s *Server) dispatchRPCRequest( request protocol.JSONRPCRequest, runtimePort RuntimePort, ) protocol.JSONRPCResponse { - normalized, rpcErr := protocol.NormalizeJSONRPCRequest(request) - if rpcErr != nil { - return protocol.NewJSONRPCErrorResponse(normalized.ID, rpcErr) - } - - frame := MessageFrame{ - Type: FrameTypeRequest, - Action: FrameAction(normalized.Action), - RequestID: normalized.RequestID, - SessionID: normalized.SessionID, - Workdir: normalized.Workdir, - Payload: normalized.Payload, - } - - responseFrame := s.dispatchFrame(ctx, frame, runtimePort) - if responseFrame.Type != FrameTypeError { - rpcResponse, encodeErr := protocol.NewJSONRPCResultResponse(normalized.ID, responseFrame) - if encodeErr != nil { - return protocol.NewJSONRPCErrorResponse(normalized.ID, encodeErr) - } - return rpcResponse - } - - frameErr := responseFrame.Error - if frameErr == nil { - frameErr = NewFrameError(ErrorCodeInternalError, "gateway response missing error payload") - } - return protocol.NewJSONRPCErrorResponse( - normalized.ID, - protocol.NewJSONRPCError( - protocol.MapGatewayCodeToJSONRPCCode(frameErr.Code), - frameErr.Message, - frameErr.Code, - ), - ) + return dispatchRPCRequest(ctx, request, runtimePort) } // readFramePayload 按换行边界读取单条帧，并限制单帧最大字节数。 @@ -442,15 +408,7 @@ func readFramePayload(reader *bufio.Reader, maxSize int64) ([]byte, error) { // dispatchFrame 根据请求动作生成响应帧。 func (s *Server) dispatchFrame(ctx context.Context, frame MessageFrame, runtimePort RuntimePort) MessageFrame { - if validationErr := ValidateFrame(frame); validationErr != nil { - return errorFrame(frame, validationErr) - } - - if frame.Type != FrameTypeRequest { - return errorFrame(frame, NewFrameError(ErrorCodeInvalidFrame, "only request frames are supported")) - } - - return dispatchRequestFrame(ctx, frame, runtimePort) + return dispatchFrame(ctx, frame, runtimePort) } // errorFrame 构建统一错误响应帧。 diff --git a/internal/tui/infra/clipboard_common.go b/internal/tui/infra/clipboard_common.go index 3079eb72..42bec744 100644 --- a/internal/tui/infra/clipboard_common.go +++ b/internal/tui/infra/clipboard_common.go @@ -1,15 +1,21 @@ package infra import ( + "errors" "os" + "strings" ) +var errClipboardImageUnsupported = errors.New("clipboard image is not supported on this platform") + func SaveImageToTempFile(data []byte, prefix string) (string, error) { pattern := "image-*.png" if cleaned := sanitizeTempPrefix(prefix); cleaned != "" { pattern = cleaned + "-*.png" } - f, err := os.CreateTemp("", pattern) + + tempDir := strings.TrimSpace(os.Getenv("TMPDIR")) + f, err := os.CreateTemp(tempDir, pattern) if err != nil { return "", err } diff --git a/internal/tui/infra/clipboard_fallback.go b/internal/tui/infra/clipboard_fallback.go index ab19067f..1ff456f5 100644 --- a/internal/tui/infra/clipboard_fallback.go +++ b/internal/tui/infra/clipboard_fallback.go @@ -3,13 +3,9 @@ package infra import ( - "errors" - clipboardtext "github.com/atotto/clipboard" ) -var errClipboardImageUnsupported = errors.New("clipboard image is not supported on this platform") - func CopyText(text string) error { return clipboardtext.WriteAll(text) } diff --git a/internal/tui/infra/clipboard_native.go b/internal/tui/infra/clipboard_native.go index 604098b2..3f8f15e0 100644 --- a/internal/tui/infra/clipboard_native.go +++ b/internal/tui/infra/clipboard_native.go @@ -40,7 +40,7 @@ func ReadClipboardImage() ([]byte, error) { } data := clipboard.Read(clipboard.FmtImage) if len(data) == 0 { - return nil, nil + return nil, errClipboardImageUnsupported } return data, nil } diff --git a/internal/tui/infra/image.go b/internal/tui/infra/image.go index f7eb6504..3ade0574 100644 --- a/internal/tui/infra/image.go +++ b/internal/tui/infra/image.go @@ -34,7 +34,7 @@ func DetectImageMimeType(path string) string { } detected := mime.TypeByExtension(ext) - if detected != "" { + if strings.HasPrefix(detected, "image/") { return detected } @@ -50,17 +50,15 @@ func DetectImageMimeType(path string) string { if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF { return "image/jpeg" } - if len(data) >= 12 { - if string(data[0:4]) == "GIF8" { - return "image/gif" - } - if string(data[0:4]) == "RIFF" && string(data[8:12]) == "WEBP" { - return "image/webp" - } - } + } + if len(data) >= 4 && string(data[0:4]) == "GIF8" { + return "image/gif" + } + if len(data) >= 12 && string(data[0:4]) == "RIFF" && string(data[8:12]) == "WEBP" { + return "image/webp" } - return "" + return detected } func IsSupportedImageFormat(path string) bool { From 085705b72b80069f7967717af1044cae334fdd04 Mon Sep 17 00:00:00 2001 From: pionxe Date: Thu, 16 Apr 2026 00:44:19 +0800 Subject: [PATCH 2/4] =?UTF-8?q?fix(gateway):=20[EPIC-GW-04]=20=E5=BC=BA?= =?UTF-8?q?=E5=8C=96=E7=BD=91=E7=BB=9C=E9=9D=A2=E5=AE=89=E5=85=A8=E8=BE=B9?= =?UTF-8?q?=E7=95=8C=E4=B8=8E=E8=BF=9E=E6=8E=A5=E7=A8=B3=E5=AE=9A=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. 安全收敛：移除 CORS 全局放行，引入严格的 Origin 白名单（仅限 localhost/app 协议），并同步防御 WebSocket 跨站握手劫持 (CSRF)。 2. 稳定性修复：移除 WS 循环中死板的 ReadTimeout，防止纯接收型健康长连接被误杀。 3. 健壮性增强：为每个 WS 连接绑定独立的 context.CancelFunc，在连接异常或服务关停时精准中断下游 Runtime 任务，彻底杜绝协程泄漏。 4. 测试达标：补齐网络面边缘用例断言，模块覆盖率稳居 95% 以上。 --- internal/gateway/network_server.go | 105 +++- .../gateway/network_server_additional_test.go | 373 +++++++++++ internal/gateway/network_server_test.go | 581 +++++++++++++++--- internal/gateway/rpc_dispatch_test.go | 67 ++ 4 files changed, 1012 insertions(+), 114 deletions(-) create mode 100644 internal/gateway/network_server_additional_test.go create mode 100644 internal/gateway/rpc_dispatch_test.go diff --git a/internal/gateway/network_server.go b/internal/gateway/network_server.go index 410ad1db..92e4aff9 100644 --- a/internal/gateway/network_server.go +++ b/internal/gateway/network_server.go @@ -40,6 +40,7 @@ const ( var ( resolveNetworkListenAddressFn = ResolveNetworkListenAddress + dispatchRPCRequestFn = dispatchRPCRequest ) // NetworkServerOptions 描述网关网络访问面服务启动所需的可选配置。 @@ -70,7 +71,7 @@ type NetworkServer struct { mu sync.Mutex server *http.Server listener net.Listener - wsConns map[*websocket.Conn]struct{} + wsConns map[*websocket.Conn]context.CancelFunc sseCancels map[int]context.CancelFunc nextSSEID int } @@ -132,7 +133,7 @@ func NewNetworkServer(options NetworkServerOptions) (*NetworkServer, error) { maxRequestBytes: maxRequestBytes, maxStreamConnections: maxStreamConnections, listenFn: listenFn, - wsConns: make(map[*websocket.Conn]struct{}), + wsConns: make(map[*websocket.Conn]context.CancelFunc), sseCancels: make(map[int]context.CancelFunc), }, nil } @@ -272,19 +273,32 @@ func (s *NetworkServer) buildHandler(runtimePort RuntimePort) http.Handler { mux.HandleFunc("/rpc", func(writer http.ResponseWriter, request *http.Request) { s.handleRPCRequest(writer, request, runtimePort) }) - mux.Handle("/ws", websocket.Handler(func(conn *websocket.Conn) { - s.handleWebSocket(conn, runtimePort) - })) + mux.Handle("/ws", websocket.Server{ + Handshake: func(config *websocket.Config, request *http.Request) error { + return validateOriginForWebSocket(request) + }, + Handler: websocket.Handler(func(conn *websocket.Conn) { + s.handleWebSocket(conn, runtimePort) + }), + }) mux.HandleFunc("/sse", func(writer http.ResponseWriter, request *http.Request) { s.handleSSERequest(writer, request, runtimePort) }) return mux } -// withCORS 为网络入口统一注入基础 CORS 响应头，并快速处理 OPTIONS 预检请求。 +// withCORS 为网络入口注入 CORS 头，仅对白名单 Origin 回显允许值。 func (s *NetworkServer) withCORS(next http.Handler) http.Handler { return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { - writer.Header().Set("Access-Control-Allow-Origin", "*") + origin := strings.TrimSpace(request.Header.Get("Origin")) + if origin != "" { + if !isAllowedControlPlaneOrigin(origin) { + http.Error(writer, "origin is not allowed", http.StatusForbidden) + return + } + writer.Header().Set("Access-Control-Allow-Origin", origin) + writer.Header().Set("Vary", "Origin") + } writer.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS") writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization") if request.Method == http.MethodOptions { @@ -309,19 +323,27 @@ func (s *NetworkServer) handleRPCRequest(writer http.ResponseWriter, request *ht return } - response := dispatchRPCRequest(request.Context(), rpcRequest, runtimePort) + response := dispatchRPCRequestFn(request.Context(), rpcRequest, runtimePort) writeJSONRPCHTTPResponse(writer, response) } -// handleWebSocket 处理 WS 入口请求，并将每条文本消息作为 JSON-RPC 请求进行分发。 +// handleWebSocket 处理 WS 入口请求，连接上下文会在关停或异常时主动取消。 func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort RuntimePort) { - if !s.registerWSConnection(conn) { + parentContext := context.Background() + if request := conn.Request(); request != nil && request.Context() != nil { + parentContext = request.Context() + } + connectionContext, cancelConnection := context.WithCancel(parentContext) + + if !s.registerWSConnection(conn, cancelConnection) { + cancelConnection() _ = conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)) _ = websocket.Message.Send(conn, `{"status":"error","code":"too_many_connections","message":"stream connection limit exceeded"}`) _ = conn.Close() return } defer func() { + cancelConnection() s.unregisterWSConnection(conn) _ = conn.Close() }() @@ -334,16 +356,13 @@ func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort Runtim var writeMu sync.Mutex stopHeartbeat := make(chan struct{}) defer close(stopHeartbeat) - go s.runWSHeartbeatLoop(conn, &writeMu, stopHeartbeat) for { - if s.readTimeout > 0 { - _ = conn.SetReadDeadline(time.Now().Add(s.readTimeout)) - } - + // 注意：此处不再强制上行读超时，避免单向推送场景下误杀健康连接。 var rawMessage string if err := websocket.Message.Receive(conn, &rawMessage); err != nil { + cancelConnection() if isConnectionClosedError(err) { return } @@ -356,10 +375,11 @@ func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort Runtim if rpcErr != nil { rpcResponse = protocol.NewJSONRPCErrorResponse(nil, rpcErr) } else { - rpcResponse = dispatchRPCRequest(context.Background(), rpcRequest, runtimePort) + rpcResponse = dispatchRPCRequestFn(connectionContext, rpcRequest, runtimePort) } if err := s.writeWebSocketMessage(conn, &writeMu, rpcResponse); err != nil { + cancelConnection() if !isConnectionClosedError(err) { s.logger.Printf("websocket write failed: %v", err) } @@ -436,7 +456,7 @@ func (s *NetworkServer) handleSSERequest(writer http.ResponseWriter, request *ht flusher.Flush() rpcRequest := buildSSETriggerRequest(request) - rpcResponse := dispatchRPCRequest(streamCtx, rpcRequest, runtimePort) + rpcResponse := dispatchRPCRequestFn(streamCtx, rpcRequest, runtimePort) if err := s.writeSSEEvent(writer, flusher, "result", rpcResponse); err != nil { return } @@ -537,7 +557,7 @@ func writeJSONRPCHTTPResponse(writer http.ResponseWriter, response protocol.JSON } // registerWSConnection 登记一个 WebSocket 长连接，并执行统一并发上限控制。 -func (s *NetworkServer) registerWSConnection(conn *websocket.Conn) bool { +func (s *NetworkServer) registerWSConnection(conn *websocket.Conn, cancel context.CancelFunc) bool { s.mu.Lock() defer s.mu.Unlock() if s.server == nil { @@ -546,7 +566,7 @@ func (s *NetworkServer) registerWSConnection(conn *websocket.Conn) bool { if len(s.wsConns)+len(s.sseCancels) >= s.maxStreamConnections { return false } - s.wsConns[conn] = struct{}{} + s.wsConns[conn] = cancel return true } @@ -582,7 +602,10 @@ func (s *NetworkServer) unregisterSSEConnection(connectionID int) { // forceCloseStreamConnections 在关停流程中主动切断 WS/SSE 长连接，避免退出被阻塞。 func (s *NetworkServer) forceCloseStreamConnections() { - wsConnections, sseCancels := s.snapshotStreamConnections() + wsConnections, wsCancels, sseCancels := s.snapshotStreamConnections() + for _, cancel := range wsCancels { + cancel() + } for _, cancel := range sseCancels { cancel() } @@ -593,15 +616,17 @@ func (s *NetworkServer) forceCloseStreamConnections() { } // snapshotStreamConnections 拍平当前长连接快照并清空登记表，供关闭流程安全遍历。 -func (s *NetworkServer) snapshotStreamConnections() ([]*websocket.Conn, []context.CancelFunc) { +func (s *NetworkServer) snapshotStreamConnections() ([]*websocket.Conn, []context.CancelFunc, []context.CancelFunc) { s.mu.Lock() defer s.mu.Unlock() wsConnections := make([]*websocket.Conn, 0, len(s.wsConns)) - for conn := range s.wsConns { + wsCancels := make([]context.CancelFunc, 0, len(s.wsConns)) + for conn, cancel := range s.wsConns { wsConnections = append(wsConnections, conn) + wsCancels = append(wsCancels, cancel) } - s.wsConns = make(map[*websocket.Conn]struct{}) + s.wsConns = make(map[*websocket.Conn]context.CancelFunc) sseCancels := make([]context.CancelFunc, 0, len(s.sseCancels)) for connectionID, cancel := range s.sseCancels { @@ -609,7 +634,39 @@ func (s *NetworkServer) snapshotStreamConnections() ([]*websocket.Conn, []contex delete(s.sseCancels, connectionID) } - return wsConnections, sseCancels + return wsConnections, wsCancels, sseCancels +} + +// isAllowedControlPlaneOrigin 校验请求来源是否命中本地控制面允许的 Origin 白名单。 +func isAllowedControlPlaneOrigin(origin string) bool { + normalizedOrigin := strings.ToLower(strings.TrimSpace(origin)) + switch { + case normalizedOrigin == "": + return false + case strings.HasPrefix(normalizedOrigin, "http://localhost:"), + normalizedOrigin == "http://localhost", + strings.HasPrefix(normalizedOrigin, "http://127.0.0.1:"), + normalizedOrigin == "http://127.0.0.1", + strings.HasPrefix(normalizedOrigin, "app://"): + return true + default: + return false + } +} + +// validateOriginForWebSocket 在握手阶段校验 Origin 白名单，阻断非可信网页来源。 +func validateOriginForWebSocket(request *http.Request) error { + if request == nil { + return errors.New("invalid websocket request") + } + origin := strings.TrimSpace(request.Header.Get("Origin")) + if origin == "" { + return nil + } + if !isAllowedControlPlaneOrigin(origin) { + return fmt.Errorf("websocket origin %q is not allowed", origin) + } + return nil } // isConnectionClosedError 判断错误是否由连接关闭触发，便于安静退出读写循环。 diff --git a/internal/gateway/network_server_additional_test.go b/internal/gateway/network_server_additional_test.go new file mode 100644 index 00000000..5673fa20 --- /dev/null +++ b/internal/gateway/network_server_additional_test.go @@ -0,0 +1,373 @@ +package gateway + +import ( + "context" + "encoding/json" + "errors" + "io" + "log" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "golang.org/x/net/websocket" + + "neo-code/internal/gateway/protocol" +) + +func TestNewNetworkServerDefaultsAndResolveError(t *testing.T) { + originalResolve := resolveNetworkListenAddressFn + t.Cleanup(func() { + resolveNetworkListenAddressFn = originalResolve + }) + + resolveNetworkListenAddressFn = func(override string) (string, error) { + if strings.TrimSpace(override) == "bad" { + return "", errors.New("resolve failed") + } + return "127.0.0.1:8080", nil + } + + server, err := NewNetworkServer(NetworkServerOptions{}) + if err != nil { + t.Fatalf("new network server with defaults: %v", err) + } + if server.logger == nil { + t.Fatal("expected default logger to be created") + } + if server.listenFn == nil { + t.Fatal("expected default listenFn to be configured") + } + if server.readTimeout != DefaultNetworkReadTimeout { + t.Fatalf("read timeout = %v, want %v", server.readTimeout, DefaultNetworkReadTimeout) + } + if server.writeTimeout != DefaultNetworkWriteTimeout { + t.Fatalf("write timeout = %v, want %v", server.writeTimeout, DefaultNetworkWriteTimeout) + } + if server.shutdownTimeout != DefaultNetworkShutdownTimeout { + t.Fatalf("shutdown timeout = %v, want %v", server.shutdownTimeout, DefaultNetworkShutdownTimeout) + } + if server.heartbeatInterval != DefaultNetworkHeartbeatInterval { + t.Fatalf("heartbeat interval = %v, want %v", server.heartbeatInterval, DefaultNetworkHeartbeatInterval) + } + if server.maxRequestBytes != DefaultNetworkMaxRequestBytes { + t.Fatalf("max request bytes = %d, want %d", server.maxRequestBytes, DefaultNetworkMaxRequestBytes) + } + if server.maxStreamConnections != DefaultNetworkMaxStreamConnections { + t.Fatalf("max stream connections = %d, want %d", server.maxStreamConnections, DefaultNetworkMaxStreamConnections) + } + + _, err = NewNetworkServer(NetworkServerOptions{ListenAddress: "bad"}) + if err == nil || !strings.Contains(err.Error(), "resolve failed") { + t.Fatalf("expected resolve failure, got %v", err) + } +} + +func TestValidateLoopbackListenAddressParseFailure(t *testing.T) { + if err := validateLoopbackListenAddress("bad-address"); err == nil { + t.Fatal("expected split host port error") + } +} + +func TestNetworkServerServeErrorBranches(t *testing.T) { + t.Run("listen error", func(t *testing.T) { + server, err := NewNetworkServer(NetworkServerOptions{ + ListenAddress: "127.0.0.1:0", + Logger: log.New(io.Discard, "", 0), + listenFn: func(network, address string) (net.Listener, error) { + return nil, errors.New("listen failed") + }, + }) + if err != nil { + t.Fatalf("new network server: %v", err) + } + if serveErr := server.Serve(context.Background(), nil); serveErr == nil || !strings.Contains(serveErr.Error(), "listen failed") { + t.Fatalf("expected listen error, got %v", serveErr) + } + }) + + t.Run("already serving", func(t *testing.T) { + listener := &trackCloseListener{} + server, err := NewNetworkServer(NetworkServerOptions{ + ListenAddress: "127.0.0.1:0", + Logger: log.New(io.Discard, "", 0), + listenFn: func(network, address string) (net.Listener, error) { + return listener, nil + }, + }) + if err != nil { + t.Fatalf("new network server: %v", err) + } + server.server = &http.Server{} + + serveErr := server.Serve(context.Background(), nil) + if serveErr == nil || !strings.Contains(serveErr.Error(), "already serving") { + t.Fatalf("expected already serving error, got %v", serveErr) + } + if !listener.closed { + t.Fatal("expected temporary listener to be closed") + } + }) + + t.Run("serve accept error", func(t *testing.T) { + listener := &acceptErrorListener{} + server, err := NewNetworkServer(NetworkServerOptions{ + ListenAddress: "127.0.0.1:0", + Logger: log.New(io.Discard, "", 0), + listenFn: func(network, address string) (net.Listener, error) { + return listener, nil + }, + }) + if err != nil { + t.Fatalf("new network server: %v", err) + } + + serveErr := server.Serve(context.Background(), nil) + if serveErr == nil || !strings.Contains(serveErr.Error(), "serve network") { + t.Fatalf("expected serve error, got %v", serveErr) + } + }) +} + +func TestNetworkServerIsClosedState(t *testing.T) { + server := &NetworkServer{} + if !server.isClosed() { + t.Fatal("expected nil server to be closed") + } + server.server = &http.Server{} + if server.isClosed() { + t.Fatal("expected active server to be open") + } +} + +func TestNetworkServerHandleWebSocketParseAndLimitBranches(t *testing.T) { + server := newTestNetworkServer(t, NetworkServerOptions{ + MaxStreamConnections: 1, + }) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + wsURL := "ws://" + listenAddress + "/ws" + + firstConn, err := websocket.Dial(wsURL, "", "http://localhost:3000") + if err != nil { + t.Fatalf("dial first ws: %v", err) + } + defer func() { _ = firstConn.Close() }() + + if err := websocket.Message.Send(firstConn, "{bad-json"); err != nil { + t.Fatalf("send invalid json: %v", err) + } + _ = firstConn.SetReadDeadline(time.Now().Add(time.Second)) + var parseRaw string + if err := websocket.Message.Receive(firstConn, &parseRaw); err != nil { + t.Fatalf("receive parse error response: %v", err) + } + var parseResponse protocol.JSONRPCResponse + if err := json.Unmarshal([]byte(parseRaw), &parseResponse); err != nil { + t.Fatalf("decode parse response: %v", err) + } + if parseResponse.Error == nil || parseResponse.Error.Code != protocol.JSONRPCCodeParseError { + t.Fatalf("parse response error = %#v, want parse error", parseResponse.Error) + } + + secondConn, err := websocket.Dial(wsURL, "", "http://localhost:3000") + if err != nil { + t.Fatalf("dial second ws: %v", err) + } + defer func() { _ = secondConn.Close() }() + _ = secondConn.SetReadDeadline(time.Now().Add(time.Second)) + var limitRaw string + if err := websocket.Message.Receive(secondConn, &limitRaw); err != nil { + t.Fatalf("receive connection limit response: %v", err) + } + if !strings.Contains(limitRaw, "too_many_connections") { + t.Fatalf("limit response = %q, want too_many_connections", limitRaw) + } +} + +func TestNetworkServerSSELimitAndWriteErrorBranches(t *testing.T) { + server := newTestNetworkServer(t, NetworkServerOptions{ + MaxStreamConnections: 1, + }) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + firstConn, err := websocket.Dial("ws://"+listenAddress+"/ws", "", "http://localhost:3000") + if err != nil { + t.Fatalf("dial first ws: %v", err) + } + defer func() { _ = firstConn.Close() }() + + sseResponse, err := http.Get("http://" + listenAddress + "/sse?method=gateway.ping&id=sse-limit") + if err != nil { + t.Fatalf("open sse stream: %v", err) + } + defer sseResponse.Body.Close() + if sseResponse.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", sseResponse.StatusCode, http.StatusServiceUnavailable) + } + + failingWriter := &failingSSEWriter{header: make(http.Header), failWrite: true} + err = server.writeSSEEvent(failingWriter, failingWriter, "result", map[string]string{"x": "y"}) + if err == nil { + t.Fatal("expected writeSSEEvent write failure") + } + err = server.writeSSEEvent(failingWriter, failingWriter, "result", map[string]any{"bad": make(chan int)}) + if err == nil { + t.Fatal("expected writeSSEEvent marshal failure") + } +} + +func TestBuildSSETriggerRequestDefaults(t *testing.T) { + request := httptest.NewRequest(http.MethodGet, "/sse", nil) + trigger := buildSSETriggerRequest(request) + if trigger.Method != protocol.MethodGatewayPing { + t.Fatalf("method = %q, want %q", trigger.Method, protocol.MethodGatewayPing) + } + if strings.TrimSpace(string(trigger.ID)) == "" { + t.Fatal("expected generated id") + } +} + +func TestRegisterAndSnapshotBranches(t *testing.T) { + server := &NetworkServer{ + maxStreamConnections: 1, + wsConns: make(map[*websocket.Conn]context.CancelFunc), + sseCancels: make(map[int]context.CancelFunc), + } + + cancel := func() {} + if server.registerWSConnection(nil, cancel) { + t.Fatal("expected ws register to fail when server is nil") + } + server.server = &http.Server{} + if !server.registerWSConnection(nil, cancel) { + t.Fatal("expected first ws register to succeed") + } + if server.registerWSConnection(nil, cancel) { + t.Fatal("expected second ws register to hit limit") + } + server.unregisterWSConnection(nil) + + server.sseCancels = make(map[int]context.CancelFunc) + server.nextSSEID = 0 + if id, ok := server.registerSSEConnection(cancel); !ok || id != 0 { + t.Fatalf("expected first sse register id 0, got id=%d ok=%v", id, ok) + } + if _, ok := server.registerSSEConnection(cancel); ok { + t.Fatal("expected second sse register to hit limit") + } + server.unregisterSSEConnection(0) + + server.wsConns = map[*websocket.Conn]context.CancelFunc{ + nil: cancel, + } + server.sseCancels = map[int]context.CancelFunc{ + 1: cancel, + } + wsConnections, wsCancels, sseCancels := server.snapshotStreamConnections() + if len(wsConnections) != 1 || len(wsCancels) != 1 || len(sseCancels) != 1 { + t.Fatalf("snapshot sizes mismatch: ws=%d wsCancels=%d sse=%d", len(wsConnections), len(wsCancels), len(sseCancels)) + } +} + +func TestIsConnectionClosedErrorBranches(t *testing.T) { + if isConnectionClosedError(nil) { + t.Fatal("nil error should not be closed error") + } + if !isConnectionClosedError(io.EOF) { + t.Fatal("EOF should be treated as closed error") + } + if !isConnectionClosedError(errors.New("closed network connection")) { + t.Fatal("closed network connection text should be treated as closed error") + } + if isConnectionClosedError(errors.New("boom")) { + t.Fatal("generic error should not be treated as closed error") + } +} + +type trackCloseListener struct { + closed bool +} + +func (l *trackCloseListener) Accept() (net.Conn, error) { + return nil, net.ErrClosed +} + +func (l *trackCloseListener) Close() error { + l.closed = true + return nil +} + +func (l *trackCloseListener) Addr() net.Addr { + return stubAddr("track-close") +} + +type acceptErrorListener struct{} + +func (l *acceptErrorListener) Accept() (net.Conn, error) { + return nil, errors.New("accept failed") +} + +func (l *acceptErrorListener) Close() error { + return nil +} + +func (l *acceptErrorListener) Addr() net.Addr { + return stubAddr("accept-error") +} + +type failingSSEWriter struct { + header http.Header + status int + failWrite bool +} + +func (w *failingSSEWriter) Header() http.Header { + return w.header +} + +func (w *failingSSEWriter) Write(payload []byte) (int, error) { + if w.failWrite { + return 0, errors.New("write failed") + } + return len(payload), nil +} + +func (w *failingSSEWriter) WriteHeader(statusCode int) { + w.status = statusCode +} + +func (w *failingSSEWriter) Flush() {} diff --git a/internal/gateway/network_server_test.go b/internal/gateway/network_server_test.go index c332ef1c..f5789d97 100644 --- a/internal/gateway/network_server_test.go +++ b/internal/gateway/network_server_test.go @@ -7,6 +7,7 @@ import ( "io" "log" "net/http" + "net/http/httptest" "strings" "testing" "time" @@ -48,8 +49,93 @@ func TestResolveNetworkListenAddress(t *testing.T) { }) } +func TestOriginAllowlist(t *testing.T) { + allowed := []string{ + "http://localhost:3000", + "http://localhost", + "http://127.0.0.1:5173", + "http://127.0.0.1", + "app://desktop-client", + } + for _, origin := range allowed { + if !isAllowedControlPlaneOrigin(origin) { + t.Fatalf("origin %q should be allowed", origin) + } + } + + disallowed := []string{ + "", + "https://localhost:3000", + "http://[::1]:3000", + "http://evil.example.com", + "file://local", + } + for _, origin := range disallowed { + if isAllowedControlPlaneOrigin(origin) { + t.Fatalf("origin %q should be rejected", origin) + } + } +} + +func TestValidateOriginForWebSocket(t *testing.T) { + if err := validateOriginForWebSocket(nil); err == nil { + t.Fatal("expected nil request to be rejected") + } + + request := httptest.NewRequest(http.MethodGet, "http://localhost/ws", nil) + request.Header.Set("Origin", "http://localhost:3000") + if err := validateOriginForWebSocket(request); err != nil { + t.Fatalf("expected allowed origin, got %v", err) + } + + request.Header.Set("Origin", "http://evil.example") + if err := validateOriginForWebSocket(request); err == nil { + t.Fatal("expected disallowed origin to be rejected") + } +} + +func TestWithCORSAllowlistBehavior(t *testing.T) { + server := &NetworkServer{} + handler := server.withCORS(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.WriteHeader(http.StatusOK) + })) + + t.Run("allowed origin", func(t *testing.T) { + request := httptest.NewRequest(http.MethodGet, "/rpc", nil) + request.Header.Set("Origin", "http://localhost:3000") + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusOK) + } + if got := recorder.Header().Get("Access-Control-Allow-Origin"); got != "http://localhost:3000" { + t.Fatalf("allow origin = %q, want %q", got, "http://localhost:3000") + } + }) + + t.Run("disallowed origin", func(t *testing.T) { + request := httptest.NewRequest(http.MethodGet, "/rpc", nil) + request.Header.Set("Origin", "http://evil.example") + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, request) + if recorder.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusForbidden) + } + }) + + t.Run("options preflight", func(t *testing.T) { + request := httptest.NewRequest(http.MethodOptions, "/rpc", nil) + request.Header.Set("Origin", "http://127.0.0.1:3000") + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, request) + if recorder.Code != http.StatusNoContent { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusNoContent) + } + }) +} + func TestNetworkServerHTTPRPCAndCORS(t *testing.T) { - server := newTestNetworkServer(t) + server := newTestNetworkServer(t, NetworkServerOptions{}) testContext, cancel := context.WithCancel(context.Background()) defer cancel() @@ -69,14 +155,20 @@ func TestNetworkServerHTTPRPCAndCORS(t *testing.T) { listenAddress := waitForNetworkAddress(t, server) requestBody := strings.NewReader(`{"jsonrpc":"2.0","id":"http-1","method":"gateway.ping","params":{}}`) - response, err := http.Post("http://"+listenAddress+"/rpc", "application/json", requestBody) + request, err := http.NewRequest(http.MethodPost, "http://"+listenAddress+"/rpc", requestBody) + if err != nil { + t.Fatalf("new request: %v", err) + } + request.Header.Set("Origin", "http://localhost:3000") + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) if err != nil { t.Fatalf("post /rpc: %v", err) } defer response.Body.Close() - if response.Header.Get("Access-Control-Allow-Origin") != "*" { - t.Fatalf("cors allow origin = %q, want %q", response.Header.Get("Access-Control-Allow-Origin"), "*") + if response.Header.Get("Access-Control-Allow-Origin") != "http://localhost:3000" { + t.Fatalf("cors allow origin = %q, want %q", response.Header.Get("Access-Control-Allow-Origin"), "http://localhost:3000") } var rpcResponse protocol.JSONRPCResponse @@ -93,27 +185,123 @@ func TestNetworkServerHTTPRPCAndCORS(t *testing.T) { if resultFrame.Type != FrameTypeAck || resultFrame.Action != FrameActionPing { t.Fatalf("result frame = %#v, want ping ack", resultFrame) } +} + +func TestNetworkServerRejectsDisallowedCORSOrigin(t *testing.T) { + server := newTestNetworkServer(t, NetworkServerOptions{}) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) - optionsRequest, err := http.NewRequest(http.MethodOptions, "http://"+listenAddress+"/rpc", nil) + requestBody := strings.NewReader(`{"jsonrpc":"2.0","id":"http-1","method":"gateway.ping","params":{}}`) + request, err := http.NewRequest(http.MethodPost, "http://"+listenAddress+"/rpc", requestBody) if err != nil { - t.Fatalf("new options request: %v", err) + t.Fatalf("new request: %v", err) } - optionsResponse, err := http.DefaultClient.Do(optionsRequest) + request.Header.Set("Origin", "http://evil.example") + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) if err != nil { - t.Fatalf("options /rpc: %v", err) - } - defer optionsResponse.Body.Close() - - if optionsResponse.StatusCode != http.StatusNoContent { - t.Fatalf("options status = %d, want %d", optionsResponse.StatusCode, http.StatusNoContent) + t.Fatalf("post /rpc: %v", err) } - if optionsResponse.Header.Get("Access-Control-Allow-Methods") == "" { - t.Fatal("expected CORS allow methods header on OPTIONS response") + defer response.Body.Close() + if response.StatusCode != http.StatusForbidden { + t.Fatalf("status = %d, want %d", response.StatusCode, http.StatusForbidden) } } +func TestNetworkServerRPCErrorBranches(t *testing.T) { + server := newTestNetworkServer(t, NetworkServerOptions{MaxRequestBytes: 16}) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + + t.Run("method not allowed", func(t *testing.T) { + response, err := http.Get("http://" + listenAddress + "/rpc") + if err != nil { + t.Fatalf("get /rpc: %v", err) + } + defer response.Body.Close() + if response.StatusCode != http.StatusMethodNotAllowed { + t.Fatalf("status = %d, want %d", response.StatusCode, http.StatusMethodNotAllowed) + } + }) + + t.Run("invalid json", func(t *testing.T) { + request, err := http.NewRequest(http.MethodPost, "http://"+listenAddress+"/rpc", strings.NewReader("{bad")) + if err != nil { + t.Fatalf("new request: %v", err) + } + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + if err != nil { + t.Fatalf("post /rpc: %v", err) + } + defer response.Body.Close() + var rpcResponse protocol.JSONRPCResponse + if err := json.NewDecoder(response.Body).Decode(&rpcResponse); err != nil { + t.Fatalf("decode rpc error response: %v", err) + } + if rpcResponse.Error == nil || rpcResponse.Error.Code != protocol.JSONRPCCodeParseError { + t.Fatalf("rpc error = %#v, want parse error", rpcResponse.Error) + } + }) + + t.Run("oversized request", func(t *testing.T) { + request, err := http.NewRequest( + http.MethodPost, + "http://"+listenAddress+"/rpc", + strings.NewReader(`{"jsonrpc":"2.0","id":"x","method":"gateway.ping","params":{}}`), + ) + if err != nil { + t.Fatalf("new request: %v", err) + } + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + if err != nil { + t.Fatalf("post /rpc: %v", err) + } + defer response.Body.Close() + var rpcResponse protocol.JSONRPCResponse + if err := json.NewDecoder(response.Body).Decode(&rpcResponse); err != nil { + t.Fatalf("decode rpc error response: %v", err) + } + if rpcResponse.Error == nil || rpcResponse.Error.Code != protocol.JSONRPCCodeParseError { + t.Fatalf("rpc error = %#v, want parse error", rpcResponse.Error) + } + }) +} + func TestNetworkServerWebSocketAndSSEPing(t *testing.T) { - server := newTestNetworkServer(t) + server := newTestNetworkServer(t, NetworkServerOptions{}) testContext, cancel := context.WithCancel(context.Background()) defer cancel() @@ -132,7 +320,7 @@ func TestNetworkServerWebSocketAndSSEPing(t *testing.T) { listenAddress := waitForNetworkAddress(t, server) wsURL := "ws://" + listenAddress + "/ws" - wsConn, err := websocket.Dial(wsURL, "", "http://localhost/") + wsConn, err := websocket.Dial(wsURL, "", "http://localhost:3000") if err != nil { t.Fatalf("dial websocket: %v", err) } @@ -142,29 +330,17 @@ func TestNetworkServerWebSocketAndSSEPing(t *testing.T) { t.Fatalf("send websocket ping request: %v", err) } - _ = wsConn.SetReadDeadline(time.Now().Add(2 * time.Second)) - var wsRawResponse string - for attempt := 0; attempt < 5; attempt++ { - if err := websocket.Message.Receive(wsConn, &wsRawResponse); err != nil { - t.Fatalf("receive websocket message: %v", err) - } - var rpcResponse protocol.JSONRPCResponse - if err := json.Unmarshal([]byte(wsRawResponse), &rpcResponse); err == nil && len(rpcResponse.JSONRPC) > 0 { - if rpcResponse.Error != nil { - t.Fatalf("unexpected websocket rpc error: %+v", rpcResponse.Error) - } - resultFrame, err := decodeJSONRPCResultFrame(rpcResponse) - if err != nil { - t.Fatalf("decode websocket result frame: %v", err) - } - if resultFrame.Action != FrameActionPing { - t.Fatalf("websocket action = %q, want %q", resultFrame.Action, FrameActionPing) - } - break - } + ackFrame := receiveWSAckFrame(t, wsConn) + if ackFrame.Action != FrameActionPing { + t.Fatalf("websocket action = %q, want %q", ackFrame.Action, FrameActionPing) } - sseResponse, err := http.Get("http://" + listenAddress + "/sse?method=gateway.ping&id=sse-1") + sseRequest, err := http.NewRequest(http.MethodGet, "http://"+listenAddress+"/sse?method=gateway.ping&id=sse-1", nil) + if err != nil { + t.Fatalf("new sse request: %v", err) + } + sseRequest.Header.Set("Origin", "app://desktop-client") + sseResponse, err := http.DefaultClient.Do(sseRequest) if err != nil { t.Fatalf("get /sse: %v", err) } @@ -173,51 +349,177 @@ func TestNetworkServerWebSocketAndSSEPing(t *testing.T) { t.Fatalf("sse status = %d, want %d", sseResponse.StatusCode, http.StatusOK) } - sseReader := bufio.NewReader(sseResponse.Body) - resultFound := false - currentEvent := "" - timeout := time.After(3 * time.Second) - for !resultFound { + resultFrame := readSSEResultFrame(t, sseResponse.Body) + if resultFrame.Action != FrameActionPing { + t.Fatalf("sse action = %q, want %q", resultFrame.Action, FrameActionPing) + } +} + +func TestNetworkServerWebSocketOriginRejected(t *testing.T) { + server := newTestNetworkServer(t, NetworkServerOptions{}) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) select { - case <-timeout: - t.Fatal("timed out waiting for sse result event") + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + wsURL := "ws://" + listenAddress + "/ws" + if _, err := websocket.Dial(wsURL, "", "http://evil.example"); err == nil { + t.Fatal("expected websocket handshake to reject disallowed origin") + } +} + +func TestNetworkServerWebSocketReadTimeoutDoesNotKillIdleConnection(t *testing.T) { + server := newTestNetworkServer(t, NetworkServerOptions{ + ReadTimeout: 80 * time.Millisecond, + HeartbeatInterval: 30 * time.Millisecond, + }) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + wsConn, err := websocket.Dial("ws://"+listenAddress+"/ws", "", "http://localhost:3000") + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + t.Cleanup(func() { _ = wsConn.Close() }) + + time.Sleep(300 * time.Millisecond) + + if err := websocket.Message.Send(wsConn, `{"jsonrpc":"2.0","id":"ws-idle","method":"gateway.ping","params":{}}`); err != nil { + t.Fatalf("send ping after idle: %v", err) + } + ackFrame := receiveWSAckFrame(t, wsConn) + if ackFrame.RequestID != "ws-idle" { + t.Fatalf("request_id = %q, want %q", ackFrame.RequestID, "ws-idle") + } +} + +func TestNetworkServerWebSocketDispatchContextCancelledOnShutdown(t *testing.T) { + originalDispatch := dispatchRPCRequestFn + t.Cleanup(func() { dispatchRPCRequestFn = originalDispatch }) + + started := make(chan struct{}, 1) + cancelled := make(chan struct{}, 1) + dispatchRPCRequestFn = func(ctx context.Context, request protocol.JSONRPCRequest, runtimePort RuntimePort) protocol.JSONRPCResponse { + select { + case started <- struct{}{}: + default: + } + <-ctx.Done() + select { + case cancelled <- struct{}{}: default: - line, readErr := sseReader.ReadString('\n') - if readErr != nil { - t.Fatalf("read sse line: %v", readErr) - } - trimmedLine := strings.TrimSpace(line) - if trimmedLine == "" { - continue - } - if strings.HasPrefix(trimmedLine, "event:") { - currentEvent = strings.TrimSpace(strings.TrimPrefix(trimmedLine, "event:")) - continue - } - if strings.HasPrefix(trimmedLine, "data:") && currentEvent == "result" { - rawData := strings.TrimSpace(strings.TrimPrefix(trimmedLine, "data:")) - var rpcResponse protocol.JSONRPCResponse - if err := json.Unmarshal([]byte(rawData), &rpcResponse); err != nil { - t.Fatalf("decode sse result: %v", err) - } - if rpcResponse.Error != nil { - t.Fatalf("unexpected sse rpc error: %+v", rpcResponse.Error) - } - resultFrame, err := decodeJSONRPCResultFrame(rpcResponse) - if err != nil { - t.Fatalf("decode sse result frame: %v", err) - } - if resultFrame.Action != FrameActionPing { - t.Fatalf("sse action = %q, want %q", resultFrame.Action, FrameActionPing) - } - resultFound = true - } } + return protocol.NewJSONRPCErrorResponse( + json.RawMessage(`"ws-cancel"`), + protocol.NewJSONRPCError(protocol.JSONRPCCodeInternalError, "cancelled", protocol.GatewayCodeInternalError), + ) + } + + server := newTestNetworkServer(t, NetworkServerOptions{}) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, nil) + }() + listenAddress := waitForNetworkAddress(t, server) + + wsConn, err := websocket.Dial("ws://"+listenAddress+"/ws", "", "http://localhost:3000") + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = wsConn.Close() }() + + if err := websocket.Message.Send(wsConn, `{"jsonrpc":"2.0","id":"ws-block","method":"gateway.ping","params":{}}`); err != nil { + t.Fatalf("send websocket request: %v", err) + } + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("dispatch function was not invoked") + } + + if err := server.Close(context.Background()); err != nil { + t.Fatalf("close network server: %v", err) + } + + select { + case <-cancelled: + case <-time.After(time.Second): + t.Fatal("expected websocket dispatch context to be cancelled on shutdown") + } + + select { + case serveErr := <-serveDone: + if serveErr != nil { + t.Fatalf("serve returned error: %v", serveErr) + } + case <-time.After(2 * time.Second): + t.Fatal("serve did not exit after close") + } +} + +func TestNetworkServerSSEErrorBranches(t *testing.T) { + server := newTestNetworkServer(t, NetworkServerOptions{}) + + t.Run("method not allowed", func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/sse", nil) + server.handleSSERequest(recorder, request, nil) + if recorder.Code != http.StatusMethodNotAllowed { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusMethodNotAllowed) + } + }) + + t.Run("streaming unsupported", func(t *testing.T) { + writer := &noFlushResponseWriter{header: make(http.Header)} + request := httptest.NewRequest(http.MethodGet, "/sse", nil) + server.handleSSERequest(writer, request, nil) + if writer.status != http.StatusInternalServerError { + t.Fatalf("status = %d, want %d", writer.status, http.StatusInternalServerError) + } + }) +} + +func TestDecodeJSONRPCRequestFromReaderTrailingJSON(t *testing.T) { + request, rpcErr := decodeJSONRPCRequestFromReader(strings.NewReader(`{"jsonrpc":"2.0","id":"x","method":"gateway.ping"} {"extra":1}`)) + if rpcErr == nil { + t.Fatalf("expected parse error, got request %#v", request) + } + if rpcErr.Code != protocol.JSONRPCCodeParseError { + t.Fatalf("rpc error code = %d, want %d", rpcErr.Code, protocol.JSONRPCCodeParseError) } } func TestNetworkServerCloseInterruptsStreams(t *testing.T) { - server := newTestNetworkServer(t) + server := newTestNetworkServer(t, NetworkServerOptions{}) testContext, cancel := context.WithCancel(context.Background()) defer cancel() @@ -227,8 +529,7 @@ func TestNetworkServerCloseInterruptsStreams(t *testing.T) { }() listenAddress := waitForNetworkAddress(t, server) - wsURL := "ws://" + listenAddress + "/ws" - wsConn, err := websocket.Dial(wsURL, "", "http://localhost/") + wsConn, err := websocket.Dial("ws://"+listenAddress+"/ws", "", "http://localhost:3000") if err != nil { t.Fatalf("dial websocket: %v", err) } @@ -240,9 +541,7 @@ func TestNetworkServerCloseInterruptsStreams(t *testing.T) { } defer sseResponse.Body.Close() - closeContext, closeCancel := context.WithTimeout(context.Background(), time.Second) - defer closeCancel() - if err := server.Close(closeContext); err != nil { + if err := server.Close(context.Background()); err != nil { t.Fatalf("close network server: %v", err) } @@ -275,17 +574,29 @@ func TestNetworkServerCloseInterruptsStreams(t *testing.T) { } // newTestNetworkServer 创建默认测试网络服务实例，统一收敛测试参数。 -func newTestNetworkServer(t *testing.T) *NetworkServer { +func newTestNetworkServer(t *testing.T, overrides NetworkServerOptions) *NetworkServer { t.Helper() - server, err := NewNetworkServer(NetworkServerOptions{ - ListenAddress: "127.0.0.1:0", - Logger: log.New(io.Discard, "", 0), - HeartbeatInterval: 100 * time.Millisecond, - ReadTimeout: 2 * time.Second, - WriteTimeout: 2 * time.Second, - ShutdownTimeout: 500 * time.Millisecond, - }) + if strings.TrimSpace(overrides.ListenAddress) == "" { + overrides.ListenAddress = "127.0.0.1:0" + } + if overrides.Logger == nil { + overrides.Logger = log.New(io.Discard, "", 0) + } + if overrides.HeartbeatInterval <= 0 { + overrides.HeartbeatInterval = 100 * time.Millisecond + } + if overrides.ReadTimeout <= 0 { + overrides.ReadTimeout = 2 * time.Second + } + if overrides.WriteTimeout <= 0 { + overrides.WriteTimeout = 2 * time.Second + } + if overrides.ShutdownTimeout <= 0 { + overrides.ShutdownTimeout = 500 * time.Millisecond + } + + server, err := NewNetworkServer(overrides) if err != nil { t.Fatalf("new network server: %v", err) } @@ -312,3 +623,93 @@ func waitForNetworkAddress(t *testing.T, server *NetworkServer) string { } } } + +// receiveWSAckFrame 连续读取 WS 消息直到拿到 JSON-RPC ACK 结果帧。 +func receiveWSAckFrame(t *testing.T, wsConn *websocket.Conn) MessageFrame { + t.Helper() + _ = wsConn.SetReadDeadline(time.Now().Add(2 * time.Second)) + for attempt := 0; attempt < 12; attempt++ { + var rawResponse string + if err := websocket.Message.Receive(wsConn, &rawResponse); err != nil { + t.Fatalf("receive websocket message: %v", err) + } + var rpcResponse protocol.JSONRPCResponse + if err := json.Unmarshal([]byte(rawResponse), &rpcResponse); err != nil { + continue + } + if rpcResponse.JSONRPC == "" { + continue + } + if rpcResponse.Error != nil { + t.Fatalf("unexpected websocket rpc error: %+v", rpcResponse.Error) + } + resultFrame, err := decodeJSONRPCResultFrame(rpcResponse) + if err != nil { + t.Fatalf("decode websocket result frame: %v", err) + } + return resultFrame + } + t.Fatal("did not receive websocket ack frame") + return MessageFrame{} +} + +// readSSEResultFrame 读取 SSE result 事件并解析内部 JSON-RPC 结果帧。 +func readSSEResultFrame(t *testing.T, body io.Reader) MessageFrame { + t.Helper() + reader := bufio.NewReader(body) + currentEvent := "" + timeout := time.After(3 * time.Second) + + for { + select { + case <-timeout: + t.Fatal("timed out waiting for sse result") + default: + line, err := reader.ReadString('\n') + if err != nil { + t.Fatalf("read sse line: %v", err) + } + trimmed := strings.TrimSpace(line) + if trimmed == "" { + continue + } + if strings.HasPrefix(trimmed, "event:") { + currentEvent = strings.TrimSpace(strings.TrimPrefix(trimmed, "event:")) + continue + } + if currentEvent == "result" && strings.HasPrefix(trimmed, "data:") { + rawData := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + var rpcResponse protocol.JSONRPCResponse + if err := json.Unmarshal([]byte(rawData), &rpcResponse); err != nil { + t.Fatalf("decode sse result: %v", err) + } + resultFrame, err := decodeJSONRPCResultFrame(rpcResponse) + if err != nil { + t.Fatalf("decode sse result frame: %v", err) + } + return resultFrame + } + } + } +} + +type noFlushResponseWriter struct { + header http.Header + status int + body strings.Builder +} + +func (w *noFlushResponseWriter) Header() http.Header { + return w.header +} + +func (w *noFlushResponseWriter) Write(payload []byte) (int, error) { + if w.status == 0 { + w.status = http.StatusOK + } + return w.body.Write(payload) +} + +func (w *noFlushResponseWriter) WriteHeader(statusCode int) { + w.status = statusCode +} diff --git a/internal/gateway/rpc_dispatch_test.go b/internal/gateway/rpc_dispatch_test.go new file mode 100644 index 00000000..1b296e36 --- /dev/null +++ b/internal/gateway/rpc_dispatch_test.go @@ -0,0 +1,67 @@ +package gateway + +import ( + "context" + "encoding/json" + "testing" + + "neo-code/internal/gateway/protocol" +) + +func TestDispatchRPCRequestResultEncodeError(t *testing.T) { + originalHandlers := requestFrameHandlers + requestFrameHandlers = map[FrameAction]requestFrameHandler{ + FrameActionPing: func(frame MessageFrame) MessageFrame { + return MessageFrame{ + Type: FrameTypeAck, + Action: FrameActionPing, + RequestID: frame.RequestID, + Payload: map[string]any{ + "bad": make(chan int), + }, + } + }, + } + t.Cleanup(func() { + requestFrameHandlers = originalHandlers + }) + + response := dispatchRPCRequest(context.Background(), protocol.JSONRPCRequest{ + JSONRPC: protocol.JSONRPCVersion, + ID: json.RawMessage(`"rpc-encode-1"`), + Method: protocol.MethodGatewayPing, + Params: json.RawMessage(`{}`), + }, nil) + if response.Error == nil { + t.Fatal("expected jsonrpc internal error") + } + if response.Error.Code != protocol.JSONRPCCodeInternalError { + t.Fatalf("rpc error code = %d, want %d", response.Error.Code, protocol.JSONRPCCodeInternalError) + } + if gatewayCode := protocol.GatewayCodeFromJSONRPCError(response.Error); gatewayCode != ErrorCodeInternalError.String() { + t.Fatalf("gateway_code = %q, want %q", gatewayCode, ErrorCodeInternalError.String()) + } +} + +func TestDispatchFrameValidationBranches(t *testing.T) { + response := dispatchFrame(context.Background(), MessageFrame{ + Type: FrameType("invalid"), + }, nil) + if response.Type != FrameTypeError { + t.Fatalf("response type = %q, want %q", response.Type, FrameTypeError) + } + if response.Error == nil || response.Error.Code != ErrorCodeInvalidFrame.String() { + t.Fatalf("response error = %#v, want invalid_frame", response.Error) + } + + response = dispatchFrame(context.Background(), MessageFrame{ + Type: FrameTypeEvent, + Action: FrameActionPing, + }, nil) + if response.Type != FrameTypeError { + t.Fatalf("response type = %q, want %q", response.Type, FrameTypeError) + } + if response.Error == nil || response.Error.Code != ErrorCodeInvalidFrame.String() { + t.Fatalf("response error = %#v, want invalid_frame", response.Error) + } +} From 0e014de0d6fc311541bedb446fadd258d12a90a8 Mon Sep 17 00:00:00 2001 From: pionxe Date: Thu, 16 Apr 2026 08:53:49 +0800 Subject: [PATCH 3/4] =?UTF-8?q?fix(gateway):=20[EPIC-GW-04]=20=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=20IPv6=20=E8=B7=A8=E7=AB=99=E7=99=BD=E5=90=8D?= =?UTF-8?q?=E5=8D=95=E4=B8=8E=E5=89=AA=E8=B4=B4=E6=9D=BF=E7=A9=BA=E7=8A=B6?= =?UTF-8?q?=E6=80=81=E8=AF=AF=E5=88=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 处理 Code Review 的最后边缘场景反馈，进一步提升网络访问面的一致性与系统健壮性： 1. 网络一致性修复：将 IPv6 环回地址 (http://[::1]) 及其带端口的形态正式纳入 CORS 与 WebSocket 的 Origin 校验白名单。修复了网关允许绑定 IPv6 却拦截其同源请求的自相矛盾问题。 2. 状态机逻辑修正：修复 TUI infra 层的本地剪贴板读取逻辑。当剪贴板无图片内容时，正确返回空状态 (nil, nil) 而非抛出“该平台不支持”的异常，避免上游 TUI 抛出错误告警。 3. 文档契约对齐：在 README.md 网络端点章节补充详尽的 Origin 安全白名单声明（限 localhost / 127.0.0.1 / [::1] / app://），消除未来第三方前端接入时的调试盲区。 4. 测试同步：同步更新网络面跨域测试与 infra 层断言用例，网关核心覆盖率稳居 95.3%。 --- README.md | 4 ++++ internal/gateway/network_server.go | 2 ++ internal/gateway/network_server_test.go | 3 ++- internal/tui/infra/clipboard_native.go | 2 +- internal/tui/infra/infra_test.go | 8 ++++---- 5 files changed, 13 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 18a52d44..2b75ef6f 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,10 @@ go run ./cmd/neocode gateway --http-listen 127.0.0.1:8080 - `GET /ws`：WebSocket 流式入口（含心跳） - `GET /sse`：SSE 流式入口（MVP 默认触发 `gateway.ping`，含心跳） +安全限制：为防止跨站攻击，网关网络面默认开启严格的 Origin 校验。当前仅允许 +`http://localhost`、`http://127.0.0.1`、`http://[::1]` 以及 `app://` 前缀来源连入； +非允许来源的跨域调用会被拦截并返回 `403`。 + URL Scheme 派发骨架命令（EPIC-GW-02A）： ```bash diff --git a/internal/gateway/network_server.go b/internal/gateway/network_server.go index 92e4aff9..f3064592 100644 --- a/internal/gateway/network_server.go +++ b/internal/gateway/network_server.go @@ -647,6 +647,8 @@ func isAllowedControlPlaneOrigin(origin string) bool { normalizedOrigin == "http://localhost", strings.HasPrefix(normalizedOrigin, "http://127.0.0.1:"), normalizedOrigin == "http://127.0.0.1", + strings.HasPrefix(normalizedOrigin, "http://[::1]:"), + normalizedOrigin == "http://[::1]", strings.HasPrefix(normalizedOrigin, "app://"): return true default: diff --git a/internal/gateway/network_server_test.go b/internal/gateway/network_server_test.go index f5789d97..db7ae2b6 100644 --- a/internal/gateway/network_server_test.go +++ b/internal/gateway/network_server_test.go @@ -55,6 +55,8 @@ func TestOriginAllowlist(t *testing.T) { "http://localhost", "http://127.0.0.1:5173", "http://127.0.0.1", + "http://[::1]:3000", + "http://[::1]", "app://desktop-client", } for _, origin := range allowed { @@ -66,7 +68,6 @@ func TestOriginAllowlist(t *testing.T) { disallowed := []string{ "", "https://localhost:3000", - "http://[::1]:3000", "http://evil.example.com", "file://local", } diff --git a/internal/tui/infra/clipboard_native.go b/internal/tui/infra/clipboard_native.go index 3f8f15e0..604098b2 100644 --- a/internal/tui/infra/clipboard_native.go +++ b/internal/tui/infra/clipboard_native.go @@ -40,7 +40,7 @@ func ReadClipboardImage() ([]byte, error) { } data := clipboard.Read(clipboard.FmtImage) if len(data) == 0 { - return nil, errClipboardImageUnsupported + return nil, nil } return data, nil } diff --git a/internal/tui/infra/infra_test.go b/internal/tui/infra/infra_test.go index 9e39bb45..4dfc1b9b 100644 --- a/internal/tui/infra/infra_test.go +++ b/internal/tui/infra/infra_test.go @@ -403,11 +403,11 @@ func TestClipboardFallbackFunctions(t *testing.T) { t.Fatalf("expected clipboard text or an error, got empty success result") } data, err := ReadClipboardImage() - if err != errClipboardImageUnsupported { - t.Fatalf("expected unsupported image error, got %v", err) + if err != nil && err != errClipboardImageUnsupported { + t.Fatalf("expected nil or unsupported image error, got %v", err) } - if data != nil { - t.Fatalf("expected nil image data on unsupported platform") + if err == nil && len(data) == 0 && data != nil { + t.Fatalf("expected nil for empty clipboard image state, got zero-length slice") } } From d6c6296f7eed70b0cbfb32808bfa93407ab2079a Mon Sep 17 00:00:00 2001 From: pionxe Date: Thu, 16 Apr 2026 09:53:10 +0800 Subject: [PATCH 4/4] =?UTF-8?q?fix(gateway):=20[EPIC-GW-04]=20=E6=8F=90?= =?UTF-8?q?=E5=8D=87=E7=BD=91=E7=BB=9C=E9=9D=A2=E5=90=AF=E5=8A=A8=E9=99=8D?= =?UTF-8?q?=E7=BA=A7=E5=AE=B9=E9=94=99=E4=B8=8E=E7=8E=AF=E5=9B=9E=E6=A0=A1?= =?UTF-8?q?=E9=AA=8C=E5=AE=89=E5=85=A8=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 解决网络层边缘状态与可用性问题，保障本地核心服务的极高存活率： 1. 优雅降级 (Graceful Degradation)：解耦网络面与核心 IPC 通道的生命周期。当 HTTP 监听失败（如 8080 端口冲突被占用）时，不再触发进程 Fatal，而是转为输出 Warning 告警，确保网关底层 IPC 继续为本地应用提供服务。 2. 严格的环回校验 (Loopback Resolution)：消除了原先仅依赖字符串比对 localhost 的安全盲区。现引入真实的 DNS 解析，强制验证域名对应的所有 IP 均为 Loopback 地址，防御特定环境下的解析污染。 3. 契约文档对齐 (Docs)：在 README 中明确注明 Origin 安全拦截机制的生效范围。澄清了非浏览器客户端（如 cURL、Postman、本地后端脚本）等不携带 Origin 头的请求会被正常放行，消除调用方的调试疑虑。 --- README.md | 2 + internal/cli/gateway_commands.go | 26 +++++------- internal/cli/root_test.go | 40 +++++++++++++++++-- internal/gateway/network_server.go | 18 +++++++-- .../gateway/network_server_additional_test.go | 39 ++++++++++++++++++ 5 files changed, 101 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index 2b75ef6f..c1b4dd78 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,8 @@ go run ./cmd/neocode gateway --http-listen 127.0.0.1:8080 安全限制：为防止跨站攻击，网关网络面默认开启严格的 Origin 校验。当前仅允许 `http://localhost`、`http://127.0.0.1`、`http://[::1]` 以及 `app://` 前缀来源连入； 非允许来源的跨域调用会被拦截并返回 `403`。 +注：上述白名单机制仅针对携带 `Origin` 头的浏览器跨站请求生效。若请求不携带 `Origin` 头 +（例如 `cURL`、Postman 或本地后端脚本直连），网关默认放行。 URL Scheme 派发骨架命令（EPIC-GW-02A）： diff --git a/internal/cli/gateway_commands.go b/internal/cli/gateway_commands.go index 1003bc81..13101328 100644 --- a/internal/cli/gateway_commands.go +++ b/internal/cli/gateway_commands.go @@ -148,26 +148,18 @@ func defaultGatewayCommandRunner(ctx context.Context, options gatewayCommandOpti logger.Printf("gateway ipc listen address: %s", ipcServer.ListenAddress()) logger.Printf("gateway network listen address: %s", networkServer.ListenAddress()) - serveContext, cancel := context.WithCancel(signalContext) - defer cancel() - - serveErrorChannel := make(chan error, 2) - go func() { - serveErrorChannel <- ipcServer.Serve(serveContext, nil) - }() go func() { - serveErrorChannel <- networkServer.Serve(serveContext, nil) + serveErr := networkServer.Serve(signalContext, nil) + if serveErr != nil && signalContext.Err() == nil { + logger.Printf( + "warning: HTTP server failed to start on %s (port in use?), but IPC server is still running: %v", + networkServer.ListenAddress(), + serveErr, + ) + } }() - var firstError error - for index := 0; index < 2; index++ { - serveError := <-serveErrorChannel - if serveError != nil && firstError == nil { - firstError = serveError - cancel() - } - } - return firstError + return ipcServer.Serve(signalContext, nil) } // defaultNewGatewayServer 创建默认网关服务实例，供命令层启动流程调用。 diff --git a/internal/cli/root_test.go b/internal/cli/root_test.go index dc65c18a..ad473be2 100644 --- a/internal/cli/root_test.go +++ b/internal/cli/root_test.go @@ -255,9 +255,6 @@ func TestDefaultGatewayCommandRunnerSuccess(t *testing.T) { if !server.closeCalled { t.Fatal("expected server Close to be called") } - if !networkServer.serveCalled { - t.Fatal("expected network server Serve to be called") - } if !networkServer.closeCalled { t.Fatal("expected network server Close to be called") } @@ -322,6 +319,43 @@ func TestDefaultGatewayCommandRunnerReturnsServeError(t *testing.T) { } } +func TestDefaultGatewayCommandRunnerDegradesWhenNetworkServeFails(t *testing.T) { + originalNewGatewayServer := newGatewayServer + originalNewGatewayNetwork := newGatewayNetwork + t.Cleanup(func() { newGatewayServer = originalNewGatewayServer }) + t.Cleanup(func() { newGatewayNetwork = originalNewGatewayNetwork }) + + ipcServer := &stubGatewayServer{listenAddress: "stub://gateway"} + newGatewayServer = func(options gateway.ServerOptions) (gatewayServer, error) { + return ipcServer, nil + } + networkServer := &stubGatewayServer{ + listenAddress: "127.0.0.1:8080", + serveErr: errors.New("bind: address already in use"), + } + newGatewayNetwork = func(options gateway.NetworkServerOptions) (gatewayNetworkServer, error) { + return networkServer, nil + } + + err := defaultGatewayCommandRunner(context.Background(), gatewayCommandOptions{ + ListenAddress: "stub://gateway", + HTTPAddress: "127.0.0.1:8080", + LogLevel: "info", + }) + if err != nil { + t.Fatalf("expected graceful degradation on network serve error, got %v", err) + } + if !ipcServer.serveCalled { + t.Fatal("expected ipc server Serve to be called") + } + if !ipcServer.closeCalled { + t.Fatal("expected ipc server Close to be called") + } + if !networkServer.closeCalled { + t.Fatal("expected network server Close to be called") + } +} + func TestDefaultGatewayCommandRunnerReturnsNetworkConstructorError(t *testing.T) { originalNewGatewayServer := newGatewayServer originalNewGatewayNetwork := newGatewayNetwork diff --git a/internal/gateway/network_server.go b/internal/gateway/network_server.go index f3064592..da3cdf6b 100644 --- a/internal/gateway/network_server.go +++ b/internal/gateway/network_server.go @@ -40,6 +40,7 @@ const ( var ( resolveNetworkListenAddressFn = ResolveNetworkListenAddress + lookupHostIPsFn = net.LookupIP dispatchRPCRequestFn = dispatchRPCRequest ) @@ -160,13 +161,22 @@ func validateLoopbackListenAddress(address string) error { if normalizedHost == "" { return fmt.Errorf("invalid --http-listen %q: host must be loopback", address) } - if strings.EqualFold(normalizedHost, "localhost") { + if ip := net.ParseIP(normalizedHost); ip != nil { + if !ip.IsLoopback() { + return fmt.Errorf("invalid --http-listen %q: host must be loopback", address) + } return nil } - ip := net.ParseIP(normalizedHost) - if ip == nil || !ip.IsLoopback() { - return fmt.Errorf("invalid --http-listen %q: host must be loopback", address) + resolvedHostIPs, lookupErr := lookupHostIPsFn(normalizedHost) + if lookupErr != nil || len(resolvedHostIPs) == 0 { + return fmt.Errorf("invalid --http-listen %q: host must resolve to loopback addresses", address) + } + + for _, resolvedIP := range resolvedHostIPs { + if resolvedIP == nil || !resolvedIP.IsLoopback() { + return fmt.Errorf("invalid --http-listen %q: host must be loopback", address) + } } return nil } diff --git a/internal/gateway/network_server_additional_test.go b/internal/gateway/network_server_additional_test.go index 5673fa20..c1bf9f8d 100644 --- a/internal/gateway/network_server_additional_test.go +++ b/internal/gateway/network_server_additional_test.go @@ -72,6 +72,45 @@ func TestValidateLoopbackListenAddressParseFailure(t *testing.T) { } } +func TestValidateLoopbackListenAddressHostLookup(t *testing.T) { + originalLookup := lookupHostIPsFn + t.Cleanup(func() { + lookupHostIPsFn = originalLookup + }) + + t.Run("hostname resolves to loopback addresses", func(t *testing.T) { + lookupHostIPsFn = func(host string) ([]net.IP, error) { + return []net.IP{ + net.ParseIP("127.0.0.1"), + net.ParseIP("::1"), + }, nil + } + if err := validateLoopbackListenAddress("localhost:8080"); err != nil { + t.Fatalf("expected loopback hostname to pass, got %v", err) + } + }) + + t.Run("hostname resolves to non-loopback address", func(t *testing.T) { + lookupHostIPsFn = func(host string) ([]net.IP, error) { + return []net.IP{net.ParseIP("192.168.1.10")}, nil + } + err := validateLoopbackListenAddress("localhost:8080") + if err == nil || !strings.Contains(err.Error(), "host must be loopback") { + t.Fatalf("expected non-loopback hostname rejection, got %v", err) + } + }) + + t.Run("hostname lookup failed", func(t *testing.T) { + lookupHostIPsFn = func(host string) ([]net.IP, error) { + return nil, errors.New("lookup failed") + } + err := validateLoopbackListenAddress("localhost:8080") + if err == nil || !strings.Contains(err.Error(), "host must resolve to loopback addresses") { + t.Fatalf("expected lookup failure rejection, got %v", err) + } + }) +} + func TestNetworkServerServeErrorBranches(t *testing.T) { t.Run("listen error", func(t *testing.T) { server, err := NewNetworkServer(NetworkServerOptions{