nginx-ui/server/modules/oauth2/controller.go

package oauth2

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/astaxie/beego/logs"
	"github.com/astaxie/beego/orm"
	"io"
	"nginx-ui/server/base"
	"nginx-ui/server/config"
	"nginx-ui/server/models"
	"nginx-ui/server/utils"
)

type Controller struct {
	base.Controller
}

type SSOReq struct {
	Code  string `Json:"code"`
	Scope string `Json:"scope"`
	State string `Json:"state"`
}

// Get 获取oauth2.0的登录url
func (c *Controller) Get() {
	state, err := utils.RandPassword(6)
	if err != nil {
		c.ErrorJson(err)
		return
	}
	url := config.OauthConfig.AuthCodeURL(state)
	c.AddRespData("redirect_url", url).AddRespData("state", state).Json()
}

// Callback 用户注册
func (c *Controller) Callback() {
	var ssoReq SSOReq
	err := json.Unmarshal(c.Ctx.Input.RequestBody, &ssoReq)
	if err != nil {
		logs.Error(err, string(c.Ctx.Input.RequestBody))
		c.ErrorJson(err)
		return
	}
	oauth := config.OauthConfig
	if len(ssoReq.Code) == 0 {
		c.SetCode(-1).SetMsg("登录失败(Code)：code is empty").Json()
		return
	}
	token, err := oauth.Exchange(context.Background(), ssoReq.Code)
	if err != nil {
		logs.Error("ExchangeToken", err)
		c.SetCode(-1).SetMsg("登录失败(Exchange)：" + err.Error()).Json()
		return
	}
	client := oauth.Client(context.Background(), token)
	resp, err := client.Get(oauth.Userinfo)
	if err != nil {
		logs.Error("GetUserinfo", err)
		c.SetCode(-1).SetMsg(fmt.Sprintf("登录失败(Userinfo)：%s", err.Error())).Json()
		return
	}
	defer resp.Body.Close()
	content, err := io.ReadAll(resp.Body)
	if err != nil {
		logs.Error("GetUserinfo Read Body", err)
		c.SetCode(-1).SetMsg(fmt.Sprintf("登录失败(Userinfo)：%s", err.Error())).Json()
		return
	}
	user := models.User{}
	err = json.Unmarshal(content, &user)
	if err != nil {
		logs.Error("GetUserinfo Unmarshal", err)
	}
	if len(user.Account) == 0 {
		c.SetCode(-1).SetMsg("登录失败,请确认userinfo接口返回了account字段").Json()
		return
	}
	if len(user.Nickname) == 0 {
		user.Nickname = user.Account
	}
	o := orm.NewOrm()
	err = o.Read(&user, "Account")
	if err != nil {
		_, err = o.Insert(&user)
	}
	user.Password = ""
	if err != nil {
		c.SetCode(-1).SetMsg(fmt.Sprintf("保存用户失败：%s", err.Error())).Json()
		return
	}
	c.SetSession("user", user)
	c.SetData(user).Json()
}