✨ github oauth2 login /gitee oauth2 login

api/oauth_api/gitee_api.go

@@ -0,0 +1,230 @@
+package oauth_api
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	ginI18n "github.com/gin-contrib/i18n"
+	"github.com/gin-gonic/gin"
+	"github.com/yitter/idgenerator-go/idgen"
+	"gorm.io/gorm"
+	"net/http"
+	"schisandra-cloud-album/common/enum"
+	"schisandra-cloud-album/common/result"
+	"schisandra-cloud-album/global"
+	"schisandra-cloud-album/model"
+	"strconv"
+	"time"
+)
+
+type GiteeUser struct {
+	AvatarURL         string      `json:"avatar_url"`
+	Bio               string      `json:"bio"`
+	Blog              string      `json:"blog"`
+	CreatedAt         time.Time   `json:"created_at"`
+	Email             string      `json:"email"`
+	EventsURL         string      `json:"events_url"`
+	Followers         int         `json:"followers"`
+	FollowersURL      string      `json:"followers_url"`
+	Following         int         `json:"following"`
+	FollowingURL      string      `json:"following_url"`
+	GistsURL          string      `json:"gists_url"`
+	HTMLURL           string      `json:"html_url"`
+	ID                int         `json:"id"`
+	Login             string      `json:"login"`
+	Name              string      `json:"name"`
+	OrganizationsURL  string      `json:"organizations_url"`
+	PublicGists       int         `json:"public_gists"`
+	PublicRepos       int         `json:"public_repos"`
+	ReceivedEventsURL string      `json:"received_events_url"`
+	Remark            string      `json:"remark"`
+	ReposURL          string      `json:"repos_url"`
+	Stared            int         `json:"stared"`
+	StarredURL        string      `json:"starred_url"`
+	SubscriptionsURL  string      `json:"subscriptions_url"`
+	Type              string      `json:"type"`
+	UpdatedAt         time.Time   `json:"updated_at"`
+	URL               string      `json:"url"`
+	Watched           int         `json:"watched"`
+	Weibo             interface{} `json:"weibo"`
+}
+
+// GetGiteeRedirectUrl 获取Gitee登录地址
+func (OAuthAPI) GetGiteeRedirectUrl(c *gin.Context) {
+	clientID := global.CONFIG.OAuth.Gitee.ClientID
+	redirectURI := global.CONFIG.OAuth.Gitee.RedirectURI
+	url := "https://gitee.com/oauth/authorize?client_id=" + clientID + "&redirect_uri=" + redirectURI + "&response_type=code"
+	result.OkWithData(url, c)
+	return
+}
+
+// GetGiteeTokenAuthUrl 获取Gitee token
+func GetGiteeTokenAuthUrl(code string) string {
+	clientId := global.CONFIG.OAuth.Gitee.ClientID
+	clientSecret := global.CONFIG.OAuth.Gitee.ClientSecret
+	redirectURI := global.CONFIG.OAuth.Gitee.RedirectURI
+	return fmt.Sprintf(
+		"https://gitee.com/oauth/token?grant_type=authorization_code&code=%s&client_id=%s&redirect_uri=%s&client_secret=%s",
+		code, clientId, redirectURI, clientSecret,
+	)
+}
+
+// GetGiteeToken 获取 token
+func GetGiteeToken(url string) (*Token, error) {
+
+	// 形成请求
+	var req *http.Request
+	var err error
+	if req, err = http.NewRequest(http.MethodPost, url, nil); err != nil {
+		return nil, err
+	}
+	req.Header.Set("accept", "application/json")
+
+	// 发送请求并获得响应
+	var httpClient = http.Client{}
+	var res *http.Response
+	if res, err = httpClient.Do(req); err != nil {
+		return nil, err
+	}
+
+	// 将响应体解析为 token，并返回
+	var token Token
+	if err = json.NewDecoder(res.Body).Decode(&token); err != nil {
+		return nil, err
+	}
+	return &token, nil
+}
+
+// GetGiteeUserInfo 获取用户信息
+func GetGiteeUserInfo(token *Token) (map[string]interface{}, error) {
+
+	// 形成请求
+	var userInfoUrl = "https://gitee.com/api/v5/user" // github用户信息获取接口
+	var req *http.Request
+	var err error
+	if req, err = http.NewRequest(http.MethodGet, userInfoUrl, nil); err != nil {
+		return nil, err
+	}
+	req.Header.Set("accept", "application/json")
+	req.Header.Set("Authorization", fmt.Sprintf("token %s", token.AccessToken))
+	// 发送请求并获取响应
+	var client = http.Client{}
+	var res *http.Response
+	if res, err = client.Do(req); err != nil {
+		return nil, err
+	}
+
+	// 将响应的数据写入 userInfo 中，并返回
+	var userInfo = make(map[string]interface{})
+	if err = json.NewDecoder(res.Body).Decode(&userInfo); err != nil {
+		return nil, err
+	}
+	return userInfo, nil
+}
+
+// GiteeCallback 处理Gitee回调
+func (OAuthAPI) GiteeCallback(c *gin.Context) {
+	var err error
+	// 获取 code
+	var code = c.Query("code")
+	if code == "" {
+		result.FailWithMessage(ginI18n.MustGetMessage(c, "ParamsError"), c)
+		return
+	}
+	// 通过 code, 获取 token
+	var tokenAuthUrl = GetGiteeTokenAuthUrl(code)
+	var token *Token
+	if token, err = GetGiteeToken(tokenAuthUrl); err != nil {
+		global.LOG.Error(err)
+		return
+	}
+
+	// 通过token，获取用户信息
+	var userInfo map[string]interface{}
+	if userInfo, err = GetGiteeUserInfo(token); err != nil {
+		global.LOG.Error(err)
+		return
+	}
+
+	userInfoBytes, err := json.Marshal(userInfo)
+	if err != nil {
+		global.LOG.Error(err)
+		return
+	}
+	var giteeUser GiteeUser
+	err = json.Unmarshal(userInfoBytes, &giteeUser)
+	if err != nil {
+		global.LOG.Error(err)
+		return
+	}
+
+	Id := strconv.Itoa(giteeUser.ID)
+	userSocial, err := userSocialService.QueryUserSocialByUUID(Id)
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		// 第一次登录，创建用户
+		uid := idgen.NextId()
+		uidStr := strconv.FormatInt(uid, 10)
+		user := model.ScaAuthUser{
+			UID:      &uidStr,
+			Username: &giteeUser.Login,
+			Nickname: &giteeUser.Name,
+			Avatar:   &giteeUser.AvatarURL,
+			Blog:     &giteeUser.Blog,
+			Email:    &giteeUser.Email,
+		}
+		addUser, err := userService.AddUser(user)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		gitee := enum.OAuthSourceGitee
+		userSocial = model.ScaAuthUserSocial{
+			UserID: &addUser.ID,
+			UUID:   &Id,
+			Source: &gitee,
+		}
+		err = userSocialService.AddUserSocial(userSocial)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		userRole := model.ScaAuthUserRole{
+			UserID: addUser.ID,
+			RoleID: enum.User,
+		}
+		err = userRoleService.AddUserRole(userRole)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		res, data := HandelUserLogin(addUser)
+		if !res {
+			return
+		}
+		tokenData, err := json.Marshal(data)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		formattedScript := fmt.Sprintf(script, tokenData, global.CONFIG.System.Web)
+		c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(formattedScript))
+	} else {
+		user, err := userService.QueryUserByUsername(giteeUser.Login)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		res, data := HandelUserLogin(user)
+		if !res {
+			return
+		}
+		tokenData, err := json.Marshal(data)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		formattedScript := fmt.Sprintf(script, tokenData, global.CONFIG.System.Web)
+		c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(formattedScript))
+	}
+	return
+}

api/oauth_api/github_api.go

@@ -0,0 +1,235 @@
+package oauth_api
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	ginI18n "github.com/gin-contrib/i18n"
+	"github.com/gin-gonic/gin"
+	"github.com/yitter/idgenerator-go/idgen"
+	"gorm.io/gorm"
+	"net/http"
+	"schisandra-cloud-album/common/enum"
+	"schisandra-cloud-album/common/result"
+	"schisandra-cloud-album/global"
+	"schisandra-cloud-album/model"
+	"strconv"
+)
+
+type Token struct {
+	AccessToken string `json:"access_token"`
+}
+type GitHubUser struct {
+	AvatarURL         string      `json:"avatar_url"`
+	Bio               interface{} `json:"bio"`
+	Blog              string      `json:"blog"`
+	Company           interface{} `json:"company"`
+	CreatedAt         string      `json:"created_at"`
+	Email             string      `json:"email"`
+	EventsURL         string      `json:"events_url"`
+	Followers         int         `json:"followers"`
+	FollowersURL      string      `json:"followers_url"`
+	Following         int         `json:"following"`
+	FollowingURL      string      `json:"following_url"`
+	GistsURL          string      `json:"gists_url"`
+	GravatarID        string      `json:"gravatar_id"`
+	Hireable          interface{} `json:"hireable"`
+	HTMLURL           string      `json:"html_url"`
+	ID                int         `json:"id"`
+	Location          interface{} `json:"location"`
+	Login             string      `json:"login"`
+	Name              string      `json:"name"`
+	NodeID            string      `json:"node_id"`
+	NotificationEmail interface{} `json:"notification_email"`
+	OrganizationsURL  string      `json:"organizations_url"`
+	PublicGists       int         `json:"public_gists"`
+	PublicRepos       int         `json:"public_repos"`
+	ReceivedEventsURL string      `json:"received_events_url"`
+	ReposURL          string      `json:"repos_url"`
+	SiteAdmin         bool        `json:"site_admin"`
+	StarredURL        string      `json:"starred_url"`
+	SubscriptionsURL  string      `json:"subscriptions_url"`
+	TwitterUsername   interface{} `json:"twitter_username"`
+	Type              string      `json:"type"`
+	UpdatedAt         string      `json:"updated_at"`
+	URL               string      `json:"url"`
+}
+
+// GetRedirectUrl 获取github登录url
+func (OAuthAPI) GetRedirectUrl(c *gin.Context) {
+	clientId := global.CONFIG.OAuth.Github.ClientID
+	redirectUrl := global.CONFIG.OAuth.Github.RedirectURI
+	url := "https://github.com/login/oauth/authorize?client_id=" + clientId + "&redirect_uri=" + redirectUrl
+	result.OkWithData(url, c)
+	return
+}
+
+// GetTokenAuthUrl 通过code获取token认证url
+func GetTokenAuthUrl(code string) string {
+	clientId := global.CONFIG.OAuth.Github.ClientID
+	clientSecret := global.CONFIG.OAuth.Github.ClientSecret
+	return fmt.Sprintf(
+		"https://github.com/login/oauth/access_token?client_id=%s&client_secret=%s&code=%s",
+		clientId, clientSecret, code,
+	)
+}
+
+// GetToken 获取 token
+func GetToken(url string) (*Token, error) {
+
+	// 形成请求
+	var req *http.Request
+	var err error
+	if req, err = http.NewRequest(http.MethodGet, url, nil); err != nil {
+		return nil, err
+	}
+	req.Header.Set("accept", "application/json")
+
+	// 发送请求并获得响应
+	var httpClient = http.Client{}
+	var res *http.Response
+	if res, err = httpClient.Do(req); err != nil {
+		return nil, err
+	}
+
+	// 将响应体解析为 token，并返回
+	var token Token
+	if err = json.NewDecoder(res.Body).Decode(&token); err != nil {
+		return nil, err
+	}
+	return &token, nil
+}
+
+// GetUserInfo 获取用户信息
+func GetUserInfo(token *Token) (map[string]interface{}, error) {
+
+	// 形成请求
+	var userInfoUrl = "https://api.github.com/user" // github用户信息获取接口
+	var req *http.Request
+	var err error
+	if req, err = http.NewRequest(http.MethodGet, userInfoUrl, nil); err != nil {
+		return nil, err
+	}
+	req.Header.Set("accept", "application/json")
+	req.Header.Set("Authorization", fmt.Sprintf("token %s", token.AccessToken))
+
+	// 发送请求并获取响应
+	var client = http.Client{}
+	var res *http.Response
+	if res, err = client.Do(req); err != nil {
+		return nil, err
+	}
+
+	// 将响应的数据写入 userInfo 中，并返回
+	var userInfo = make(map[string]interface{})
+	if err = json.NewDecoder(res.Body).Decode(&userInfo); err != nil {
+		return nil, err
+	}
+	return userInfo, nil
+}
+
+// Callback 登录回调函数
+func (OAuthAPI) Callback(c *gin.Context) {
+	var err error
+	// 获取 code
+	var code = c.Query("code")
+	if code == "" {
+		result.FailWithMessage(ginI18n.MustGetMessage(c, "ParamsError"), c)
+		return
+	}
+	// 通过 code, 获取 token
+	var tokenAuthUrl = GetTokenAuthUrl(code)
+	var token *Token
+	if token, err = GetToken(tokenAuthUrl); err != nil {
+		global.LOG.Error(err)
+		return
+	}
+
+	// 通过token，获取用户信息
+	var userInfo map[string]interface{}
+	if userInfo, err = GetUserInfo(token); err != nil {
+		global.LOG.Error(err)
+		return
+	}
+	//json 转 struct
+	userInfoBytes, err := json.Marshal(userInfo)
+	if err != nil {
+		global.LOG.Error(err)
+		return
+	}
+	var gitHubUser GitHubUser
+	err = json.Unmarshal(userInfoBytes, &gitHubUser)
+	if err != nil {
+		global.LOG.Error(err)
+		return
+	}
+	Id := strconv.Itoa(gitHubUser.ID)
+	userSocial, err := userSocialService.QueryUserSocialByUUID(Id)
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		// 第一次登录，创建用户
+		uid := idgen.NextId()
+		uidStr := strconv.FormatInt(uid, 10)
+		user := model.ScaAuthUser{
+			UID:      &uidStr,
+			Username: &gitHubUser.Login,
+			Nickname: &gitHubUser.Name,
+			Avatar:   &gitHubUser.AvatarURL,
+			Blog:     &gitHubUser.Blog,
+			Email:    &gitHubUser.Email,
+		}
+		addUser, err := userService.AddUser(user)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		github := enum.OAuthSourceGithub
+		userSocial = model.ScaAuthUserSocial{
+			UserID: &addUser.ID,
+			UUID:   &Id,
+			Source: &github,
+		}
+		err = userSocialService.AddUserSocial(userSocial)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		userRole := model.ScaAuthUserRole{
+			UserID: addUser.ID,
+			RoleID: enum.User,
+		}
+		err = userRoleService.AddUserRole(userRole)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		res, data := HandelUserLogin(addUser)
+		if !res {
+			return
+		}
+		tokenData, err := json.Marshal(data)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		formattedScript := fmt.Sprintf(script, tokenData, global.CONFIG.System.Web)
+		c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(formattedScript))
+	} else {
+		user, err := userService.QueryUserByUsername(gitHubUser.Login)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		res, data := HandelUserLogin(user)
+		if !res {
+			return
+		}
+		tokenData, err := json.Marshal(data)
+		if err != nil {
+			global.LOG.Error(err)
+			return
+		}
+		formattedScript := fmt.Sprintf(script, tokenData, global.CONFIG.System.Web)
+		c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(formattedScript))
+	}
+	return
+}

api/oauth_api/oauth.go

@@ -1,3 +1,80 @@
 package oauth_api
 
+import (
+	"encoding/json"
+	"schisandra-cloud-album/api/user_api/dto"
+	"schisandra-cloud-album/common/constant"
+	"schisandra-cloud-album/common/redis"
+	"schisandra-cloud-album/model"
+	"schisandra-cloud-album/service"
+	"schisandra-cloud-album/utils"
+	"time"
+)
+
 type OAuthAPI struct{}
+
+var userService = service.Service.UserService
+var userRoleService = service.Service.UserRoleService
+var userSocialService = service.Service.UserSocialService
+var rolePermissionService = service.Service.RolePermissionService
+var permissionServiceService = service.Service.PermissionService
+var roleService = service.Service.RoleService
+
+var script = `
+        <script>
+        window.opener.postMessage('%s', '%s');
+        window.close();
+        </script>
+        `
+
+// HandelUserLogin 处理用户登录
+func HandelUserLogin(user model.ScaAuthUser) (bool, map[string]interface{}) {
+	ids, err := userRoleService.GetUserRoleIdsByUserId(user.ID)
+	if err != nil {
+		return false, nil
+	}
+	permissionIds := rolePermissionService.QueryPermissionIdsByRoleId(ids)
+	permissions, err := permissionServiceService.GetPermissionsByIds(permissionIds)
+	if err != nil {
+		return false, nil
+	}
+	serializedPermissions, err := json.Marshal(permissions)
+	if err != nil {
+		return false, nil
+	}
+	wrong := redis.Set(constant.UserAuthPermissionRedisKey+*user.UID, serializedPermissions, 0).Err()
+	if wrong != nil {
+		return false, nil
+	}
+	roleList, err := roleService.GetRoleListByIds(ids)
+	if err != nil {
+		return false, nil
+	}
+	serializedRoleList, err := json.Marshal(roleList)
+	if err != nil {
+		return false, nil
+	}
+	er := redis.Set(constant.UserAuthRoleRedisKey+*user.UID, serializedRoleList, 0).Err()
+	if er != nil {
+		return false, nil
+	}
+	accessToken, refreshToken, expiresAt := utils.GenerateAccessTokenAndRefreshToken(utils.JWTPayload{UserID: user.UID, RoleID: ids})
+
+	data := dto.ResponseData{
+		AccessToken:  accessToken,
+		RefreshToken: refreshToken,
+		ExpiresAt:    expiresAt,
+		UID:          user.UID,
+	}
+	fail := redis.Set(constant.UserLoginTokenRedisKey+*user.UID, data, time.Hour*24*7).Err()
+	if fail != nil {
+		return false, nil
+	}
+	responseData := map[string]interface{}{
+		"code":    0,
+		"message": "success",
+		"data":    data,
+		"success": true,
+	}
+	return true, responseData
+}

api/oauth_api/qq_api.go

@@ -0,0 +1 @@
+package oauth_api

api/oauth_api/wechat_api.go

@@ -23,20 +23,12 @@ import (
 	"schisandra-cloud-album/common/result"
 	"schisandra-cloud-album/global"
 	"schisandra-cloud-album/model"
-	"schisandra-cloud-album/service"
 	"schisandra-cloud-album/utils"
 	"strconv"
 	"strings"
 	"time"
 )
 
-var userService = service.Service.UserService
-var userRoleService = service.Service.UserRoleService
-var userSocialService = service.Service.UserSocialService
-var rolePermissionService = service.Service.RolePermissionService
-var permissionServiceService = service.Service.PermissionService
-var roleService = service.Service.RoleService
-
 // GenerateClientId 生成客户端ID
 // @Summary 生成客户端ID
 // @Description 生成客户端ID
@@ -281,11 +273,16 @@ func handelUserLogin(user model.ScaAuthUser, clientId string) bool {
 		UID:          user.UID,
 	}
 	fail := redis.Set(constant.UserLoginTokenRedisKey+*user.UID, data, time.Hour*24*7).Err()
-	w := redis.Set(constant.UserLoginWechatRedisKey+clientId, data, time.Minute*5).Err()
-	if fail != nil || w != nil {
+	if fail != nil {
 		return false
 	}
-	res := websocket_api.SendMessageData(clientId, data)
+	responseData := map[string]interface{}{
+		"code":    0,
+		"message": "success",
+		"data":    data,
+		"success": true,
+	}
+	res := websocket_api.SendMessageData(clientId, responseData)
 	if !res {
 		return false
 	}

api/user_api/user_api.go

@@ -43,7 +43,7 @@ func (UserAPI) GetUserList(c *gin.Context) {
 // @Router /api/auth/user/query_by_username [get]
 func (UserAPI) QueryUserByUsername(c *gin.Context) {
 	username := c.Query("username")
-	user := userService.QueryUserByUsername(username)
+	user, _ := userService.QueryUserByUsername(username)
 	if reflect.DeepEqual(user, model.ScaAuthUser{}) {
 		result.FailWithMessage(ginI18n.MustGetMessage(c, "NotFoundUser"), c)
 		return
@@ -117,7 +117,7 @@ func (UserAPI) AddUser(c *gin.Context) {
 		return
 	}
 
-	username := userService.QueryUserByUsername(addUserRequest.Username)
+	username, _ := userService.QueryUserByUsername(addUserRequest.Username)
 	if !reflect.DeepEqual(username, model.ScaAuthUser{}) {
 		result.FailWithMessage(ginI18n.MustGetMessage(c, "UsernameExists"), c)
 		return
@@ -214,7 +214,7 @@ func (UserAPI) AccountLogin(c *gin.Context) {
 	}
 	isUsername := utils.IsUsername(account)
 	if isUsername {
-		user := userService.QueryUserByUsername(account)
+		user, _ := userService.QueryUserByUsername(account)
 		if reflect.DeepEqual(user, model.ScaAuthUser{}) {
 			result.FailWithMessage(ginI18n.MustGetMessage(c, "UsernameNotRegister"), c)
 			return

api/websocket_api/gws_api.go

@@ -0,0 +1,104 @@
+package websocket_api
+
+import (
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/lxzan/gws"
+	"net/http"
+	"schisandra-cloud-album/global"
+	"time"
+)
+
+const (
+	PingInterval         = 5 * time.Second  // 客户端心跳间隔
+	HeartbeatWaitTimeout = 10 * time.Second // 心跳等待超时时间
+)
+
+func (WebsocketAPI) NewGWSServer(c *gin.Context) {
+	var handler = NewWebSocket()
+	upgrader := gws.NewUpgrader(handler, &gws.ServerOption{
+		HandshakeTimeout: 5 * time.Second, // 握手超时时间
+		ReadBufferSize:   1024,            // 读缓冲区大小
+		ParallelEnabled:  true,            // 开启并行消息处理
+		Recovery:         gws.Recovery,    // 开启异常恢复
+		CheckUtf8Enabled: true,            // 开启UTF8校验
+		PermessageDeflate: gws.PermessageDeflate{
+			Enabled: true, // 开启压缩
+		},
+		Authorize: func(r *http.Request, session gws.SessionStorage) bool {
+			var clientId = r.URL.Query().Get("client_id")
+			if clientId == "" {
+				return false
+			}
+			session.Store("client_id", clientId)
+			return true
+		},
+	})
+	socket, err := upgrader.Upgrade(c.Writer, c.Request)
+	if err != nil {
+		return
+	}
+	go func() {
+		socket.ReadLoop() // 此处阻塞会使请求上下文不能顺利被GC
+	}()
+}
+func MustLoad[T any](session gws.SessionStorage, key string) (v T) {
+	if value, exist := session.Load(key); exist {
+		v = value.(T)
+	}
+	return
+}
+
+func NewWebSocket() *WebSocket {
+	return &WebSocket{
+		sessions: gws.NewConcurrentMap[string, *gws.Conn](64, 128),
+	}
+}
+
+type WebSocket struct {
+	gws.BuiltinEventHandler
+	sessions *gws.ConcurrentMap[string, *gws.Conn] // 使用内置的ConcurrentMap存储连接, 可以减少锁冲突
+}
+
+func (c *WebSocket) OnOpen(socket *gws.Conn) {
+	name := MustLoad[string](socket.Session(), "client_id")
+	if conn, ok := c.sessions.Load(name); ok {
+		conn.WriteClose(1000, []byte("connection is replaced"))
+	}
+	c.sessions.Store(name, socket)
+	global.LOG.Printf("%s connected\n", name)
+}
+
+func (c *WebSocket) OnClose(socket *gws.Conn, err error) {
+	name := MustLoad[string](socket.Session(), "client_id")
+	sharding := c.sessions.GetSharding(name)
+	c.sessions.Delete(name)
+	sharding.Lock()
+	defer sharding.Unlock()
+
+	global.LOG.Printf("onerror, name=%s, msg=%s\n", name, err.Error())
+}
+
+func (c *WebSocket) OnPing(socket *gws.Conn, payload []byte) {
+	_ = socket.SetDeadline(time.Now().Add(PingInterval + HeartbeatWaitTimeout))
+	_ = socket.WritePong(payload)
+}
+
+func (c *WebSocket) OnPong(socket *gws.Conn, payload []byte) {}
+
+func (c *WebSocket) OnMessage(socket *gws.Conn, message *gws.Message) {
+	defer message.Close()
+	name := MustLoad[string](socket.Session(), "client_id")
+	if conn, ok := c.sessions.Load(name); ok {
+		_ = conn.WriteMessage(gws.OpcodeText, message.Bytes())
+	}
+}
+
+// SendMessageToClient 向指定客户端发送消息
+func (c *WebSocket) SendMessageToClient(clientId string, message []byte) error {
+	conn, ok := c.sessions.Load(clientId)
+	if ok {
+		return conn.WriteMessage(gws.OpcodeText, message)
+	}
+	return fmt.Errorf("client %s not found", clientId)
+}

api/websocket_api/websocket_api.go

@@ -1,11 +1,11 @@
 package websocket_api
 
 import (
-	"fmt"
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
 	"log"
 	"net/http"
+	"schisandra-cloud-album/common/result"
 	"schisandra-cloud-album/global"
 	"sync"
 	"time"
@@ -40,9 +40,7 @@ func (WebsocketAPI) DeleteClient(context *gin.Context) {
 		}
 		deleteClient(id)
 	} else {
-		context.JSON(http.StatusOK, gin.H{
-			"mesg": "未找到该客户端",
-		})
+		result.FailWithMessage("客户端不存在", context)
 	}
 	// 关闭其消息通道
 	_, exist = getMsgChannel(id)
@@ -86,7 +84,7 @@ func WsHandler(w http.ResponseWriter, r *http.Request, id string) {
 	var err error
 	var exist bool
 	// 创建一个定时器用于服务端心跳
-	//pingTicker := time.NewTicker(time.Second * 10)
+	pingTicker := time.NewTicker(time.Second * 10)
 	conn, err = wsupgrader.Upgrade(w, r, nil)
 	if err != nil {
 		global.LOG.Println(err)
@@ -110,7 +108,7 @@ func WsHandler(w http.ResponseWriter, r *http.Request, id string) {
 	for {
 		_, _, err := conn.ReadMessage()
 		if err != nil {
-			fmt.Println(err)
+			global.LOG.Error(err)
 			return
 		}
 		select {
@@ -118,30 +116,29 @@ func WsHandler(w http.ResponseWriter, r *http.Request, id string) {
 			// 从消息通道接收消息，然后推送给前端
 			err = conn.WriteJSON(content)
 			if err != nil {
-				global.LOG.Error(err)
-				err := conn.Close()
+				err = conn.Close()
 				if err != nil {
 					return
 				}
 				deleteClient(id)
 				break
 			}
-			//case <-pingTicker.C:
-			//	// 服务端心跳:每20秒ping一次客户端，查看其是否在线
-			//	err := conn.SetWriteDeadline(time.Now().Add(time.Second * 20))
-			//	if err != nil {
-			//		return
-			//	}
-			//	err = conn.WriteMessage(websocket.PongMessage, []byte("pong"))
-			//	if err != nil {
-			//		log.Println("send pong err:", err)
-			//		err := conn.Close()
-			//		if err != nil {
-			//			return
-			//		}
-			//		deleteClient(id)
-			//		return
-			//	}
+		case <-pingTicker.C:
+			// 服务端心跳:每20秒ping一次客户端，查看其是否在线
+			err := conn.SetWriteDeadline(time.Now().Add(time.Second * 20))
+			if err != nil {
+				return
+			}
+			err = conn.WriteMessage(websocket.PingMessage, []byte{})
+			if err != nil {
+				log.Println("send pong err:", err)
+				err := conn.Close()
+				if err != nil {
+					return
+				}
+				deleteClient(id)
+				return
+			}
 		}
 	}