优化

internal/worker/task_executor.go

@@ -0,0 +1,170 @@
+package worker
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+	"syscall"
+	"time"
+
+	"kra/internal/biz"
+	"kra/internal/conf"
+)
+
+type TaskExecutor struct {
+	tasks   *biz.TaskUsecase
+	media   *biz.MediaUsecase
+	runtime *conf.Runtime
+}
+
+func NewTaskExecutor(tasks *biz.TaskUsecase, media *biz.MediaUsecase, runtime *conf.Runtime) *TaskExecutor {
+	return &TaskExecutor{tasks: tasks, media: media, runtime: runtime}
+}
+
+func privateIP(ip net.IP) bool {
+	return ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalMulticast() || ip.IsLinkLocalUnicast() || ip.IsUnspecified()
+}
+
+func taskHTTPClient(allowPrivate bool) *http.Client {
+	dialer := &net.Dialer{Timeout: 10 * time.Second, Control: func(_ string, address string, _ syscall.RawConn) error {
+		if allowPrivate {
+			return nil
+		}
+		host, _, err := net.SplitHostPort(address)
+		if err != nil {
+			return fmt.Errorf("解析拨号地址失败: %w", err)
+		}
+		ip := net.ParseIP(host)
+		if ip == nil {
+			return fmt.Errorf("非法拨号 IP: %s", host)
+		}
+		if privateIP(ip) {
+			return fmt.Errorf("目标解析为内网/环回/链路本地地址, 已被 SSRF 防护拒绝(可在任务上开启\"允许内网\"豁免): %s", ip)
+		}
+		return nil
+	}}
+	return &http.Client{Timeout: 30 * time.Second, Transport: &http.Transport{Proxy: nil, DialContext: dialer.DialContext}}
+}
+
+func (e *TaskExecutor) runHTTP(ctx context.Context, task *biz.TimedTask) (string, error) {
+	parsed, err := url.Parse(task.HTTPURL)
+	if err != nil {
+		return "", err
+	}
+	if parsed.Scheme != "http" && parsed.Scheme != "https" {
+		return "", fmt.Errorf("仅允许 http/https, 实际为 %q", parsed.Scheme)
+	}
+	method := strings.ToUpper(task.HTTPMethod)
+	if method == "" {
+		method = http.MethodGet
+	}
+	request, err := http.NewRequestWithContext(ctx, method, task.HTTPURL, bytes.NewBufferString(task.HTTPBody))
+	if err != nil {
+		return "", err
+	}
+	headers := map[string]string{}
+	_ = json.Unmarshal(task.HTTPHeader, &headers)
+	for key, value := range headers {
+		request.Header.Set(key, value)
+	}
+	response, err := taskHTTPClient(task.HTTPAllowPrivate).Do(request)
+	if err != nil {
+		return "", err
+	}
+	defer response.Body.Close()
+	body, _ := io.ReadAll(io.LimitReader(response.Body, 1<<20))
+	output := fmt.Sprintf("HTTP %d: %s", response.StatusCode, string(body))
+	if response.StatusCode < 200 || response.StatusCode >= 300 {
+		return output, fmt.Errorf("非 2xx 响应: %d", response.StatusCode)
+	}
+	return output, nil
+}
+
+var errTaskTimeout = errors.New("任务执行超时")
+
+func truncateTaskText(value string) string {
+	const limit = 4000
+	if len(value) <= limit {
+		return value
+	}
+	return value[:limit] + "...(截断)"
+}
+
+func (e *TaskExecutor) runMethod(task *biz.TimedTask) error {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	defer cancel()
+	done := make(chan error, 1)
+	go func() {
+		defer func() {
+			if recovered := recover(); recovered != nil {
+				done <- fmt.Errorf("panic: %v", recovered)
+			}
+		}()
+		var err error
+		switch task.MethodName {
+		case biz.TaskMethodClearDB:
+			err = e.tasks.CleanupLogs(ctx)
+		case biz.TaskMethodUploads:
+			ttl := 24
+			config := e.runtime.Admin()
+			if config != nil && config.Media != nil && config.Media.SessionTtl > 0 {
+				ttl = int(config.Media.SessionTtl)
+			}
+			err = e.media.CleanupStale(ctx, ttl)
+		}
+		done <- err
+	}()
+	select {
+	case err := <-done:
+		if errors.Is(err, context.DeadlineExceeded) {
+			return errTaskTimeout
+		}
+		return err
+	case <-ctx.Done():
+		return errTaskTimeout
+	}
+}
+
+func (e *TaskExecutor) Run(ctx context.Context, task *biz.TimedTask, trigger string) (log *biz.TimedTaskLog) {
+	started := time.Now()
+	log = &biz.TimedTaskLog{TaskID: task.ID, TaskName: task.Name, TriggerType: trigger, StartedAt: started, Status: "success"}
+	defer func() {
+		if recovered := recover(); recovered != nil {
+			log.Status = "fail"
+			log.ErrorMsg = truncateTaskText(fmt.Sprint(recovered))
+		}
+		log.Output = truncateTaskText(log.Output)
+		log.FinishedAt = time.Now()
+		log.DurationMS = log.FinishedAt.Sub(started).Milliseconds()
+		logCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
+		defer cancel()
+		_ = e.tasks.RecordTaskLog(logCtx, log)
+	}()
+	var err error
+	switch task.ExecutorType {
+	case biz.TaskExecutorMethod:
+		err = e.runMethod(task)
+	case biz.TaskExecutorHTTP:
+		runCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		log.Output, err = e.runHTTP(runCtx, task)
+		cancel()
+	default:
+		err = fmt.Errorf("未知执行器类型: %s", task.ExecutorType)
+	}
+	if err != nil {
+		if errors.Is(err, errTaskTimeout) || errors.Is(err, context.DeadlineExceeded) {
+			log.Status = "timeout"
+		} else {
+			log.Status = "fail"
+		}
+		log.ErrorMsg = truncateTaskText(err.Error())
+	}
+	return log
+}

internal/worker/task_executor_test.go

@@ -0,0 +1,17 @@
+package worker
+
+import (
+	"net"
+	"testing"
+)
+
+func TestPrivateIP(t *testing.T) {
+	for _, value := range []string{"127.0.0.1", "10.0.0.1", "172.16.0.1", "192.168.1.1", "169.254.1.1", "::1"} {
+		if !privateIP(net.ParseIP(value)) {
+			t.Fatalf("expected %s to be private", value)
+		}
+	}
+	if privateIP(net.ParseIP("8.8.8.8")) {
+		t.Fatal("public IP was classified as private")
+	}
+}

internal/worker/task_scheduler.go

@@ -8,14 +8,13 @@ import (
 	"sync"
 	"time"
 
-	"kra/internal/biz"
-	"kra/internal/service"
-
 	"github.com/robfig/cron/v3"
+	"kra/internal/biz"
 )
 
 type TaskScheduler struct {
-	service     *service.TaskService
+	tasks       *biz.TaskUsecase
+	executor    *TaskExecutor
 	logger      *slog.Logger
 	standard    *cron.Cron
 	seconds     *cron.Cron
@@ -33,8 +32,8 @@ type scheduledEntry struct {
 	entry   cron.EntryID
 }
 
-func NewTaskScheduler(service *service.TaskService, logger *slog.Logger) *TaskScheduler {
-	return &TaskScheduler{service: service, logger: logger, standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, subscribers: map[chan []byte]struct{}{}}
+func NewTaskScheduler(tasks *biz.TaskUsecase, executor *TaskExecutor, logger *slog.Logger) *TaskScheduler {
+	return &TaskScheduler{tasks: tasks, executor: executor, logger: logger, standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, subscribers: map[chan []byte]struct{}{}}
 }
 
 func (s *TaskScheduler) Start(ctx context.Context) error {
@@ -44,7 +43,7 @@ func (s *TaskScheduler) Start(ctx context.Context) error {
 	s.ctxMu.Unlock()
 	s.standard.Start()
 	s.seconds.Start()
-	items, err := s.service.ScheduledTasks(ctx)
+	items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil)
 	if err == nil {
 		for _, task := range items {
 			if task.Enabled {
@@ -104,7 +103,7 @@ func (s *TaskScheduler) Reload(ctx context.Context) error {
 		delete(s.entries, id)
 	}
 	s.mu.Unlock()
-	items, err := s.service.ScheduledTasks(ctx)
+	items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil)
 	if err != nil {
 		return err
 	}
@@ -128,7 +127,7 @@ func (s *TaskScheduler) executionContext() context.Context {
 }
 
 func (s *TaskScheduler) run(task *biz.TimedTask, trigger string) {
-	log := s.service.Run(s.executionContext(), task, trigger)
+	log := s.executor.Run(s.executionContext(), task, trigger)
 	if log.Status != "success" {
 		s.Broadcast(map[string]any{"taskId": task.ID, "taskName": task.Name, "status": log.Status, "errorMsg": log.ErrorMsg, "time": time.Now()})
 	}
@@ -160,7 +159,7 @@ func (s *TaskScheduler) Schedule(task *biz.TimedTask) error {
 }
 
 func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error {
-	task, err := s.service.Task(ctx, id)
+	task, err := s.tasks.FindTask(ctx, id)
 	if err != nil {
 		return err
 	}
@@ -168,7 +167,7 @@ func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error {
 }
 
 func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error {
-	task, err := s.service.Task(ctx, id)
+	task, err := s.tasks.FindTask(ctx, id)
 	if err != nil {
 		return err
 	}