You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

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