middleware.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. package server
  2. import (
  3. "io"
  4. "github.com/gin-gonic/gin"
  5. "github.com/langgenius/dify-plugin-daemon/internal/db"
  6. "github.com/langgenius/dify-plugin-daemon/internal/server/constants"
  7. "github.com/langgenius/dify-plugin-daemon/internal/types/entities"
  8. "github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
  9. "github.com/langgenius/dify-plugin-daemon/internal/types/models"
  10. "github.com/langgenius/dify-plugin-daemon/internal/utils/log"
  11. )
  12. func CheckingKey(key string) gin.HandlerFunc {
  13. return func(c *gin.Context) {
  14. // get header X-Api-Key
  15. if c.GetHeader(constants.X_API_KEY) != key {
  16. c.JSON(200, entities.NewErrorResponse(-401, "Unauthorized"))
  17. c.Abort()
  18. return
  19. }
  20. c.Next()
  21. }
  22. }
  23. func (app *App) FetchPluginInstallation() gin.HandlerFunc {
  24. return func(ctx *gin.Context) {
  25. plugin_id := ctx.Request.Header.Get(constants.X_PLUGIN_ID)
  26. if plugin_id == "" {
  27. ctx.AbortWithStatusJSON(400, gin.H{"error": "Invalid request, plugin_id is required"})
  28. return
  29. }
  30. tenant_id := ctx.Param("tenant_id")
  31. if tenant_id == "" {
  32. ctx.AbortWithStatusJSON(400, gin.H{"error": "Invalid request, tenant_id is required"})
  33. return
  34. }
  35. // fetch plugin installation
  36. installation, err := db.GetOne[models.PluginInstallation](
  37. db.Equal("tenant_id", tenant_id),
  38. db.Equal("plugin_id", plugin_id),
  39. )
  40. if err != nil {
  41. ctx.AbortWithStatusJSON(400, gin.H{"error": "Invalid request, " + err.Error()})
  42. return
  43. }
  44. identity, err := plugin_entities.NewPluginUniqueIdentifier(installation.PluginUniqueIdentifier)
  45. if err != nil {
  46. ctx.AbortWithStatusJSON(400, gin.H{"error": "Invalid request, " + err.Error()})
  47. return
  48. }
  49. ctx.Set(constants.CONTEXT_KEY_PLUGIN_INSTALLATION, installation)
  50. ctx.Set(constants.CONTEXT_KEY_PLUGIN_UNIQUE_IDENTIFIER, identity)
  51. ctx.Next()
  52. }
  53. }
  54. // RedirectPluginInvoke redirects the request to the correct cluster node
  55. func (app *App) RedirectPluginInvoke() gin.HandlerFunc {
  56. return func(ctx *gin.Context) {
  57. // get plugin unique identifier
  58. identity_any, ok := ctx.Get(constants.CONTEXT_KEY_PLUGIN_UNIQUE_IDENTIFIER)
  59. if !ok {
  60. ctx.AbortWithStatusJSON(500, gin.H{"error": "Internal server error, plugin unique identifier not found"})
  61. return
  62. }
  63. identity, ok := identity_any.(plugin_entities.PluginUniqueIdentifier)
  64. if !ok {
  65. ctx.AbortWithStatusJSON(500, gin.H{"error": "Internal server error, failed to parse plugin unique identifier"})
  66. return
  67. }
  68. // check if plugin in current node
  69. if !app.cluster.IsPluginOnCurrentNode(
  70. identity,
  71. ) {
  72. app.redirectPluginInvokeByPluginIdentifier(ctx, identity)
  73. ctx.Abort()
  74. } else {
  75. ctx.Next()
  76. }
  77. }
  78. }
  79. func (app *App) redirectPluginInvokeByPluginIdentifier(
  80. ctx *gin.Context,
  81. plugin_unique_identifier plugin_entities.PluginUniqueIdentifier,
  82. ) {
  83. // try find the correct node
  84. nodes, err := app.cluster.FetchPluginAvailableNodesById(plugin_unique_identifier.String())
  85. if err != nil {
  86. ctx.AbortWithStatusJSON(500, gin.H{"error": "Internal server error"})
  87. return
  88. } else if len(nodes) == 0 {
  89. ctx.AbortWithStatusJSON(404, gin.H{"error": "No available node"})
  90. return
  91. }
  92. // redirect to the correct node
  93. node_id := nodes[0]
  94. status_code, header, body, err := app.cluster.RedirectRequest(node_id, ctx.Request)
  95. if err != nil {
  96. log.Error("redirect request failed: %s", err.Error())
  97. ctx.AbortWithStatusJSON(500, gin.H{"error": "Internal server error"})
  98. return
  99. }
  100. // set status code
  101. ctx.Writer.WriteHeader(status_code)
  102. // set header
  103. for key, values := range header {
  104. for _, value := range values {
  105. ctx.Writer.Header().Set(key, value)
  106. }
  107. }
  108. for {
  109. buf := make([]byte, 1024)
  110. n, err := body.Read(buf)
  111. if err != nil && err != io.EOF {
  112. break
  113. } else if err != nil {
  114. ctx.Writer.Write(buf[:n])
  115. break
  116. }
  117. if n > 0 {
  118. ctx.Writer.Write(buf[:n])
  119. }
  120. }
  121. }
  122. func (app *App) InitClusterID() gin.HandlerFunc {
  123. return func(ctx *gin.Context) {
  124. ctx.Set(constants.CONTEXT_KEY_CLUSTER_ID, app.cluster.ID())
  125. ctx.Next()
  126. }
  127. }