Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

397 linhas
11 KiB

  1. package middleware
  2. import (
  3. "fmt"
  4. "net"
  5. "net/http"
  6. "strconv"
  7. "strings"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/constant"
  10. "github.com/QuantumNous/new-api/logger"
  11. "github.com/QuantumNous/new-api/model"
  12. "github.com/QuantumNous/new-api/service"
  13. "github.com/QuantumNous/new-api/setting/ratio_setting"
  14. "github.com/QuantumNous/new-api/types"
  15. "github.com/gin-contrib/sessions"
  16. "github.com/gin-gonic/gin"
  17. )
  18. func validUserInfo(username string, role int) bool {
  19. // check username is empty
  20. if strings.TrimSpace(username) == "" {
  21. return false
  22. }
  23. if !common.IsValidateRole(role) {
  24. return false
  25. }
  26. return true
  27. }
  28. // parseTokenKey 解析令牌字符串,提取 key
  29. // 支持格式:sk-{key}
  30. func parseTokenKey(key string) (string, []string) {
  31. key = strings.TrimPrefix(key, "sk-")
  32. parts := strings.Split(key, "-")
  33. key = parts[0]
  34. return key, parts
  35. }
  36. func authHelper(c *gin.Context, minRole int) {
  37. session := sessions.Default(c)
  38. username := session.Get("username")
  39. role := session.Get("role")
  40. id := session.Get("id")
  41. status := session.Get("status")
  42. useAccessToken := false
  43. if username == nil {
  44. // Check access token
  45. accessToken := c.Request.Header.Get("Authorization")
  46. if accessToken == "" {
  47. c.JSON(http.StatusUnauthorized, gin.H{
  48. "success": false,
  49. "message": "无权进行此操作,未登录且未提供 access token",
  50. })
  51. c.Abort()
  52. return
  53. }
  54. user := model.ValidateAccessToken(accessToken)
  55. if user != nil && user.Username != "" {
  56. if !validUserInfo(user.Username, user.Role) {
  57. c.JSON(http.StatusOK, gin.H{
  58. "success": false,
  59. "message": "无权进行此操作,用户信息无效",
  60. })
  61. c.Abort()
  62. return
  63. }
  64. // Token is valid
  65. username = user.Username
  66. role = user.Role
  67. id = user.Id
  68. status = user.Status
  69. useAccessToken = true
  70. } else {
  71. c.JSON(http.StatusOK, gin.H{
  72. "success": false,
  73. "message": "无权进行此操作,access token 无效",
  74. })
  75. c.Abort()
  76. return
  77. }
  78. }
  79. // get header New-Api-User
  80. apiUserIdStr := c.Request.Header.Get("New-Api-User")
  81. if apiUserIdStr == "" {
  82. c.JSON(http.StatusUnauthorized, gin.H{
  83. "success": false,
  84. "message": "无权进行此操作,未提供 New-Api-User",
  85. })
  86. c.Abort()
  87. return
  88. }
  89. apiUserId, err := strconv.Atoi(apiUserIdStr)
  90. if err != nil {
  91. c.JSON(http.StatusUnauthorized, gin.H{
  92. "success": false,
  93. "message": "无权进行此操作,New-Api-User 格式错误",
  94. })
  95. c.Abort()
  96. return
  97. }
  98. if id != apiUserId {
  99. c.JSON(http.StatusUnauthorized, gin.H{
  100. "success": false,
  101. "message": "无权进行此操作,New-Api-User 与登录用户不匹配",
  102. })
  103. c.Abort()
  104. return
  105. }
  106. if status.(int) == common.UserStatusDisabled {
  107. c.JSON(http.StatusOK, gin.H{
  108. "success": false,
  109. "message": "用户已被封禁",
  110. })
  111. c.Abort()
  112. return
  113. }
  114. if role.(int) < minRole {
  115. c.JSON(http.StatusOK, gin.H{
  116. "success": false,
  117. "message": "无权进行此操作,权限不足",
  118. })
  119. c.Abort()
  120. return
  121. }
  122. if !validUserInfo(username.(string), role.(int)) {
  123. c.JSON(http.StatusOK, gin.H{
  124. "success": false,
  125. "message": "无权进行此操作,用户信息无效",
  126. })
  127. c.Abort()
  128. return
  129. }
  130. // 防止不同newapi版本冲突,导致数据不通用
  131. c.Header("Auth-Version", "864b7076dbcd0a3c01b5520316720ebf")
  132. c.Set("username", username)
  133. c.Set("role", role)
  134. c.Set("id", id)
  135. c.Set("group", session.Get("group"))
  136. c.Set("user_group", session.Get("group"))
  137. c.Set("use_access_token", useAccessToken)
  138. c.Next()
  139. }
  140. func TryUserAuth() func(c *gin.Context) {
  141. return func(c *gin.Context) {
  142. session := sessions.Default(c)
  143. id := session.Get("id")
  144. if id != nil {
  145. c.Set("id", id)
  146. }
  147. c.Next()
  148. }
  149. }
  150. func UserAuth() func(c *gin.Context) {
  151. return func(c *gin.Context) {
  152. authHelper(c, common.RoleCommonUser)
  153. }
  154. }
  155. func AdminAuth() func(c *gin.Context) {
  156. return func(c *gin.Context) {
  157. authHelper(c, common.RoleAdminUser)
  158. }
  159. }
  160. func RootAuth() func(c *gin.Context) {
  161. return func(c *gin.Context) {
  162. authHelper(c, common.RoleRootUser)
  163. }
  164. }
  165. func WssAuth(c *gin.Context) {
  166. }
  167. // TokenOrUserAuth allows either session-based user auth or API token auth.
  168. // Used for endpoints that need to be accessible from both the dashboard and API clients.
  169. func TokenOrUserAuth() func(c *gin.Context) {
  170. return func(c *gin.Context) {
  171. // Try session auth first (dashboard users)
  172. session := sessions.Default(c)
  173. if id := session.Get("id"); id != nil {
  174. if status, ok := session.Get("status").(int); ok && status == common.UserStatusEnabled {
  175. c.Set("id", id)
  176. c.Next()
  177. return
  178. }
  179. }
  180. // Fall back to token auth (API clients)
  181. TokenAuth()(c)
  182. }
  183. }
  184. // TokenAuthReadOnly 宽松版本的令牌认证中间件,用于只读查询接口。
  185. // 只验证令牌 key 是否存在,不检查令牌状态、过期时间和额度。
  186. // 即使令牌已过期、已耗尽或已禁用,也允许访问。
  187. // 仍然检查用户是否被封禁。
  188. func TokenAuthReadOnly() func(c *gin.Context) {
  189. return func(c *gin.Context) {
  190. key := c.Request.Header.Get("Authorization")
  191. if key == "" {
  192. c.JSON(http.StatusUnauthorized, gin.H{
  193. "success": false,
  194. "message": "未提供 Authorization 请求头",
  195. })
  196. c.Abort()
  197. return
  198. }
  199. if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
  200. key = strings.TrimSpace(key[7:])
  201. }
  202. key = strings.TrimPrefix(key, "sk-")
  203. parts := strings.Split(key, "-")
  204. key = parts[0]
  205. token, err := model.GetTokenByKey(key, false)
  206. if err != nil {
  207. c.JSON(http.StatusUnauthorized, gin.H{
  208. "success": false,
  209. "message": "无效的令牌",
  210. })
  211. c.Abort()
  212. return
  213. }
  214. userCache, err := model.GetUserCache(token.UserId)
  215. if err != nil {
  216. c.JSON(http.StatusInternalServerError, gin.H{
  217. "success": false,
  218. "message": err.Error(),
  219. })
  220. c.Abort()
  221. return
  222. }
  223. if userCache.Status != common.UserStatusEnabled {
  224. c.JSON(http.StatusForbidden, gin.H{
  225. "success": false,
  226. "message": "用户已被封禁",
  227. })
  228. c.Abort()
  229. return
  230. }
  231. c.Set("id", token.UserId)
  232. c.Set("token_id", token.Id)
  233. c.Set("token_key", token.Key)
  234. c.Next()
  235. }
  236. }
  237. func TokenAuth() func(c *gin.Context) {
  238. return func(c *gin.Context) {
  239. // 先检测是否为ws
  240. if c.Request.Header.Get("Sec-WebSocket-Protocol") != "" {
  241. // Sec-WebSocket-Protocol: realtime, openai-insecure-api-key.sk-xxx, openai-beta.realtime-v1
  242. // read sk from Sec-WebSocket-Protocol
  243. key := c.Request.Header.Get("Sec-WebSocket-Protocol")
  244. parts := strings.Split(key, ",")
  245. for _, part := range parts {
  246. part = strings.TrimSpace(part)
  247. if strings.HasPrefix(part, "openai-insecure-api-key") {
  248. key = strings.TrimPrefix(part, "openai-insecure-api-key.")
  249. break
  250. }
  251. }
  252. c.Request.Header.Set("Authorization", "Bearer "+key)
  253. }
  254. // 检查path包含/v1/messages 或 /v1/models
  255. if strings.Contains(c.Request.URL.Path, "/v1/messages") || strings.Contains(c.Request.URL.Path, "/v1/models") {
  256. anthropicKey := c.Request.Header.Get("x-api-key")
  257. if anthropicKey != "" {
  258. c.Request.Header.Set("Authorization", "Bearer "+anthropicKey)
  259. }
  260. }
  261. // gemini api 从query中获取key
  262. if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models") ||
  263. strings.HasPrefix(c.Request.URL.Path, "/v1beta/openai/models") ||
  264. strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
  265. skKey := c.Query("key")
  266. if skKey != "" {
  267. c.Request.Header.Set("Authorization", "Bearer "+skKey)
  268. }
  269. // 从x-goog-api-key header中获取key
  270. xGoogKey := c.Request.Header.Get("x-goog-api-key")
  271. if xGoogKey != "" {
  272. c.Request.Header.Set("Authorization", "Bearer "+xGoogKey)
  273. }
  274. }
  275. key := c.Request.Header.Get("Authorization")
  276. parts := make([]string, 0)
  277. if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
  278. key = strings.TrimSpace(key[7:])
  279. }
  280. if key == "" || key == "midjourney-proxy" {
  281. key = c.Request.Header.Get("mj-api-secret")
  282. if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
  283. key = strings.TrimSpace(key[7:])
  284. }
  285. }
  286. key, parts = parseTokenKey(key)
  287. token, err := model.ValidateUserToken(key)
  288. if token != nil {
  289. id := c.GetInt("id")
  290. if id == 0 {
  291. c.Set("id", token.UserId)
  292. }
  293. }
  294. if err != nil {
  295. abortWithOpenAiMessage(c, http.StatusUnauthorized, err.Error())
  296. return
  297. }
  298. allowIps := token.GetIpLimits()
  299. if len(allowIps) > 0 {
  300. clientIp := c.ClientIP()
  301. logger.LogDebug(c, "Token has IP restrictions, checking client IP %s", clientIp)
  302. ip := net.ParseIP(clientIp)
  303. if ip == nil {
  304. abortWithOpenAiMessage(c, http.StatusForbidden, "无法解析客户端 IP 地址")
  305. return
  306. }
  307. if common.IsIpInCIDRList(ip, allowIps) == false {
  308. abortWithOpenAiMessage(c, http.StatusForbidden, "您的 IP 不在令牌允许访问的列表中", types.ErrorCodeAccessDenied)
  309. return
  310. }
  311. logger.LogDebug(c, "Client IP %s passed the token IP restrictions check", clientIp)
  312. }
  313. userCache, err := model.GetUserCache(token.UserId)
  314. if err != nil {
  315. abortWithOpenAiMessage(c, http.StatusInternalServerError, err.Error())
  316. return
  317. }
  318. userEnabled := userCache.Status == common.UserStatusEnabled
  319. if !userEnabled {
  320. abortWithOpenAiMessage(c, http.StatusForbidden, "用户已被封禁")
  321. return
  322. }
  323. userCache.WriteContext(c)
  324. userGroup := userCache.Group
  325. tokenGroup := token.Group
  326. if tokenGroup != "" {
  327. // check common.UserUsableGroups[userGroup]
  328. if _, ok := service.GetUserUsableGroups(userGroup)[tokenGroup]; !ok {
  329. abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("无权访问 %s 分组", tokenGroup))
  330. return
  331. }
  332. // check group in common.GroupRatio
  333. if !ratio_setting.ContainsGroupRatio(tokenGroup) {
  334. if tokenGroup != "auto" {
  335. abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被弃用", tokenGroup))
  336. return
  337. }
  338. }
  339. userGroup = tokenGroup
  340. }
  341. common.SetContextKey(c, constant.ContextKeyUsingGroup, userGroup)
  342. err = SetupContextForToken(c, token, parts...)
  343. if err != nil {
  344. return
  345. }
  346. c.Next()
  347. }
  348. }
  349. func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) error {
  350. if token == nil {
  351. return fmt.Errorf("token is nil")
  352. }
  353. c.Set("id", token.UserId)
  354. c.Set("token_id", token.Id)
  355. c.Set("token_key", token.Key)
  356. c.Set("token_name", token.Name)
  357. c.Set("token_unlimited_quota", token.UnlimitedQuota)
  358. if !token.UnlimitedQuota {
  359. c.Set("token_quota", token.RemainQuota)
  360. }
  361. if token.ModelLimitsEnabled {
  362. c.Set("token_model_limit_enabled", true)
  363. c.Set("token_model_limit", token.GetModelLimitsMap())
  364. } else {
  365. c.Set("token_model_limit_enabled", false)
  366. }
  367. common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group)
  368. common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry)
  369. return nil
  370. }