yms-daemon/internal/daemonserver/server.go

// Package daemonserver 通过固定的本地 Unix Socket 接受命令行客户端请求，并将其分派给事务编排器执行。
// 服务端负责监听并校验请求、把请求映射到后端更新或重启动作、将执行进度流式写回客户端，
// 最终返回携带事务终态与错误信息的结果。
package daemonserver

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

	"yms-daemon/internal/backendupdate"
	"yms-daemon/internal/daemonapi"
	"yms-daemon/internal/transaction"
)

// maximumRequestBytes 单次请求体的最大字节数，超过该上限的请求会被拒绝，以防内存被异常输入耗尽。
const maximumRequestBytes = 1 << 20

// backendUpdater 抽象了后端更新与重启的具体实现，是 Server 与事务编排器之间的解耦接口。
// 它由 backendupdate 包在运行时注入，便于在测试中用轻量替身替换。
type backendUpdater interface {
	// UpdateRepack 使用重新打包的 ZIP 制品执行后端更新，file 为制品绝对路径。
	UpdateRepack(context.Context, string, backendupdate.ProgressReporter) (transaction.Transaction, error)
	// UpdateNativeJAR 使用原生后端 JAR 文件执行后端更新，file 为制品绝对路径。
	UpdateNativeJAR(context.Context, string, backendupdate.ProgressReporter) (transaction.Transaction, error)
	// UpdateContainerImage 使用容器镜像执行后端更新，imageReference 为镜像引用，startLog 控制是否输出启动日志。
	UpdateContainerImage(context.Context, string, bool, backendupdate.ProgressReporter) (transaction.Transaction, error)
	// Restart 重启后端服务而不更换制品。
	Restart(context.Context, backendupdate.ProgressReporter) (transaction.Transaction, error)
}

// backendDiagnoser 抽象了后端状态诊断与对账能力，供 status/doctor/reconcile 操作使用。
// 它由 backendstatus 包在运行时注入，便于在测试中用轻量替身替换。
type backendDiagnoser interface {
	// Diagnose 执行只读诊断，返回结构化诊断结果。
	Diagnose(context.Context) (daemonapi.Diagnosis, error)
	// Reconcile 生成修复计划或执行自动修复，apply 为 true 时执行修复动作。
	Reconcile(context.Context, bool) (daemonapi.Diagnosis, error)
}

// Server 持有本地 Unix Socket 并把收到的请求分派给事务编排器。
// 它负责连接的生命周期管理、请求解码与校验、执行进度回写以及最终结果的返回。
type Server struct {
	// socketPath 守护进程监听的 Unix Socket 绝对路径。
	socketPath string
	// updater 后端更新与重启动作的实际执行者。
	updater backendUpdater
	// diagnoser 后端状态诊断与对账动作的实际执行者。
	diagnoser backendDiagnoser
	// logger 用于记录服务运行与请求处理过程中的日志。
	logger *slog.Logger
}

// New 构造一个 Server 实例并校验入参。
// socketPath 必须是规范化的绝对路径，updater 与 diagnoser 不能为空，logger 为空时回退为默认日志器。
// 当路径非法或依赖为空时返回 error。
func New(socketPath string, updater backendUpdater, diagnoser backendDiagnoser, logger *slog.Logger) (*Server, error) {
	if !filepath.IsAbs(socketPath) || filepath.Clean(socketPath) != socketPath {
		return nil, errors.New("daemon Unix Socket path must be a clean absolute path")
	}
	if updater == nil {
		return nil, errors.New("backend updater is required")
	}
	if diagnoser == nil {
		return nil, errors.New("backend diagnoser is required")
	}
	if logger == nil {
		logger = slog.Default()
	}
	return &Server{socketPath: socketPath, updater: updater, diagnoser: diagnoser, logger: logger}, nil
}

// Serve 监听 Unix Socket 直到 ctx 被取消，返回前会等待所有连接处理完成并清理 Socket 文件。
// 它先准备 Socket 目录并移除残留的旧 Socket，然后以 0o600 权限监听；
// 每个被接受的连接都由独立的 goroutine 处理，连接在响应写完后由 handle 关闭。
// 当监听或准备 Socket 路径失败时返回 error，ctx 取消后正常返回 nil。
func (s *Server) Serve(ctx context.Context) error {
	if err := prepareSocketPath(s.socketPath); err != nil {
		return err
	}
	listener, err := net.Listen("unix", s.socketPath)
	if err != nil {
		return fmt.Errorf("listen on daemon Unix Socket %s: %w", s.socketPath, err)
	}
	if err := os.Chmod(s.socketPath, 0o600); err != nil {
		_ = listener.Close()
		return fmt.Errorf("set daemon Unix Socket permissions: %w", err)
	}
	defer func() {
		_ = listener.Close()
		_ = os.Remove(s.socketPath)
	}()

	s.logger.InfoContext(ctx, "daemon server listening", "socket", s.socketPath)
	var connections sync.WaitGroup
	defer connections.Wait()
	go func() {
		<-ctx.Done()
		_ = listener.Close()
	}()
	for {
		connection, err := listener.Accept()
		if err != nil {
			if ctx.Err() != nil {
				return nil
			}
			return fmt.Errorf("accept daemon Unix Socket connection: %w", err)
		}
		connections.Add(1)
		go func() {
			defer connections.Done()
			s.handle(ctx, connection)
		}()
	}
}

// handle 处理单个已接受的连接：读取并校验请求、执行对应操作、把进度与结果写回客户端。
// 它限制请求读取时限为 10 秒，仅接受 service 为 backend 的请求；
// 对于 update 操作会按 inputType 分派到不同更新实现，对 restart 操作校验其不携带多余输入。
// 执行期间产生的进度通过 report 回调实时写回，一旦写回失败便停止后续进度发送，
// 最终无论成功失败都会写入一条 ResponseResult 响应并记录日志。
func (s *Server) handle(ctx context.Context, connection net.Conn) {
	defer connection.Close()
	request, err := s.readRequest(connection)
	if err != nil {
		s.writeError(connection, err.Error())
		return
	}
	if request.Service != "backend" {
		s.writeError(connection, "service must be backend")
		return
	}
	progress := s.progressReporter(ctx, connection)
	if request.Operation == daemonapi.OperationStatus || request.Operation == daemonapi.OperationDoctor || request.Operation == daemonapi.OperationReconcile {
		s.handleDiagnosis(ctx, connection, request)
		return
	}
	record, updateErr := s.handleBackendOperation(ctx, request, progress)
	if updateErr != nil && record.ID == "" {
		s.writeError(connection, updateErr.Error())
		return
	}
	response := daemonapi.Response{Kind: daemonapi.ResponseResult, TransactionID: record.ID, State: string(record.State)}
	if updateErr != nil {
		response.Error = updateErr.Error()
		s.logger.ErrorContext(ctx, "backend operation failed", "operation", request.Operation, "transaction_id", record.ID, "state", record.State, "error", updateErr)
	} else {
		s.logger.InfoContext(ctx, "backend operation completed", "operation", request.Operation, "transaction_id", record.ID, "state", record.State)
	}
	if err := s.writeResponse(connection, response); err != nil {
		s.logger.ErrorContext(ctx, "write daemon operation result", "operation", request.Operation, "transaction_id", record.ID, "error", err)
	}
}

func (s *Server) readRequest(connection net.Conn) (daemonapi.Request, error) {
	_ = connection.SetReadDeadline(time.Now().Add(10 * time.Second))
	request, err := decodeRequest(connection)
	_ = connection.SetReadDeadline(time.Time{})
	return request, err
}

func (s *Server) writeError(connection net.Conn, message string) {
	_ = s.writeResponse(connection, daemonapi.Response{Kind: daemonapi.ResponseResult, Error: message})
}

func (s *Server) progressReporter(ctx context.Context, connection net.Conn) backendupdate.ProgressReporter {
	progressWritable := true
	return func(progress backendupdate.Progress) {
		if !progressWritable {
			return
		}
		err := s.writeResponse(connection, daemonapi.Response{Kind: daemonapi.ResponseProgress, TransactionID: progress.TransactionID, State: string(progress.State), Message: progress.Message})
		if err != nil {
			progressWritable = false
			s.logger.WarnContext(ctx, "write daemon update progress", "error", err)
		}
	}
}

func (s *Server) handleBackendOperation(ctx context.Context, request daemonapi.Request, report backendupdate.ProgressReporter) (transaction.Transaction, error) {
	switch request.Operation {
	case daemonapi.OperationUpdate:
		return s.handleUpdate(ctx, request, report)
	case daemonapi.OperationRestart:
		if request.InputType != "" || request.File != "" || request.ImageReference != "" {
			return transaction.Transaction{}, errors.New("restart does not accept inputType or file")
		}
		return s.updater.Restart(ctx, report)
	default:
		return transaction.Transaction{}, errors.New("operation must be update, restart, status, doctor, or reconcile")
	}
}

func (s *Server) handleUpdate(ctx context.Context, request daemonapi.Request, report backendupdate.ProgressReporter) (transaction.Transaction, error) {
	switch request.InputType {
	case daemonapi.InputTypeRepackZIP:
		if request.File == "" || request.ImageReference != "" {
			return transaction.Transaction{}, errors.New("repack-zip requires file and does not accept imageReference")
		}
		return s.updater.UpdateRepack(ctx, request.File, report)
	case daemonapi.InputTypeNativeJAR:
		if request.File == "" || request.ImageReference != "" {
			return transaction.Transaction{}, errors.New("native-jar requires file and does not accept imageReference")
		}
		return s.updater.UpdateNativeJAR(ctx, request.File, report)
	case daemonapi.InputTypeContainerImage:
		if request.File != "" || request.ImageReference == "" {
			return transaction.Transaction{}, errors.New("container-image requires imageReference and does not accept file")
		}
		return s.updater.UpdateContainerImage(ctx, request.ImageReference, request.StartLog, report)
	default:
		return transaction.Transaction{}, errors.New("inputType must be repack-zip, native-jar, or container-image")
	}
}

func (s *Server) handleDiagnosis(ctx context.Context, connection net.Conn, request daemonapi.Request) {
	if request.InputType != "" || request.File != "" || request.ImageReference != "" || (request.Operation != daemonapi.OperationReconcile && request.Apply) {
		s.writeError(connection, request.Operation+" does not accept inputType, file, imageReference, or apply")
		return
	}
	var diagnosis daemonapi.Diagnosis
	var err error
	if request.Operation == daemonapi.OperationReconcile {
		diagnosis, err = s.diagnoser.Reconcile(ctx, request.Apply)
	} else {
		diagnosis, err = s.diagnoser.Diagnose(ctx)
	}
	s.writeDiagnosis(ctx, connection, request.Operation, diagnosis, err)
}

// writeResponse 把一条响应以换行分隔的 JSON 编码写入连接，失败时返回底层写入错误。
func (s *Server) writeResponse(connection net.Conn, response daemonapi.Response) error {
	return json.NewEncoder(connection).Encode(response)
}

// writeDiagnosis 把 status/doctor/reconcile 的诊断结果写回客户端。
// 诊断失败时返回携带错误信息的响应并清空诊断结果，同时记录相应日志。
func (s *Server) writeDiagnosis(ctx context.Context, connection net.Conn, operation string, diagnosis daemonapi.Diagnosis, err error) {
	response := daemonapi.Response{Kind: daemonapi.ResponseResult, Diagnosis: &diagnosis}
	if err != nil {
		response.Diagnosis = nil
		response.Error = err.Error()
		s.logger.ErrorContext(ctx, "backend diagnosis failed", "operation", operation, "error", err)
	} else {
		s.logger.InfoContext(ctx, "backend diagnosis completed", "operation", operation, "healthy", diagnosis.Healthy, "items", len(diagnosis.Items))
	}
	if writeErr := s.writeResponse(connection, response); writeErr != nil {
		s.logger.ErrorContext(ctx, "write daemon diagnosis result", "operation", operation, "error", writeErr)
	}
}

// decodeRequest 从 reader 中解码并校验一条客户端请求。
// 它限制请求体大小不超过 maximumRequestBytes，禁止未知字段，并要求请求只能包含一个 JSON 值，
// 且所有字符串字段不得带有前后空白字符；任一条不满足时返回 error。
func decodeRequest(reader io.Reader) (daemonapi.Request, error) {
	limited := io.LimitReader(reader, maximumRequestBytes+1)
	decoder := json.NewDecoder(limited)
	decoder.DisallowUnknownFields()
	var request daemonapi.Request
	if err := decoder.Decode(&request); err != nil {
		return daemonapi.Request{}, fmt.Errorf("decode daemon request: %w", err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		if err == nil {
			return daemonapi.Request{}, errors.New("daemon request contains multiple JSON values")
		}
		return daemonapi.Request{}, fmt.Errorf("decode daemon request trailing content: %w", err)
	}
	if strings.TrimSpace(request.Operation) != request.Operation || strings.TrimSpace(request.Service) != request.Service || strings.TrimSpace(request.InputType) != request.InputType || strings.TrimSpace(request.File) != request.File || strings.TrimSpace(request.ImageReference) != request.ImageReference {
		return daemonapi.Request{}, errors.New("daemon request fields must not contain surrounding whitespace")
	}
	return request, nil
}

// prepareSocketPath 在监听前准备 Socket 路径：创建父目录，若目标位置已被占用则校验其是否为 Socket，
// 若是残留 Socket 则移除，若被非 Socket 文件占用则返回 error。
func prepareSocketPath(socketPath string) error {
	if err := os.MkdirAll(filepath.Dir(socketPath), 0o755); err != nil {
		return fmt.Errorf("create daemon Unix Socket directory: %w", err)
	}
	info, err := os.Lstat(socketPath)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("inspect daemon Unix Socket path: %w", err)
	}
	if info.Mode()&os.ModeSocket == 0 {
		return fmt.Errorf("daemon Unix Socket path is occupied by a non-socket file: %s", socketPath)
	}
	if err := os.Remove(socketPath); err != nil {
		return fmt.Errorf("remove stale daemon Unix Socket: %w", err)
	}
	return nil
}