fix(cli): time out when the DMS server doesn't answer (#3791)

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)

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)
+	}
+}