diff --git a/controller/heicode_oauth.go b/controller/heicode_oauth.go new file mode 100644 index 00000000..57040938 --- /dev/null +++ b/controller/heicode_oauth.go @@ -0,0 +1,214 @@ +package controller + +import ( + "errors" + "fmt" + "html" + "net/http" + "net/url" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + + "github.com/gin-contrib/sessions" + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +// HeiCode-specific browser-login bridge. +// +// Flow: +// +// 1. HeiCode desktop opens +// GET /heicode/oauth/authorize?state=&redirect_uri=&provider_id= +// 2. If the user is not signed in to new-api, we render a small HTML page +// that links to /login and polls /heicode/oauth/session every 1.5s +// until the user has authenticated. Then it auto-redirects back. +// 3. Once the session is valid, we find or create a token named "HeiCode" +// for the current user and 302-redirect back to redirect_uri with the +// token in the query string: +// ?state=&token=sk-XXXX&provider_id= +// +// The redirect_uri is restricted to loopback addresses for safety. + +const heicodeTokenName = "HeiCode" + +func isLoopbackHost(host string) bool { + host = strings.ToLower(strings.TrimSpace(host)) + return host == "127.0.0.1" || host == "localhost" || host == "::1" +} + +// HeicodeOAuthAuthorize handles GET /heicode/oauth/authorize. +func HeicodeOAuthAuthorize(c *gin.Context) { + state := strings.TrimSpace(c.Query("state")) + redirectURI := strings.TrimSpace(c.Query("redirect_uri")) + providerID := strings.TrimSpace(c.Query("provider_id")) + + if state == "" || redirectURI == "" { + c.String(http.StatusBadRequest, "missing state or redirect_uri") + return + } + + parsed, err := url.Parse(redirectURI) + if err != nil || parsed.Scheme == "" || parsed.Host == "" { + c.String(http.StatusBadRequest, "invalid redirect_uri") + return + } + if !isLoopbackHost(parsed.Hostname()) { + c.String(http.StatusBadRequest, "redirect_uri must be a loopback address") + return + } + + session := sessions.Default(c) + rawID := session.Get("id") + if rawID == nil { + renderHeicodeLoginRequired(c, state, redirectURI, providerID) + return + } + + userID, ok := rawID.(int) + if !ok || userID <= 0 { + renderHeicodeLoginRequired(c, state, redirectURI, providerID) + return + } + + tokenKey, err := getOrCreateHeicodeToken(userID) + if err != nil { + c.String(http.StatusInternalServerError, "failed to issue HeiCode token: "+err.Error()) + return + } + + q := parsed.Query() + q.Set("state", state) + q.Set("token", "sk-"+tokenKey) + if providerID != "" { + q.Set("provider_id", providerID) + } + parsed.RawQuery = q.Encode() + c.Redirect(http.StatusFound, parsed.String()) +} + +// HeicodeOAuthSession handles GET /heicode/oauth/session — used by the +// "please log in" fallback page to poll session readiness from JS. +func HeicodeOAuthSession(c *gin.Context) { + session := sessions.Default(c) + id := session.Get("id") + loggedIn := false + var userID int + if v, ok := id.(int); ok && v > 0 { + loggedIn = true + userID = v + } + c.JSON(http.StatusOK, gin.H{ + "logged_in": loggedIn, + "user_id": userID, + }) +} + +func getOrCreateHeicodeToken(userID int) (string, error) { + var existing model.Token + err := model.DB.Where("user_id = ? AND name = ?", userID, heicodeTokenName). + Order("id desc"). + First(&existing).Error + if err == nil { + if existing.Status != common.TokenStatusEnabled { + existing.Status = common.TokenStatusEnabled + if updateErr := existing.Update(); updateErr != nil { + return "", updateErr + } + } + return existing.GetFullKey(), nil + } + if !errors.Is(err, gorm.ErrRecordNotFound) { + return "", err + } + + rawKey, err := common.GenerateKey() + if err != nil { + return "", err + } + now := common.GetTimestamp() + tok := model.Token{ + UserId: userID, + Name: heicodeTokenName, + Key: rawKey, + Status: common.TokenStatusEnabled, + CreatedTime: now, + AccessedTime: now, + ExpiredTime: -1, + UnlimitedQuota: true, + } + if err := tok.Insert(); err != nil { + return "", err + } + return tok.GetFullKey(), nil +} + +func renderHeicodeLoginRequired(c *gin.Context, state, redirectURI, providerID string) { + authorizeURL := fmt.Sprintf( + "/heicode/oauth/authorize?state=%s&redirect_uri=%s&provider_id=%s", + url.QueryEscape(state), + url.QueryEscape(redirectURI), + url.QueryEscape(providerID), + ) + loginURL := "/login" + + body := fmt.Sprintf(` + + + + +HeiCode 登录授权 + + + +
+

请先登录新 API 控制台

+

HeiCode 客户端正在等待你完成授权。请在新打开的页面中登录到 %s，登录成功后本页会自动跳转回 HeiCode。

+ 打开登录页 +

如果浏览器没有自动跳转，请手动在登录后回到此页面，或刷新本页。

+

提供方：%s

+
+ + +`, + html.EscapeString(providerID), + html.EscapeString(loginURL), + html.EscapeString(providerID), + authorizeURL, + ) + + c.Header("Content-Type", "text/html; charset=utf-8") + c.Status(http.StatusOK) + _, _ = c.Writer.WriteString(body) +} diff --git a/docker-compose.override.yml b/docker-compose.override.yml new file mode 100644 index 00000000..c4e27665 --- /dev/null +++ b/docker-compose.override.yml @@ -0,0 +1,15 @@ +# Override file: build new-api from local source so the HeiCode OAuth bridge +# (controller/heicode_oauth.go + router/heicode-router.go) is included. +# +# Usage: +# docker compose up -d --build new-api +# +# Remove this file (or rename it) to fall back to the upstream +# calciumion/new-api:latest image declared in docker-compose.yml. + +services: + new-api: + image: new-api-heicode:local + build: + context: . + dockerfile: Dockerfile diff --git a/router/heicode-router.go b/router/heicode-router.go new file mode 100644 index 00000000..f73932b0 --- /dev/null +++ b/router/heicode-router.go @@ -0,0 +1,21 @@ +package router + +import ( + "github.com/QuantumNous/new-api/controller" + + "github.com/gin-gonic/gin" +) + +// SetHeicodeOAuthRouter wires up HeiCode-specific browser-login bridge endpoints. +// These endpoints are mounted at /heicode/* (no /api prefix) so the redirect_uri +// from HeiCode desktop matches: http(s):///heicode/oauth/authorize. +func SetHeicodeOAuthRouter(router *gin.Engine) { + heicode := router.Group("/heicode") + { + oauth := heicode.Group("/oauth") + { + oauth.GET("/authorize", controller.HeicodeOAuthAuthorize) + oauth.GET("/session", controller.HeicodeOAuthSession) + } + } +} diff --git a/router/main.go b/router/main.go index d3769bd5..95c9fdc1 100644 --- a/router/main.go +++ b/router/main.go @@ -17,6 +17,7 @@ func SetRouter(router *gin.Engine, assets ThemeAssets) { SetDashboardRouter(router) SetRelayRouter(router) SetVideoRouter(router) + SetHeicodeOAuthRouter(router) frontendBaseUrl := os.Getenv("FRONTEND_BASE_URL") if common.IsMasterNode && frontendBaseUrl != "" { frontendBaseUrl = ""