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

411 строки
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 MasterNodeOnly() func(c *gin.Context) {
  166. return func(c *gin.Context) {
  167. if !common.IsMasterNode {
  168. c.JSON(http.StatusForbidden, gin.H{
  169. "success": false,
  170. "message": "only master node can access this endpoint",
  171. })
  172. c.Abort()
  173. return
  174. }
  175. c.Next()
  176. }
  177. }
  178. func WssAuth(c *gin.Context) {
  179. }
  180. // TokenOrUserAuth allows either session-based user auth or API token auth.
  181. // Used for endpoints that need to be accessible from both the dashboard and API clients.
  182. func TokenOrUserAuth() func(c *gin.Context) {
  183. return func(c *gin.Context) {
  184. // Try session auth first (dashboard users)
  185. session := sessions.Default(c)
  186. if id := session.Get("id"); id != nil {
  187. if status, ok := session.Get("status").(int); ok && status == common.UserStatusEnabled {
  188. c.Set("id", id)
  189. c.Next()
  190. return
  191. }
  192. }
  193. // Fall back to token auth (API clients)
  194. TokenAuth()(c)
  195. }
  196. }
  197. // TokenAuthReadOnly 宽松版本的令牌认证中间件,用于只读查询接口。
  198. // 只验证令牌 key 是否存在,不检查令牌状态、过期时间和额度。
  199. // 即使令牌已过期、已耗尽或已禁用,也允许访问。
  200. // 仍然检查用户是否被封禁。
  201. func TokenAuthReadOnly() func(c *gin.Context) {
  202. return func(c *gin.Context) {
  203. key := c.Request.Header.Get("Authorization")
  204. if key == "" {
  205. c.JSON(http.StatusUnauthorized, gin.H{
  206. "success": false,
  207. "message": "未提供 Authorization 请求头",
  208. })
  209. c.Abort()
  210. return
  211. }
  212. if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
  213. key = strings.TrimSpace(key[7:])
  214. }
  215. key = strings.TrimPrefix(key, "sk-")
  216. parts := strings.Split(key, "-")
  217. key = parts[0]
  218. token, err := model.GetTokenByKey(key, false)
  219. if err != nil {
  220. c.JSON(http.StatusUnauthorized, gin.H{
  221. "success": false,
  222. "message": "无效的令牌",
  223. })
  224. c.Abort()
  225. return
  226. }
  227. userCache, err := model.GetUserCache(token.UserId)
  228. if err != nil {
  229. c.JSON(http.StatusInternalServerError, gin.H{
  230. "success": false,
  231. "message": err.Error(),
  232. })
  233. c.Abort()
  234. return
  235. }
  236. if userCache.Status != common.UserStatusEnabled {
  237. c.JSON(http.StatusForbidden, gin.H{
  238. "success": false,
  239. "message": "用户已被封禁",
  240. })
  241. c.Abort()
  242. return
  243. }
  244. c.Set("id", token.UserId)
  245. c.Set("token_id", token.Id)
  246. c.Set("token_key", token.Key)
  247. c.Next()
  248. }
  249. }
  250. func TokenAuth() func(c *gin.Context) {
  251. return func(c *gin.Context) {
  252. // 先检测是否为ws
  253. if c.Request.Header.Get("Sec-WebSocket-Protocol") != "" {
  254. // Sec-WebSocket-Protocol: realtime, openai-insecure-api-key.sk-xxx, openai-beta.realtime-v1
  255. // read sk from Sec-WebSocket-Protocol
  256. key := c.Request.Header.Get("Sec-WebSocket-Protocol")
  257. parts := strings.Split(key, ",")
  258. for _, part := range parts {
  259. part = strings.TrimSpace(part)
  260. if strings.HasPrefix(part, "openai-insecure-api-key") {
  261. key = strings.TrimPrefix(part, "openai-insecure-api-key.")
  262. break
  263. }
  264. }
  265. c.Request.Header.Set("Authorization", "Bearer "+key)
  266. }
  267. // 检查path包含/v1/messages 或 /v1/models
  268. if strings.Contains(c.Request.URL.Path, "/v1/messages") || strings.Contains(c.Request.URL.Path, "/v1/models") {
  269. anthropicKey := c.Request.Header.Get("x-api-key")
  270. if anthropicKey != "" {
  271. c.Request.Header.Set("Authorization", "Bearer "+anthropicKey)
  272. }
  273. }
  274. // gemini api 从query中获取key
  275. if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models") ||
  276. strings.HasPrefix(c.Request.URL.Path, "/v1beta/openai/models") ||
  277. strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
  278. skKey := c.Query("key")
  279. if skKey != "" {
  280. c.Request.Header.Set("Authorization", "Bearer "+skKey)
  281. }
  282. // 从x-goog-api-key header中获取key
  283. xGoogKey := c.Request.Header.Get("x-goog-api-key")
  284. if xGoogKey != "" {
  285. c.Request.Header.Set("Authorization", "Bearer "+xGoogKey)
  286. }
  287. }
  288. key := c.Request.Header.Get("Authorization")
  289. parts := make([]string, 0)
  290. if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
  291. key = strings.TrimSpace(key[7:])
  292. }
  293. if key == "" || key == "midjourney-proxy" {
  294. key = c.Request.Header.Get("mj-api-secret")
  295. if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
  296. key = strings.TrimSpace(key[7:])
  297. }
  298. }
  299. key, parts = parseTokenKey(key)
  300. token, err := model.ValidateUserToken(key)
  301. if token != nil {
  302. id := c.GetInt("id")
  303. if id == 0 {
  304. c.Set("id", token.UserId)
  305. }
  306. }
  307. if err != nil {
  308. abortWithOpenAiMessage(c, http.StatusUnauthorized, err.Error())
  309. return
  310. }
  311. allowIps := token.GetIpLimits()
  312. if len(allowIps) > 0 {
  313. clientIp := c.ClientIP()
  314. logger.LogDebug(c, "Token has IP restrictions, checking client IP %s", clientIp)
  315. ip := net.ParseIP(clientIp)
  316. if ip == nil {
  317. abortWithOpenAiMessage(c, http.StatusForbidden, "无法解析客户端 IP 地址")
  318. return
  319. }
  320. if common.IsIpInCIDRList(ip, allowIps) == false {
  321. abortWithOpenAiMessage(c, http.StatusForbidden, "您的 IP 不在令牌允许访问的列表中", types.ErrorCodeAccessDenied)
  322. return
  323. }
  324. logger.LogDebug(c, "Client IP %s passed the token IP restrictions check", clientIp)
  325. }
  326. userCache, err := model.GetUserCache(token.UserId)
  327. if err != nil {
  328. abortWithOpenAiMessage(c, http.StatusInternalServerError, err.Error())
  329. return
  330. }
  331. userEnabled := userCache.Status == common.UserStatusEnabled
  332. if !userEnabled {
  333. abortWithOpenAiMessage(c, http.StatusForbidden, "用户已被封禁")
  334. return
  335. }
  336. userCache.WriteContext(c)
  337. userGroup := userCache.Group
  338. tokenGroup := token.Group
  339. if tokenGroup != "" {
  340. // check common.UserUsableGroups[userGroup]
  341. if _, ok := service.GetUserUsableGroups(userGroup)[tokenGroup]; !ok {
  342. abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("无权访问 %s 分组", tokenGroup))
  343. return
  344. }
  345. // check group in common.GroupRatio
  346. if !ratio_setting.ContainsGroupRatio(tokenGroup) {
  347. if tokenGroup != "auto" {
  348. abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被弃用", tokenGroup))
  349. return
  350. }
  351. }
  352. userGroup = tokenGroup
  353. }
  354. common.SetContextKey(c, constant.ContextKeyUsingGroup, userGroup)
  355. err = SetupContextForToken(c, token, parts...)
  356. if err != nil {
  357. return
  358. }
  359. c.Next()
  360. }
  361. }
  362. func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) error {
  363. if token == nil {
  364. return fmt.Errorf("token is nil")
  365. }
  366. c.Set("id", token.UserId)
  367. c.Set("token_id", token.Id)
  368. c.Set("token_key", token.Key)
  369. c.Set("token_name", token.Name)
  370. c.Set("token_unlimited_quota", token.UnlimitedQuota)
  371. if !token.UnlimitedQuota {
  372. c.Set("token_quota", token.RemainQuota)
  373. }
  374. if token.ModelLimitsEnabled {
  375. c.Set("token_model_limit_enabled", true)
  376. c.Set("token_model_limit", token.GetModelLimitsMap())
  377. } else {
  378. c.Set("token_model_limit_enabled", false)
  379. }
  380. common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group)
  381. common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry)
  382. return nil
  383. }