server.go

package shush

import (
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"
)

// Serve starts the agent server on the given Unix socket path.
// It blocks until a stop request is received or a signal is caught.
func Serve(socketPath string, errOut io.Writer) error {
	return ServeWithCapability(socketPath, errOut, ResolveCapability(nil))
}

// ServeWithCapability starts the agent server and optionally requires
// a capability token on all incoming requests when capability is non-empty.
func ServeWithCapability(socketPath string, errOut io.Writer, capability string) error {
	if strings.TrimSpace(socketPath) == "" {
		return errors.New("empty agent socket path")
	}
	capability = strings.TrimSpace(capability)

	if err := os.MkdirAll(filepath.Dir(socketPath), 0o700); err != nil {
		return fmt.Errorf("create agent directory: %w", err)
	}

	if info, err := os.Stat(socketPath); err == nil {
		if info.Mode()&os.ModeSocket == 0 {
			return fmt.Errorf("agent socket path exists and is not a socket: %s", socketPath)
		}
		probe := &Client{SocketPath: socketPath, Capability: capability}
		if _, err := probe.do(request{Action: actionPing}); err == nil {
			return fmt.Errorf("agent already running at %s", socketPath)
		} else if isCapabilityAuthError(err) {
			return fmt.Errorf("agent already running at %s (capability mismatch)", socketPath)
		}
		if err := os.Remove(socketPath); err != nil {
			return fmt.Errorf("remove stale socket %s: %w", socketPath, err)
		}
	} else if !os.IsNotExist(err) {
		return fmt.Errorf("stat socket path %s: %w", socketPath, err)
	}

	listener, err := net.Listen("unix", socketPath)
	if err != nil {
		return fmt.Errorf("listen on %s: %w", socketPath, err)
	}
	defer func() {
		_ = listener.Close()
		_ = os.Remove(socketPath)
	}()
	if err := os.Chmod(socketPath, 0o600); err != nil {
		return fmt.Errorf("chmod socket %s: %w", socketPath, err)
	}

	state := &agentState{
		entries:           map[string]entry{},
		tokenEntries:      map[string]tokenEntry{},
		requireCapability: capability != "",
	}
	if state.requireCapability {
		state.capabilityHash = sha256.Sum256([]byte(capability))
	}

	stopCh := make(chan struct{})
	var stopOnce sync.Once
	stop := func() {
		stopOnce.Do(func() {
			close(stopCh)
			_ = listener.Close()
		})
	}

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	defer signal.Stop(signals)
	go func() {
		select {
		case <-signals:
			stop()
		case <-stopCh:
		}
	}()

	for {
		conn, acceptErr := listener.Accept()
		if acceptErr != nil {
			select {
			case <-stopCh:
				return nil
			default:
				if errOut != nil {
					fmt.Fprintf(errOut, "agent accept error: %v\n", acceptErr)
				}
				continue
			}
		}
		go handleConn(conn, state, stop)
	}
}

func handleConn(conn net.Conn, state *agentState, stopFn func()) {
	defer conn.Close()

	_ = conn.SetDeadline(time.Now().Add(DefaultReqTimeout))

	decoder := json.NewDecoder(conn)
	encoder := json.NewEncoder(conn)

	if err := ensureSameUID(conn); err != nil {
		_ = encoder.Encode(response{
			OK:    false,
			Error: err.Error(),
		})
		return
	}

	var req request
	if err := decoder.Decode(&req); err != nil {
		_ = encoder.Encode(response{
			OK:    false,
			Error: "invalid request",
		})
		return
	}
	if err := state.authorizeCapability(req.Capability); err != nil {
		_ = encoder.Encode(response{
			OK:    false,
			Error: err.Error(),
		})
		return
	}

	resp := state.apply(req)
	_ = encoder.Encode(resp)

	if req.Action == actionStop && resp.OK {
		stopFn()
	}
}

func ensureSameUID(conn net.Conn) error {
	peerUID, err := socketPeerUID(conn)
	if err != nil {
		return fmt.Errorf("forbidden: cannot verify peer uid (%w)", err)
	}
	currentUID := uint32(os.Getuid())
	if peerUID != currentUID {
		return fmt.Errorf("forbidden: peer uid %d does not match agent uid %d", peerUID, currentUID)
	}
	return nil
}

func (s *agentState) authorizeCapability(capability string) error {
	if !s.requireCapability {
		return nil
	}
	capability = strings.TrimSpace(capability)
	if capability == "" {
		return errors.New("forbidden: capability required")
	}
	sum := sha256.Sum256([]byte(capability))
	if subtle.ConstantTimeCompare(sum[:], s.capabilityHash[:]) != 1 {
		return errors.New("forbidden: invalid capability")
	}
	return nil
}

func isCapabilityAuthError(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "forbidden: capability")
}

func (s *agentState) apply(req request) response {
	switch req.Action {
	case actionPing:
		return response{OK: true}
	case actionGet:
		return s.handleGet(req)
	case actionSet:
		return s.handleSet(req)
	case actionClear:
		return s.handleClear(req)
	case actionSetToken:
		return s.handleSetToken(req)
	case actionGetToken:
		return s.handleGetToken(req)
	case actionClearToken:
		return s.handleClearToken(req)
	case actionStop:
		s.clearAll()
		return response{OK: true}
	default:
		return response{
			OK:    false,
			Error: "unsupported action",
		}
	}
}

func (s *agentState) handleGet(req request) response {
	key := strings.TrimSpace(req.Key)
	if key == "" {
		return response{OK: false, Error: "key is required"}
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	e, ok := s.entries[key]
	if !ok {
		return response{OK: true, Found: false}
	}
	if e.expiresAt <= time.Now().Unix() {
		Wipe(e.secret)
		delete(s.entries, key)
		return response{OK: true, Found: false}
	}
	return response{
		OK:        true,
		Found:     true,
		SecretB64: base64.StdEncoding.EncodeToString(e.secret),
	}
}

func (s *agentState) handleSet(req request) response {
	key := strings.TrimSpace(req.Key)
	if key == "" {
		return response{OK: false, Error: "key is required"}
	}
	if req.ExpiresAt <= time.Now().Unix() {
		return response{OK: false, Error: "expires_at must be in the future"}
	}

	secret, err := base64.StdEncoding.DecodeString(req.SecretB64)
	if err != nil || len(secret) == 0 {
		return response{OK: false, Error: "invalid secret_b64"}
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if old, ok := s.entries[key]; ok {
		Wipe(old.secret)
	}
	s.entries[key] = entry{
		secret:    secret,
		expiresAt: req.ExpiresAt,
	}
	return response{OK: true}
}

func (s *agentState) handleClear(req request) response {
	key := strings.TrimSpace(req.Key)
	keyPrefix := strings.TrimSpace(req.KeyPrefix)
	keyPattern := strings.TrimSpace(req.KeyPattern)

	if key != "" && (keyPrefix != "" || keyPattern != "") {
		return response{OK: false, Error: "key cannot be combined with key_prefix or key_pattern"}
	}
	if keyPrefix != "" && keyPattern != "" {
		return response{OK: false, Error: "key_prefix cannot be combined with key_pattern"}
	}
	if keyPattern != "" {
		if _, err := path.Match(keyPattern, "sample"); err != nil {
			return response{OK: false, Error: "invalid key_pattern"}
		}
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if key == "" && keyPrefix == "" && keyPattern == "" {
		for k, e := range s.entries {
			Wipe(e.secret)
			delete(s.entries, k)
		}
		return response{OK: true}
	}
	if keyPrefix != "" {
		for k, e := range s.entries {
			if strings.HasPrefix(k, keyPrefix) {
				Wipe(e.secret)
				delete(s.entries, k)
			}
		}
		return response{OK: true}
	}
	if keyPattern != "" {
		for k, e := range s.entries {
			matched, _ := path.Match(keyPattern, k)
			if matched {
				Wipe(e.secret)
				delete(s.entries, k)
			}
		}
		return response{OK: true}
	}

	if e, ok := s.entries[key]; ok {
		Wipe(e.secret)
		delete(s.entries, key)
	}
	return response{OK: true}
}

func (s *agentState) handleSetToken(req request) response {
	token := strings.TrimSpace(req.Token)
	if token == "" {
		return response{OK: false, Error: "token is required"}
	}
	if req.ExpiresAt <= time.Now().Unix() {
		return response{OK: false, Error: "expires_at must be in the future"}
	}
	if req.MaxUses <= 0 {
		return response{OK: false, Error: "max_uses must be positive"}
	}

	secret, err := base64.StdEncoding.DecodeString(req.SecretB64)
	if err != nil || len(secret) == 0 {
		return response{OK: false, Error: "invalid secret_b64"}
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if old, ok := s.tokenEntries[token]; ok {
		Wipe(old.secret)
	}
	s.tokenEntries[token] = tokenEntry{
		secret:    secret,
		expiresAt: req.ExpiresAt,
		remaining: req.MaxUses,
	}
	return response{OK: true}
}

func (s *agentState) handleGetToken(req request) response {
	token := strings.TrimSpace(req.Token)
	if token == "" {
		return response{OK: false, Error: "token is required"}
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	te, ok := s.tokenEntries[token]
	if !ok {
		return response{OK: true, Found: false}
	}
	if te.expiresAt <= time.Now().Unix() || te.remaining <= 0 {
		Wipe(te.secret)
		delete(s.tokenEntries, token)
		return response{OK: true, Found: false}
	}

	secretCopy := append([]byte(nil), te.secret...)
	te.remaining--
	if te.remaining <= 0 {
		Wipe(te.secret)
		delete(s.tokenEntries, token)
	} else {
		s.tokenEntries[token] = te
	}

	return response{
		OK:        true,
		Found:     true,
		SecretB64: base64.StdEncoding.EncodeToString(secretCopy),
	}
}

func (s *agentState) handleClearToken(req request) response {
	token := strings.TrimSpace(req.Token)

	s.mu.Lock()
	defer s.mu.Unlock()

	if token == "" {
		for k, te := range s.tokenEntries {
			Wipe(te.secret)
			delete(s.tokenEntries, k)
		}
		return response{OK: true}
	}

	if te, ok := s.tokenEntries[token]; ok {
		Wipe(te.secret)
		delete(s.tokenEntries, token)
	}
	return response{OK: true}
}

func (s *agentState) clearAll() {
	_ = s.handleClear(request{})
	_ = s.handleClearToken(request{})
}