patterm/internal/app/daemon_net.go

package app

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/hjbdev/patterm/internal/mcp"
	"github.com/hjbdev/patterm/internal/preset"
	"github.com/hjbdev/patterm/internal/protocol"
)

type DaemonOptions struct {
	ProjectDir  string
	SocketPath  string
	PidPath     string
	ListenAddr  string
	Token       string
	TokenOut    io.Writer
	ListenReady chan string
	Cols        uint16
	Rows        uint16
}

type DaemonStatus struct {
	PID      int
	Socket   string
	Projects []protocol.Project
}

func RuntimeDaemonPaths() (socketPath, pidPath string, err error) {
	base := os.Getenv("XDG_RUNTIME_DIR")
	if base == "" {
		base = os.TempDir()
	}
	dir := filepath.Join(base, "patterm")
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return "", "", err
	}
	return filepath.Join(dir, "daemon.sock"), filepath.Join(dir, "daemon.pid"), nil
}

func RunDaemon(ctx context.Context, opts DaemonOptions) error {
	if opts.ProjectDir == "" {
		cwd, err := os.Getwd()
		if err != nil {
			return err
		}
		opts.ProjectDir = cwd
	}
	if opts.SocketPath == "" || opts.PidPath == "" {
		socket, pid, err := RuntimeDaemonPaths()
		if err != nil {
			return err
		}
		if opts.SocketPath == "" {
			opts.SocketPath = socket
		}
		if opts.PidPath == "" {
			opts.PidPath = pid
		}
	}
	if opts.Cols == 0 {
		opts.Cols = 80
	}
	if opts.Rows == 0 {
		opts.Rows = 24
	}
	lockPath, err := prepareDaemonSocket(opts.SocketPath, opts.PidPath)
	if err != nil {
		return err
	}
	defer os.Remove(lockPath)
	ln, err := net.Listen("unix", opts.SocketPath)
	if err != nil {
		return fmt.Errorf("daemon: listen %s: %w", opts.SocketPath, err)
	}
	defer ln.Close()
	defer os.Remove(opts.SocketPath)
	if err := os.Chmod(opts.SocketPath, 0o600); err != nil {
		return err
	}
	if err := os.WriteFile(opts.PidPath, []byte(strconv.Itoa(os.Getpid())+"\n"), 0o600); err != nil {
		return err
	}
	defer os.Remove(opts.PidPath)

	presets, err := preset.Load()
	if err != nil {
		return fmt.Errorf("daemon: load presets: %w", err)
	}
	appSettings, _, err := loadSettings()
	if err != nil {
		logf("daemon settings load: %v", err)
	}
	mcpSrv, err := mcp.Start()
	if err != nil {
		return fmt.Errorf("daemon: mcp start: %w", err)
	}
	defer mcpSrv.Close()

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	registry := newProjectRegistry(presets, appSettings, mcpSrv, opts.Cols, opts.Rows)
	defer registry.Shutdown()
	mcpSrv.SetHost(registry)
	if _, err := registry.Open(ctx, opts.ProjectDir); err != nil {
		return err
	}

	var tcpLn net.Listener
	tcpToken := opts.Token
	if opts.ListenAddr != "" {
		addr := normalizeListenAddr(opts.ListenAddr)
		tcpToken, err = ensureDaemonToken(tcpToken)
		if err != nil {
			return err
		}
		tcpLn, err = net.Listen("tcp", addr)
		if err != nil {
			return fmt.Errorf("daemon: listen tcp %s: %w", addr, err)
		}
		defer tcpLn.Close()
		if opts.ListenReady != nil {
			select {
			case opts.ListenReady <- tcpLn.Addr().String():
			default:
			}
		}
		out := opts.TokenOut
		if out == nil {
			out = os.Stderr
		}
		fmt.Fprintf(out, "patterm daemon listening on %s\npatterm token: %s\n", tcpLn.Addr().String(), tcpToken)
	}

	var wg sync.WaitGroup
	go func() {
		<-ctx.Done()
		_ = ln.Close()
		if tcpLn != nil {
			_ = tcpLn.Close()
		}
	}()
	errCh := make(chan error, 2)
	go acceptDaemonLoop(ctx, &wg, ln, "", cancel, registry, errCh)
	if tcpLn != nil {
		go acceptDaemonLoop(ctx, &wg, tcpLn, tcpToken, cancel, registry, errCh)
	}
	select {
	case <-ctx.Done():
	case err := <-errCh:
		cancel()
		wg.Wait()
		return err
	}
	wg.Wait()
	return nil
}

func acceptDaemonLoop(ctx context.Context, wg *sync.WaitGroup, ln net.Listener, authToken string, stop func(), registry *ProjectRegistry, errCh chan<- error) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) || ctx.Err() != nil {
				return
			}
			select {
			case errCh <- err:
			default:
			}
			return
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			handleDaemonConn(ctx, stop, registry, protocol.NewConnTransport(conn), authToken)
		}()
	}
}

func normalizeListenAddr(addr string) string {
	addr = strings.TrimSpace(addr)
	if addr == "" {
		return ""
	}
	if _, _, err := net.SplitHostPort(addr); err == nil {
		return addr
	}
	if strings.HasPrefix(addr, ":") {
		return addr
	}
	if _, err := strconv.Atoi(addr); err == nil {
		return ":" + addr
	}
	return addr
}

func ensureDaemonToken(token string) (string, error) {
	if strings.TrimSpace(token) != "" {
		return strings.TrimSpace(token), nil
	}
	return LoadOrCreateClientToken()
}

func prepareDaemonSocket(socketPath, pidPath string) (string, error) {
	if err := os.MkdirAll(filepath.Dir(socketPath), 0o700); err != nil {
		return "", err
	}
	lockPath := pidPath + ".lock"
	if data, err := os.ReadFile(pidPath); err == nil {
		if pid, perr := strconv.Atoi(strings.TrimSpace(string(data))); perr == nil && pid > 0 {
			if sigErr := syscallSignal0(pid); sigErr == nil {
				return "", fmt.Errorf("daemon already running with pid %d", pid)
			}
		}
	}
	_ = os.Remove(socketPath)
	_ = os.Remove(pidPath)
	_ = os.Remove(lockPath)
	f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
	if err != nil {
		return "", fmt.Errorf("daemon: lock %s: %w", lockPath, err)
	}
	_, _ = f.WriteString(strconv.Itoa(os.Getpid()) + "\n")
	_ = f.Close()
	return lockPath, nil
}

func syscallSignal0(pid int) error {
	return syscall.Kill(pid, 0)
}

func handleDaemonConn(ctx context.Context, stop func(), registry *ProjectRegistry, t protocol.Transport, authToken string) {
	defer t.Close()
	f, err := t.Recv()
	if err != nil {
		return
	}
	switch f.Type {
	case protocol.FrameList:
		_ = sendProjectList(t, registry, "")
		return
	case protocol.FrameStop:
		_ = sendProjectList(t, registry, "")
		stop()
		return
	case protocol.FrameAttach:
		if authToken != "" {
			attach, err := protocol.Decode[protocol.Attach](f)
			if err != nil {
				_ = sendProtocolError(t, err.Error())
				return
			}
			if attach.Token != authToken {
				_ = sendProtocolError(t, "auth denied")
				return
			}
		}
		handleDaemonAttach(ctx, registry, t, f)
	default:
		_ = sendProtocolError(t, fmt.Sprintf("first frame must be attach, list, or stop; got %q", f.Type))
	}
}

func handleDaemonAttach(ctx context.Context, registry *ProjectRegistry, t protocol.Transport, first protocol.Frame) {
	attach, err := protocol.Decode[protocol.Attach](first)
	if err != nil {
		_ = sendProtocolError(t, err.Error())
		return
	}
	project := registry.Project(attach.ProjectKey)
	if project == nil && attach.ProjectPath != "" {
		project, err = registry.Open(ctx, attach.ProjectPath)
		if err != nil {
			_ = sendProtocolError(t, err.Error())
			return
		}
	}
	if project == nil {
		project = registry.DefaultProject()
	}
	if project == nil {
		_ = sendProtocolError(t, "no project open")
		return
	}
	clientID := fmt.Sprintf("c-%d", time.Now().UnixNano())
	view := ClientView{
		ID:          clientID,
		ProjectKey:  project.Key,
		ProjectName: project.Name,
		Cols:        attach.TermSize.Cols,
		Rows:        attach.TermSize.Rows,
	}
	if child := firstRunningTopLevel(project.Session.Children()); child != nil {
		view.FocusChild(child.ID)
		project.ClaimPaneDisplay(clientID, child.ID, attach.TermSize)
	}
	sub := newClientSubscriber(project, clientID, defaultClientSubscriberQueue)
	project.Session.SubscribeClient(sub)
	defer project.Session.UnsubscribeClient(sub)
	defer project.ReleaseClientDisplays(clientID)

	_ = sendHello(t, project, view.ID)
	_ = sendProjectList(t, registry, project.Key)
	_ = sendChrome(t, project, view)
	if view.FocusedID != "" {
		_ = sendSnapshot(t, project, clientID, view.FocusedID)
	}

	// Close the transport when the daemon context is cancelled (shutdown or
	// `daemon stop`). Without this the t.Recv() loop below blocks forever on a
	// still-connected client and the accept loop's wg.Wait() never returns.
	go func() {
		<-ctx.Done()
		_ = t.Close()
	}()

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			f, ok := sub.Recv()
			if !ok {
				return
			}
			if err := t.Send(f); err != nil {
				return
			}
		}
	}()

	for {
		f, err := t.Recv()
		if err != nil {
			return
		}
		switch f.Type {
		case protocol.FrameDetach:
			return
		case protocol.FrameInput:
			msg, err := protocol.Decode[protocol.Input](f)
			if err == nil {
				if c := project.Session.FindChild(msg.PaneID); c != nil {
					_ = c.InjectAsUser(msg.Bytes)
				}
			}
		case protocol.FrameResize:
			msg, err := protocol.Decode[protocol.Resize](f)
			if err == nil {
				view.Resize(msg.Size.Cols, msg.Size.Rows)
				if view.FocusedID != "" {
					if _, _, ok := project.PaneDisplay(view.FocusedID); !ok {
						project.ClaimPaneDisplay(clientID, view.FocusedID, msg.Size)
					}
				}
				project.ResizeClientDisplays(clientID, msg.Size)
			}
		case protocol.FrameFocus:
			msg, err := protocol.Decode[protocol.Focus](f)
			if err == nil && msg.PaneID != "" {
				view.FocusChild(msg.PaneID)
				project.ClaimPaneDisplay(clientID, msg.PaneID, protocol.Size{Cols: view.Cols, Rows: view.Rows})
				_ = sendChrome(t, project, view)
				_ = sendSnapshot(t, project, clientID, msg.PaneID)
			}
		case protocol.FramePaletteCommand:
			if child := handleDaemonPaletteCommand(project, f); child != nil {
				view.FocusChild(child.ID)
				project.ClaimPaneDisplay(clientID, child.ID, protocol.Size{Cols: view.Cols, Rows: view.Rows})
				_ = sendChrome(t, project, view)
				_ = sendSnapshot(t, project, clientID, child.ID)
			}
		}
		select {
		case <-done:
			return
		default:
		}
	}
}

func handleDaemonPaletteCommand(project *Project, f protocol.Frame) *Child {
	msg, err := protocol.Decode[protocol.PaletteCommand](f)
	if err != nil {
		return nil
	}
	switch msg.Kind {
	case "spawn_command":
		var p struct {
			Argv    []string `json:"argv"`
			Name    string   `json:"name"`
			WorkDir string   `json:"working_dir"`
			Shell   bool     `json:"shell"`
		}
		if err := json.Unmarshal(msg.Data, &p); err != nil || len(p.Argv) == 0 {
			return nil
		}
		name := p.Name
		if name == "" {
			name = strings.Join(p.Argv, " ")
		}
		c, err := project.Launcher.LaunchCommandArgv(p.Argv, name, "", p.WorkDir, nil, p.Shell)
		if err != nil {
			return nil
		}
		return c
	}
	return nil
}

func sendHello(t protocol.Transport, p *Project, clientID string) error {
	f, err := protocol.NewFrame(protocol.FrameHello, protocol.Hello{Version: 1, DaemonID: strconv.Itoa(os.Getpid()), ClientID: clientID, ProjectKey: p.Key})
	if err != nil {
		return err
	}
	return t.Send(f)
}

func sendProjectList(t protocol.Transport, registry *ProjectRegistry, current string) error {
	summaries := registry.Summaries(current)
	projects := make([]protocol.Project, 0, len(summaries))
	for _, p := range summaries {
		projects = append(projects, protocol.Project{Key: p.Key, Path: p.Dir, Name: p.Name, TabCount: p.TabCount})
	}
	f, err := protocol.NewFrame(protocol.FrameProjectList, protocol.ProjectList{Projects: projects})
	if err != nil {
		return err
	}
	return t.Send(f)
}

func sendChrome(t protocol.Transport, p *Project, view ClientView) error {
	pads, _ := p.Pads.List()
	model := buildChromeModel(p.Key, view, p.Session.Children(), pads)
	b, err := json.Marshal(model)
	if err != nil {
		return err
	}
	f, err := protocol.NewFrame(protocol.FrameChrome, protocol.Chrome{ProjectKey: p.Key, Model: b})
	if err != nil {
		return err
	}
	return t.Send(f)
}

func sendSnapshot(t protocol.Transport, p *Project, clientID, paneID string) error {
	b, err := p.Session.SerializeChild(paneID)
	if err != nil {
		return nil
	}
	size, ownerID, _ := p.PaneDisplay(paneID)
	f, err := protocol.NewFrame(protocol.FramePaneSnapshot, protocol.PaneSnapshot{
		PaneID:       paneID,
		Bytes:        b,
		Size:         size,
		DisplayOwner: ownerID == "" || ownerID == clientID,
	})
	if err != nil {
		return err
	}
	return t.Send(f)
}

func sendProtocolError(t protocol.Transport, msg string) error {
	f, err := protocol.NewFrame(protocol.FrameError, protocol.Error{Message: msg})
	if err != nil {
		return err
	}
	return t.Send(f)
}