diff --git a/core/cmd/dms/server_client.go b/core/cmd/dms/server_client.go index 56c4c6b42..56a9023ac 100644 --- a/core/cmd/dms/server_client.go +++ b/core/cmd/dms/server_client.go @@ -4,9 +4,11 @@ import ( "bufio" "encoding/json" "fmt" + "io" "net" "os" "path/filepath" + "time" "github.com/AvengeMedia/DankMaterialShell/core/internal/server" "github.com/AvengeMedia/DankMaterialShell/core/internal/utils" @@ -17,18 +19,37 @@ import ( // overhead in the line-delimited IPC response. const maxIPCMessageSize = 96 * 1024 * 1024 -func sendServerRequest(req ipc.Request) (*ipc.Response[any], error) { - socketPath := getServerSocketPath() +// serverHandshakeTimeout bounds connecting, the greeting and sending the +// request, so a stuck server can't hang the CLI. Replies get no limit because +// some requests (matugen queue --wait) legitimately take long. +var serverHandshakeTimeout = 2 * time.Second - conn, err := net.Dial("unix", socketPath) +func dialServer() (net.Conn, *bufio.Scanner, error) { + conn, err := net.DialTimeout("unix", getServerSocketPath(), serverHandshakeTimeout) if err != nil { - return nil, fmt.Errorf("failed to connect to server (is it running?): %w", err) + return nil, nil, fmt.Errorf("failed to connect to server (is it running?): %w", err) } - defer conn.Close() + _ = conn.SetDeadline(time.Now().Add(serverHandshakeTimeout)) scanner := bufio.NewScanner(conn) scanner.Buffer(make([]byte, bufio.MaxScanTokenSize), maxIPCMessageSize) - scanner.Scan() // discard initial capabilities message + if !scanner.Scan() { // initial capabilities message + conn.Close() + err := scanner.Err() + if err == nil { + err = io.EOF + } + return nil, nil, fmt.Errorf("server did not respond: %w", err) + } + return conn, scanner, nil +} + +func sendServerRequest(req ipc.Request) (*ipc.Response[any], error) { + conn, scanner, err := dialServer() + if err != nil { + return nil, err + } + defer conn.Close() reqData, err := json.Marshal(req) if err != nil { @@ -43,6 +64,7 @@ func sendServerRequest(req ipc.Request) (*ipc.Response[any], error) { return nil, fmt.Errorf("failed to write newline: %w", err) } + _ = conn.SetDeadline(time.Time{}) if !scanner.Scan() { return nil, fmt.Errorf("failed to read response") } @@ -58,18 +80,12 @@ func sendServerRequest(req ipc.Request) (*ipc.Response[any], error) { // sendServerRequestFireAndForget sends a request without waiting for a response. // Useful for commands that trigger UI or async operations. func sendServerRequestFireAndForget(req ipc.Request) error { - socketPath := getServerSocketPath() - - conn, err := net.Dial("unix", socketPath) + conn, _, err := dialServer() if err != nil { - return fmt.Errorf("failed to connect to server (is it running?): %w", err) + return err } defer conn.Close() - scanner := bufio.NewScanner(conn) - scanner.Buffer(make([]byte, bufio.MaxScanTokenSize), maxIPCMessageSize) - scanner.Scan() // discard initial capabilities message - reqData, err := json.Marshal(req) if err != nil { return fmt.Errorf("failed to marshal request: %w", err) diff --git a/core/cmd/dms/server_client_test.go b/core/cmd/dms/server_client_test.go new file mode 100644 index 000000000..479099e6e --- /dev/null +++ b/core/cmd/dms/server_client_test.go @@ -0,0 +1,86 @@ +package main + +import ( + "bufio" + "net" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/AvengeMedia/dankgo/ipc" +) + +func fakeServer(t *testing.T, handle func(net.Conn)) { + t.Helper() + dir := t.TempDir() + t.Setenv("XDG_RUNTIME_DIR", dir) + t.Setenv("WAYLAND_DISPLAY", "") + + ln, err := net.Listen("unix", filepath.Join(dir, "danklinux.sock")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { ln.Close() }) + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func() { + defer conn.Close() + handle(conn) + }() + } + }() +} + +func TestServerRequestTimesOutWithoutGreeting(t *testing.T) { + old := serverHandshakeTimeout + serverHandshakeTimeout = 100 * time.Millisecond + t.Cleanup(func() { serverHandshakeTimeout = old }) + + fakeServer(t, func(conn net.Conn) { time.Sleep(time.Second) }) + + start := time.Now() + if _, err := sendServerRequest(ipc.Request{Method: "test"}); err == nil { + t.Fatal("sendServerRequest: expected an error from a silent server") + } + if err := sendServerRequestFireAndForget(ipc.Request{Method: "test"}); err == nil { + t.Fatal("sendServerRequestFireAndForget: expected an error from a silent server") + } + if elapsed := time.Since(start); elapsed > 900*time.Millisecond { + t.Fatalf("requests took %v, want them to give up after the handshake timeout", elapsed) + } +} + +func TestServerRequestWaitsForSlowReply(t *testing.T) { + old := serverHandshakeTimeout + serverHandshakeTimeout = 100 * time.Millisecond + t.Cleanup(func() { serverHandshakeTimeout = old }) + + fakeServer(t, func(conn net.Conn) { + conn.Write([]byte("{}\n")) + bufio.NewReader(conn).ReadBytes('\n') + time.Sleep(300 * time.Millisecond) + conn.Write([]byte(`{"id":1,"result":"ok"}` + "\n")) + }) + + resp, err := sendServerRequest(ipc.Request{ID: 1, Method: "test"}) + if err != nil { + t.Fatalf("sendServerRequest: %v", err) + } + if resp.Result == nil || *resp.Result != "ok" { + t.Fatalf("unexpected response: %+v", resp) + } +} + +func TestServerRequestReportsClosedConnection(t *testing.T) { + fakeServer(t, func(conn net.Conn) {}) + + _, err := sendServerRequest(ipc.Request{Method: "test"}) + if err == nil || strings.Contains(err.Error(), "%!") { + t.Fatalf("want a readable error for a server that hangs up, got %v", err) + } +}