sherlock/internal/socket/socket.go

// Package socket implements a tiny JSON-over-newline RPC transport on
// top of a Unix domain socket. The broker is the server; the sherlock
// CLI and (later) per-service MCPs are clients.
//
// Framing: each rpc.Request and rpc.Response is one JSON object
// terminated by a single '\n'. No length prefixes, no chunked frames.
// This keeps the protocol trivially debuggable with `socat - UNIX-CONNECT:...`.
package socket

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"time"

	"gitea.alexandru.macocian.me/Charlie/sherlock/internal/rpc"
)

// Max bytes for a single framed JSON object. Plenty of headroom for
// status / token payloads; guards against runaway parser allocations.
const maxFrameBytes = 1 << 20 // 1 MiB

// Handler turns a request into a response. Implementations should not
// write directly to the connection — the server takes care of framing.
type Handler interface {
	Handle(ctx context.Context, req rpc.Request) rpc.Response
}

// HandlerFunc adapts a function to the Handler interface.
type HandlerFunc func(ctx context.Context, req rpc.Request) rpc.Response

func (f HandlerFunc) Handle(ctx context.Context, req rpc.Request) rpc.Response {
	return f(ctx, req)
}

// Server listens on a Unix socket and dispatches each framed request
// to a Handler on its own goroutine per connection.
type Server struct {
	ln       net.Listener
	handler  Handler
	wg       sync.WaitGroup
	closing  atomic.Bool
	onAccept func() // test hook: called once per accepted conn before dispatch.
}

// Listen creates the socket at path, removing any stale leftover file.
// Mode 0600 ensures only the owner can dial.
func Listen(path string, handler Handler) (*Server, error) {
	if handler == nil {
		return nil, errors.New("socket: nil handler")
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return nil, fmt.Errorf("socket: mkdir parent: %w", err)
	}
	// Best-effort cleanup of a stale socket file. We deliberately don't
	// touch it if another broker is actively listening — Listen below
	// will fail with EADDRINUSE in that case.
	if fi, err := os.Stat(path); err == nil && fi.Mode()&os.ModeSocket != 0 {
		_ = os.Remove(path)
	}
	ln, err := net.Listen("unix", path)
	if err != nil {
		return nil, fmt.Errorf("socket: listen %s: %w", path, err)
	}
	if err := os.Chmod(path, 0o600); err != nil {
		_ = ln.Close()
		return nil, fmt.Errorf("socket: chmod %s: %w", path, err)
	}
	return &Server{ln: ln, handler: handler}, nil
}

// Addr returns the socket path.
func (s *Server) Addr() string { return s.ln.Addr().String() }

// Serve accepts connections until Close is called or ctx is cancelled.
// Returns nil on a clean shutdown.
func (s *Server) Serve(ctx context.Context) error {
	go func() {
		<-ctx.Done()
		_ = s.Close()
	}()
	for {
		conn, err := s.ln.Accept()
		if err != nil {
			if s.closing.Load() {
				s.wg.Wait()
				return nil
			}
			return fmt.Errorf("socket: accept: %w", err)
		}
		if s.onAccept != nil {
			s.onAccept()
		}
		s.wg.Add(1)
		go s.serveConn(ctx, conn)
	}
}

// Close stops accepting new connections and removes the socket file.
// In-flight handlers continue until they return; Serve waits for them.
func (s *Server) Close() error {
	if !s.closing.CompareAndSwap(false, true) {
		return nil
	}
	err := s.ln.Close()
	_ = os.Remove(s.Addr())
	return err
}

func (s *Server) serveConn(ctx context.Context, conn net.Conn) {
	defer s.wg.Done()
	defer func() { _ = conn.Close() }()

	scanner := bufio.NewScanner(conn)
	scanner.Buffer(make([]byte, 64*1024), maxFrameBytes)
	w := bufio.NewWriter(conn)

	for scanner.Scan() {
		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}
		var req rpc.Request
		var resp rpc.Response
		if err := json.Unmarshal(line, &req); err != nil {
			resp = rpc.Response{Error: rpc.Errorf(rpc.CodeBadRequest, "decode request: %v", err)}
		} else {
			resp = s.handler.Handle(ctx, req)
			resp.ID = req.ID
		}
		if err := writeFrame(w, resp); err != nil {
			return
		}
	}
	// Drain a non-EOF scanner error silently — the client likely hung up.
}

func writeFrame(w *bufio.Writer, resp rpc.Response) error {
	b, err := json.Marshal(resp)
	if err != nil {
		// Last-ditch: send a generic internal error. If even that fails
		// we have nothing useful to do but close the connection.
		fallback, _ := json.Marshal(rpc.Response{
			ID:    resp.ID,
			Error: rpc.NewError(rpc.CodeInternal, "marshal response failed"),
		})
		if _, err := w.Write(fallback); err != nil {
			return err
		}
		if err := w.WriteByte('\n'); err != nil {
			return err
		}
		return w.Flush()
	}
	if _, err := w.Write(b); err != nil {
		return err
	}
	if err := w.WriteByte('\n'); err != nil {
		return err
	}
	return w.Flush()
}

// Client is a thin RPC client over the UDS. It is safe to use a single
// Client from multiple goroutines: Call serialises writes and matches
// responses by request ID.
type Client struct {
	conn   net.Conn
	mu     sync.Mutex
	rd     *bufio.Reader
	nextID atomic.Uint64
}

// Dial connects to the broker UDS at path. The caller owns the returned
// Client and must Close it when done.
func Dial(path string) (*Client, error) {
	conn, err := net.Dial("unix", path)
	if err != nil {
		return nil, err
	}
	return &Client{conn: conn, rd: bufio.NewReader(conn)}, nil
}

// Close closes the underlying socket.
func (c *Client) Close() error { return c.conn.Close() }

// Call sends a request and unmarshals the response.
//
// We hold the per-client mutex for the whole round-trip. Sherlock RPC
// calls are short (sub-second) and never streamed, so request pipelining
// would be pointless — keeping the protocol synchronous keeps the client
// trivially correct and lets us reuse a single connection across calls.
func (c *Client) Call(ctx context.Context, method string, params, result any) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	id := c.nextID.Add(1)
	pb, err := rpc.MarshalParams(params)
	if err != nil {
		return fmt.Errorf("socket: marshal params: %w", err)
	}
	req := rpc.Request{ID: id, Method: method, Params: pb}
	reqBytes, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("socket: marshal request: %w", err)
	}

	// Honour ctx with a connection-level deadline. Cheap and avoids a
	// reader goroutine just for cancellation.
	if dl, ok := ctx.Deadline(); ok {
		_ = c.conn.SetDeadline(dl)
		defer func() { _ = c.conn.SetDeadline(time.Time{}) }()
	}

	if _, err := c.conn.Write(append(reqBytes, '\n')); err != nil {
		return fmt.Errorf("socket: write: %w", err)
	}

	line, err := c.rd.ReadBytes('\n')
	if err != nil {
		if errors.Is(err, io.EOF) {
			return fmt.Errorf("socket: broker closed connection")
		}
		return fmt.Errorf("socket: read: %w", err)
	}
	var resp rpc.Response
	if err := json.Unmarshal(line, &resp); err != nil {
		return fmt.Errorf("socket: decode response: %w", err)
	}
	if resp.ID != id {
		return fmt.Errorf("socket: response id mismatch: got %d want %d", resp.ID, id)
	}
	if resp.Error != nil {
		return resp.Error
	}
	if result == nil || len(resp.Result) == 0 {
		return nil
	}
	if err := json.Unmarshal(resp.Result, result); err != nil {
		return fmt.Errorf("socket: decode result: %w", err)
	}
	return nil
}