feat(task): replace built-in task adaptors with a sandboxed JS plugin system (#7076)

2026-09-11 06:30:21 +00:00 · 2026-08-29 18:51:57 +08:00
parent 7037ac15bd
commit eb48396d5f
336 changed files with 52333 additions and 6369 deletions
@@ -0,0 +1,588 @@
+package router
+
+import (
+	"bufio"
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"os"
+	"sort"
+	"strings"
+	"sync"
+	"sync/atomic"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/controller"
+	"github.com/QuantumNous/new-api/logger"
+	"github.com/QuantumNous/new-api/middleware"
+	"github.com/QuantumNous/new-api/pkg/jsplugin"
+	"github.com/gin-gonic/gin"
+)
+
+type pluginDispatchStateKey struct{}
+
+type pluginDispatchState struct {
+	generation *jsplugin.RoutingGeneration
+	hit        atomic.Bool
+	writer     *gatedResponseWriter
+	requestID  string
+	language   string
+}
+
+func (s *pluginDispatchState) markHit() {
+	s.hit.Store(true)
+	s.writer.activate()
+}
+
+type pluginRouteHandlers func(*jsplugin.RoutingGeneration, jsplugin.RouteBinding) []gin.HandlerFunc
+
+type pluginGenerationBuilder struct {
+	staticRoutes   []gin.RouteInfo
+	trustedProxies []string
+	routeHandlers  pluginRouteHandlers
+	registerRoute  func(*gin.Engine, jsplugin.RouteBinding, []gin.HandlerFunc)
+	configure      func(*gin.Engine) error
+}
+
+type pluginRouteDispatcher struct {
+	registry *jsplugin.Registry
+}
+
+func SetPluginRouter(outer *gin.Engine) gin.HandlerFunc {
+	trustedProxies, _, err := common.ResolveTrustedProxies(os.Getenv("TRUSTED_PROXIES"))
+	dispatcher := &pluginRouteDispatcher{registry: jsplugin.DefaultRegistry}
+	if err != nil {
+		common.SysError("configure plugin router trusted proxies: " + err.Error())
+		return dispatcher.dispatch
+	}
+	builder := newPluginGenerationBuilder(outer.Routes(), trustedProxies, productionPluginRouteHandlers)
+	if err = jsplugin.DefaultRegistry.SetGenerationPreparer(builder.prepare); err != nil {
+		common.SysError("build initial plugin router: " + err.Error())
+	}
+	return dispatcher.dispatch
+}
+
+func newPluginGenerationBuilder(staticRoutes []gin.RouteInfo, trustedProxies []string, handlers pluginRouteHandlers) *pluginGenerationBuilder {
+	builder := &pluginGenerationBuilder{
+		staticRoutes:   append([]gin.RouteInfo(nil), staticRoutes...),
+		trustedProxies: append([]string(nil), trustedProxies...),
+		routeHandlers:  handlers,
+	}
+	builder.registerRoute = func(engine *gin.Engine, binding jsplugin.RouteBinding, routeHandlers []gin.HandlerFunc) {
+		engine.Handle(binding.Route.Method, binding.Route.Path, routeHandlers...)
+	}
+	builder.configure = func(engine *gin.Engine) error {
+		return common.ConfigureTrustedProxies(engine, builder.trustedProxies)
+	}
+	return builder
+}
+
+func productionPluginRouteHandlers(generation *jsplugin.RoutingGeneration, binding jsplugin.RouteBinding) []gin.HandlerFunc {
+	pinRoute := func(c *gin.Context) {
+		pinnedGeneration := generation
+		if state, _ := c.Request.Context().Value(pluginDispatchStateKey{}).(*pluginDispatchState); state != nil && state.generation != nil {
+			pinnedGeneration = state.generation
+		}
+		c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
+			Generation: pinnedGeneration,
+			Plugin:     binding.Plugin,
+		})
+		c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{
+			Generation: pinnedGeneration,
+			Plugin:     binding.Plugin,
+			Route:      binding.Route,
+		})
+		logger.LogDebug(
+			c,
+			"task_plugin subsystem=router event=route_matched generation=%d plugin=%q version=%q method=%q route_type=%q",
+			pinnedGeneration.Number,
+			binding.Plugin.Meta.Key,
+			binding.Plugin.Meta.Version,
+			binding.Route.Method,
+			binding.Route.Type,
+		)
+		c.Next()
+		logger.LogDebug(
+			c,
+			"task_plugin subsystem=router event=route_complete generation=%d plugin=%q method=%q status=%d",
+			pinnedGeneration.Number,
+			binding.Plugin.Meta.Key,
+			binding.Route.Method,
+			c.Writer.Status(),
+		)
+	}
+	return []gin.HandlerFunc{
+		pinRoute,
+		middleware.TokenAuth(),
+		middleware.SystemPerformanceCheck(),
+		middleware.ModelRequestRateLimit(),
+		middleware.PrepareTaskPluginRoute(),
+		middleware.Distribute(),
+		controller.RelayTask,
+	}
+}
+
+func (b *pluginGenerationBuilder) prepare(candidate, current *jsplugin.RoutingGeneration) (jsplugin.PreparedRoutingGeneration, error) {
+	blocked := make(map[*jsplugin.LoadedPlugin]string)
+	for {
+		accepted, routingErrors := b.admitPlugins(candidate, current, blocked)
+		plugins := sortedPlugins(accepted)
+		filtered, err := candidate.RebuildWithPlugins(plugins)
+		if err != nil {
+			return jsplugin.PreparedRoutingGeneration{}, err
+		}
+		engine, offender, buildErr := b.buildInnerEngine(filtered)
+		if buildErr == nil {
+			return jsplugin.PreparedRoutingGeneration{
+				Generation: filtered.WithRuntime(engine),
+				Errors:     routingErrors,
+			}, nil
+		}
+		if offender == "" {
+			return jsplugin.PreparedRoutingGeneration{}, buildErr
+		}
+
+		failedPlugin := accepted[offender]
+		if failedPlugin == nil {
+			return jsplugin.PreparedRoutingGeneration{}, fmt.Errorf("public route rebuild attributed failure to absent plugin %q: %w", offender, buildErr)
+		}
+		blocked[failedPlugin] = fmt.Sprintf("plugin %s rejected while rebuilding public routes: %v", offender, buildErr)
+	}
+}
+
+func (b *pluginGenerationBuilder) admitPlugins(
+	candidate, current *jsplugin.RoutingGeneration,
+	blocked map[*jsplugin.LoadedPlugin]string,
+) (map[string]*jsplugin.LoadedPlugin, map[string]string) {
+	accepted := make(map[string]*jsplugin.LoadedPlugin)
+	currentByKey := make(map[string]*jsplugin.LoadedPlugin)
+	if current != nil && current.RuntimeHandler() != nil {
+		for _, plugin := range current.Plugins() {
+			currentByKey[plugin.Meta.Key] = plugin
+		}
+	}
+
+	unchangedKeys := make([]string, 0)
+	changedKeys := make([]string, 0)
+	newKeys := make([]string, 0)
+	for _, plugin := range candidate.Plugins() {
+		incumbent, existed := currentByKey[plugin.Meta.Key]
+		switch {
+		case existed && incumbent == plugin:
+			unchangedKeys = append(unchangedKeys, plugin.Meta.Key)
+		case existed:
+			changedKeys = append(changedKeys, plugin.Meta.Key)
+		default:
+			newKeys = append(newKeys, plugin.Meta.Key)
+		}
+	}
+	sort.Strings(unchangedKeys)
+	sort.Strings(changedKeys)
+	sort.Strings(newKeys)
+
+	routingErrors := make(map[string]string)
+	orderedKeys := append(unchangedKeys, changedKeys...)
+	orderedKeys = append(orderedKeys, newKeys...)
+	rejectedKeys := make([]string, 0)
+	for _, key := range orderedKeys {
+		plugin, _ := candidate.Get(key)
+		if blockedError := blocked[plugin]; blockedError != "" {
+			routingErrors[key] = blockedError
+		} else if err := b.validatePlugin(plugin, accepted); err != nil {
+			routingErrors[key] = fmt.Sprintf("plugin %s rejected from public routes: %v", key, err)
+		} else {
+			accepted[key] = plugin
+			continue
+		}
+		rejectedKeys = append(rejectedKeys, key)
+	}
+
+	for _, key := range rejectedKeys {
+		plugin, _ := candidate.Get(key)
+		incumbent := currentByKey[key]
+		if incumbent == nil || incumbent == plugin || !candidate.RetainsIncumbent(key) {
+			continue
+		}
+		if blockedError := blocked[incumbent]; blockedError != "" {
+			routingErrors[key] = blockedError
+			continue
+		}
+		if err := b.validatePlugin(incumbent, accepted); err == nil {
+			accepted[key] = incumbent
+		}
+	}
+	return accepted, routingErrors
+}
+
+func (b *pluginGenerationBuilder) validatePlugin(plugin *jsplugin.LoadedPlugin, accepted map[string]*jsplugin.LoadedPlugin) error {
+	for _, route := range plugin.Meta.Routes {
+		for _, staticRoute := range b.staticRoutes {
+			if routeIntersectsStaticRoute(route.Path, staticRoute.Path) {
+				return fmt.Errorf("route %s %s intersects static route %s %s", route.Method, route.Path, staticRoute.Method, staticRoute.Path)
+			}
+		}
+	}
+	for index, left := range plugin.Meta.Routes {
+		for _, right := range plugin.Meta.Routes[index+1:] {
+			if left.Method != right.Method {
+				continue
+			}
+			if routePatternsIntersect(left.Path, right.Path) {
+				return fmt.Errorf("routes %s %s and %s overlap", left.Method, left.Path, right.Path)
+			}
+			if !routesGinCompatible(left.Path, right.Path) {
+				return fmt.Errorf("routes %s %s and %s use incompatible wildcard names", left.Method, left.Path, right.Path)
+			}
+		}
+	}
+	for _, other := range accepted {
+		if other.Meta.Key == plugin.Meta.Key {
+			continue
+		}
+		for _, route := range plugin.Meta.Routes {
+			for _, otherRoute := range other.Meta.Routes {
+				if routePatternsIntersect(route.Path, otherRoute.Path) {
+					return fmt.Errorf("route %s %s overlaps plugin %s route %s %s", route.Method, route.Path, other.Meta.Key, otherRoute.Method, otherRoute.Path)
+				}
+				if route.Method == otherRoute.Method && !routesGinCompatible(route.Path, otherRoute.Path) {
+					return fmt.Errorf("route %s %s is structurally incompatible with plugin %s route %s", route.Method, route.Path, other.Meta.Key, otherRoute.Path)
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func (b *pluginGenerationBuilder) buildInnerEngine(generation *jsplugin.RoutingGeneration) (engine *gin.Engine, offender string, err error) {
+	currentPlugin := ""
+	defer func() {
+		if recovered := recover(); recovered != nil {
+			offender = currentPlugin
+			err = fmt.Errorf("inner Gin registration panic: %v", recovered)
+			engine = nil
+		}
+	}()
+
+	engine = gin.New()
+	engine.RedirectTrailingSlash = false
+	engine.HandleMethodNotAllowed = true
+	if err = b.configure(engine); err != nil {
+		return nil, "", err
+	}
+	engine.Use(importPluginDispatchState())
+	engine.Use(pluginRouteRecovery())
+	engine.Use(middleware.BodyStorageCleanup())
+	engine.NoMethod(func(c *gin.Context) {
+		markPluginRouteHit(c)
+		generation := uint64(0)
+		if state, _ := c.Request.Context().Value(pluginDispatchStateKey{}).(*pluginDispatchState); state != nil && state.generation != nil {
+			generation = state.generation.Number
+		}
+		logger.LogDebug(
+			c,
+			"task_plugin subsystem=router event=method_not_allowed generation=%d request_method=%q status=%d",
+			generation,
+			c.Request.Method,
+			http.StatusMethodNotAllowed,
+		)
+		c.AbortWithStatus(http.StatusMethodNotAllowed)
+	})
+
+	for _, binding := range generation.Routes() {
+		currentPlugin = binding.Plugin.Meta.Key
+		b.registerRoute(engine, binding, b.routeHandlers(generation, binding))
+	}
+	currentPlugin = ""
+	return engine, "", nil
+}
+
+func importPluginDispatchState() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		state, _ := c.Request.Context().Value(pluginDispatchStateKey{}).(*pluginDispatchState)
+		if state != nil {
+			if c.FullPath() != "" {
+				state.markHit()
+			}
+			if state.requestID != "" {
+				c.Set(common.RequestIdKey, state.requestID)
+			}
+			if state.language != "" {
+				c.Set(string(constant.ContextKeyLanguage), state.language)
+			}
+		}
+		c.Set(middleware.RouteTagKey, "relay")
+		c.Next()
+	}
+}
+
+func markPluginRouteHit(c *gin.Context) {
+	state, _ := c.Request.Context().Value(pluginDispatchStateKey{}).(*pluginDispatchState)
+	if state != nil {
+		state.markHit()
+	}
+}
+
+func pluginRouteRecovery() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		defer func() {
+			if recover() == nil {
+				return
+			}
+			common.SysError("panic recovered in plugin route")
+			if pinnedValue, exists := c.Get(jsplugin.ContextKeyPinnedRoute); exists {
+				if pinned, ok := pinnedValue.(jsplugin.PinnedRoute); ok && pinned.Plugin != nil && pinned.Generation != nil {
+					logger.LogDebug(
+						c,
+						"task_plugin subsystem=router event=panic_recovered generation=%d plugin=%q method=%q",
+						pinned.Generation.Number,
+						pinned.Plugin.Meta.Key,
+						pinned.Route.Method,
+					)
+				}
+			}
+			c.Abort()
+			if !c.Writer.Written() {
+				c.JSON(http.StatusInternalServerError, gin.H{
+					"error": gin.H{
+						"message": "internal plugin route error",
+						"type":    "plugin_route_error",
+					},
+				})
+			}
+		}()
+		c.Next()
+	}
+}
+
+func (d *pluginRouteDispatcher) dispatch(c *gin.Context) {
+	generation := d.registry.Generation()
+	if generation == nil || generation.RuntimeHandler() == nil {
+		c.Next()
+		return
+	}
+
+	previousTag, hadPreviousTag := c.Get(middleware.RouteTagKey)
+	c.Set(middleware.RouteTagKey, "relay")
+	originalContext := c.Request.Context()
+	state := &pluginDispatchState{
+		generation: generation,
+		requestID:  c.GetString(common.RequestIdKey),
+		language:   c.GetString(string(constant.ContextKeyLanguage)),
+	}
+	gatedWriter := newGatedResponseWriter(c.Writer)
+	state.writer = gatedWriter
+	c.Request = c.Request.WithContext(context.WithValue(originalContext, pluginDispatchStateKey{}, state))
+
+	generation.RuntimeHandler().ServeHTTP(gatedWriter, c.Request)
+	if state.hit.Load() {
+		if !c.Writer.Written() {
+			c.Writer.WriteHeaderNow()
+		}
+		c.Abort()
+		return
+	}
+
+	c.Request = c.Request.WithContext(originalContext)
+	if hadPreviousTag {
+		c.Set(middleware.RouteTagKey, previousTag)
+	} else {
+		delete(c.Keys, middleware.RouteTagKey)
+	}
+	c.Next()
+}
+
+type gatedResponseWriter struct {
+	underlying    gin.ResponseWriter
+	privateHeader http.Header
+	active        atomic.Bool
+	activateOnce  sync.Once
+	pendingStatus int
+}
+
+func newGatedResponseWriter(underlying gin.ResponseWriter) *gatedResponseWriter {
+	return &gatedResponseWriter{
+		underlying:    underlying,
+		privateHeader: underlying.Header().Clone(),
+	}
+}
+
+func (w *gatedResponseWriter) activate() {
+	w.activateOnce.Do(func() {
+		target := w.underlying.Header()
+		for key := range target {
+			target.Del(key)
+		}
+		for key, values := range w.privateHeader {
+			target[key] = append([]string(nil), values...)
+		}
+		w.active.Store(true)
+		if w.pendingStatus != 0 {
+			w.underlying.WriteHeader(w.pendingStatus)
+		}
+	})
+}
+
+func (w *gatedResponseWriter) Header() http.Header {
+	if w.active.Load() {
+		return w.underlying.Header()
+	}
+	return w.privateHeader
+}
+
+func (w *gatedResponseWriter) WriteHeader(statusCode int) {
+	if w.active.Load() {
+		w.underlying.WriteHeader(statusCode)
+		return
+	}
+	if w.pendingStatus == 0 {
+		w.pendingStatus = statusCode
+	}
+}
+
+func (w *gatedResponseWriter) Write(data []byte) (int, error) {
+	if !w.active.Load() {
+		return len(data), nil
+	}
+	return w.underlying.Write(data)
+}
+
+func (w *gatedResponseWriter) Flush() {
+	if w.active.Load() {
+		w.underlying.Flush()
+	}
+}
+
+func (w *gatedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if !w.active.Load() {
+		return nil, nil, errors.New("cannot hijack an unmatched plugin route")
+	}
+	return w.underlying.Hijack()
+}
+
+func (w *gatedResponseWriter) CloseNotify() <-chan bool {
+	if w.active.Load() {
+		return w.underlying.CloseNotify()
+	}
+	never := make(chan bool)
+	return never
+}
+
+func (w *gatedResponseWriter) Push(target string, options *http.PushOptions) error {
+	if !w.active.Load() {
+		return http.ErrNotSupported
+	}
+	pusher := w.underlying.Pusher()
+	if pusher == nil {
+		return http.ErrNotSupported
+	}
+	return pusher.Push(target, options)
+}
+
+func sortedPluginKeys(plugins map[string]*jsplugin.LoadedPlugin) []string {
+	keys := make([]string, 0, len(plugins))
+	for key := range plugins {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	return keys
+}
+
+func sortedPlugins(plugins map[string]*jsplugin.LoadedPlugin) []*jsplugin.LoadedPlugin {
+	keys := sortedPluginKeys(plugins)
+	sorted := make([]*jsplugin.LoadedPlugin, 0, len(keys))
+	for _, key := range keys {
+		sorted = append(sorted, plugins[key])
+	}
+	return sorted
+}
+
+type routePatternSegment struct {
+	value    string
+	dynamic  bool
+	catchAll bool
+}
+
+func parseRoutePattern(routePath string) []routePatternSegment {
+	parts := strings.Split(strings.TrimPrefix(routePath, "/"), "/")
+	segments := make([]routePatternSegment, 0, len(parts))
+	for _, part := range parts {
+		segments = append(segments, routePatternSegment{
+			value:    part,
+			dynamic:  strings.HasPrefix(part, ":") || strings.HasPrefix(part, "*"),
+			catchAll: strings.HasPrefix(part, "*"),
+		})
+	}
+	return segments
+}
+
+func routePatternsIntersect(leftPath, rightPath string) bool {
+	left := parseRoutePattern(leftPath)
+	right := parseRoutePattern(rightPath)
+	for index := 0; ; index++ {
+		leftDone := index >= len(left)
+		rightDone := index >= len(right)
+		if leftDone || rightDone {
+			return leftDone && rightDone
+		}
+		if left[index].catchAll || right[index].catchAll {
+			return true
+		}
+		if !left[index].dynamic && !right[index].dynamic && left[index].value != right[index].value {
+			return false
+		}
+		if (left[index].dynamic && right[index].value == "") || (right[index].dynamic && left[index].value == "") {
+			return false
+		}
+	}
+}
+
+func routesGinCompatible(leftPath, rightPath string) bool {
+	left := parseRoutePattern(leftPath)
+	right := parseRoutePattern(rightPath)
+	limit := len(left)
+	if len(right) < limit {
+		limit = len(right)
+	}
+	for index := 0; index < limit; index++ {
+		leftSegment := left[index]
+		rightSegment := right[index]
+		if !leftSegment.dynamic && !rightSegment.dynamic {
+			if leftSegment.value != rightSegment.value {
+				return true
+			}
+			continue
+		}
+		if leftSegment.dynamic && rightSegment.dynamic {
+			if leftSegment.value != rightSegment.value {
+				return false
+			}
+			continue
+		}
+		return true
+	}
+	return true
+}
+
+func routeIntersectsStaticRoute(pluginPath, staticPath string) bool {
+	if routePatternsIntersect(pluginPath, staticPath) {
+		return true
+	}
+	if catchAllIndex := strings.LastIndex(staticPath, "/*"); catchAllIndex >= 0 && catchAllIndex+2 < len(staticPath) {
+		prefix := staticPath[:catchAllIndex]
+		if routePatternsIntersect(pluginPath, prefix) || routePatternsIntersect(pluginPath, prefix+"/") {
+			return true
+		}
+		return false
+	}
+	if staticPath == "/" {
+		return false
+	}
+	alternate := strings.TrimSuffix(staticPath, "/")
+	if alternate == staticPath {
+		alternate += "/"
+	}
+	return routePatternsIntersect(pluginPath, alternate)
+}