package controller import ( "errors" "fmt" "html" "net/http" "net/url" "strings" "github.com/heicode/manager/common" "github.com/heicode/manager/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 heicode, 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) } // Include the authenticated user's identity so the desktop client can // display it in the title bar / settings without making a separate // /api/user/self call (that endpoint requires a session cookie which the // desktop client doesn't have). This is the missing piece that left the // UserPill empty for OAuth-route logins. var user model.User if dbErr := model.DB.First(&user, userID).Error; dbErr == nil { if email := strings.TrimSpace(user.Email); email != "" { q.Set("email", email) } display := strings.TrimSpace(user.DisplayName) if display == "" { display = strings.TrimSpace(user.Username) } if display != "" { q.Set("name", display) } switch { case user.Role >= 100: q.Set("role", "root") case user.Role >= 10: q.Set("role", "admin") default: q.Set("role", "user") } if g := strings.TrimSpace(user.Group); g != "" { q.Set("channel_id", g) } q.Set("user_id", fmt.Sprintf("%d", userID)) } 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 := "/sign-in?redirect=" + url.QueryEscape(authorizeURL) body := fmt.Sprintf(` Heicode 登录授权

请先登录 Heicode 控制台

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

打开登录页

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

租户:%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) }