Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

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