chore(服务端): 初始化工程仓库
This commit is contained in:
leon_mac
@@ -0,0 +1,266 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"perler-beads-server/internal/domain"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
wsAuthTimeout = 5 * time.Second
|
||||
wsHeartbeat = 20 * time.Second
|
||||
wsReadTimeout = 45 * time.Second
|
||||
wsWriteTimeout = 5 * time.Second
|
||||
wsRedisTimeout = 3 * time.Second
|
||||
wsMessageLimit = 4096
|
||||
)
|
||||
|
||||
type wsClose struct {
|
||||
code int
|
||||
reason string
|
||||
}
|
||||
type taskSocket struct {
|
||||
conn *websocket.Conn
|
||||
user domain.User
|
||||
token string
|
||||
changed chan struct{}
|
||||
pong chan string
|
||||
stop chan wsClose
|
||||
}
|
||||
type taskHub struct {
|
||||
mu sync.Mutex
|
||||
stores map[string]map[*taskSocket]struct{}
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newTaskHub() *taskHub { return &taskHub{stores: map[string]map[*taskSocket]struct{}{}} }
|
||||
func (h *taskHub) add(client *taskSocket) bool {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.closed {
|
||||
return false
|
||||
}
|
||||
if h.stores[client.user.StoreID] == nil {
|
||||
h.stores[client.user.StoreID] = map[*taskSocket]struct{}{}
|
||||
}
|
||||
h.stores[client.user.StoreID][client] = struct{}{}
|
||||
return true
|
||||
}
|
||||
func (h *taskHub) remove(client *taskSocket) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
store := h.stores[client.user.StoreID]
|
||||
delete(store, client)
|
||||
if len(store) == 0 {
|
||||
delete(h.stores, client.user.StoreID)
|
||||
}
|
||||
}
|
||||
|
||||
// Capacity one merges invalidations; slow browsers never block a committed business request.
|
||||
func (h *taskHub) notify(storeID string) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
for client := range h.stores[storeID] {
|
||||
select {
|
||||
case client.changed <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
func (h *taskHub) shutdown() {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.closed = true
|
||||
for _, store := range h.stores {
|
||||
for client := range store {
|
||||
client.close(1001, "server shutting down")
|
||||
}
|
||||
}
|
||||
}
|
||||
func (c *taskSocket) close(code int, reason string) {
|
||||
select {
|
||||
case c.stop <- wsClose{code, reason}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
func (s *Server) notifyTasks(storeID string) {
|
||||
if s.wsHub != nil {
|
||||
s.wsHub.notify(storeID)
|
||||
}
|
||||
}
|
||||
func (s *Server) CloseWebSockets() {
|
||||
if s.wsHub != nil {
|
||||
s.wsHub.shutdown()
|
||||
}
|
||||
}
|
||||
func (s *Server) wsOrigin(r *http.Request) bool {
|
||||
origin := r.Header.Get("Origin")
|
||||
if origin == "" || len(r.Header.Values("Origin")) != 1 {
|
||||
return false
|
||||
}
|
||||
u, e := url.Parse(origin)
|
||||
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.Fragment != "" {
|
||||
return false
|
||||
}
|
||||
if strings.EqualFold(u.Host, r.Host) {
|
||||
return true
|
||||
}
|
||||
for _, allowed := range s.Origins {
|
||||
if allowed != "" && origin == allowed {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
func (s *Server) wsUser(token string) (domain.User, error) {
|
||||
var u domain.User
|
||||
if token == "" || len(token) > 256 || s.Redis == nil {
|
||||
return u, errors.New("invalid session")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), wsRedisTimeout)
|
||||
defer cancel()
|
||||
h := sha256.Sum256([]byte(token))
|
||||
b, e := s.Redis.Get(ctx, "session:"+hex.EncodeToString(h[:])).Bytes()
|
||||
if e != nil {
|
||||
return u, e
|
||||
}
|
||||
if e = json.Unmarshal(b, &u); e != nil {
|
||||
return u, e
|
||||
}
|
||||
if u.ID == "" || u.StoreID == "" || (u.Role != "admin" && u.Role != "staff") {
|
||||
return u, errors.New("staff session required")
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
func (s *Server) websocket(c *gin.Context) {
|
||||
if !s.wsOrigin(c.Request) {
|
||||
c.Status(http.StatusForbidden)
|
||||
return
|
||||
}
|
||||
if c.Request.URL.RawQuery != "" {
|
||||
c.Status(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
upgrader := websocket.Upgrader{HandshakeTimeout: wsAuthTimeout, ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: s.wsOrigin}
|
||||
conn, e := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if e != nil {
|
||||
return
|
||||
}
|
||||
conn.SetReadLimit(wsMessageLimit)
|
||||
conn.SetReadDeadline(time.Now().Add(wsAuthTimeout))
|
||||
conn.SetCloseHandler(func(code int, text string) error { return &websocket.CloseError{Code: code, Text: text} })
|
||||
var auth struct {
|
||||
Type string `json:"type"`
|
||||
Token string `json:"token"`
|
||||
}
|
||||
kind, payload, e := conn.ReadMessage()
|
||||
if e != nil {
|
||||
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(4403, "authentication required"), time.Now().Add(wsWriteTimeout))
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
if kind != websocket.TextMessage || json.Unmarshal(payload, &auth) != nil || auth.Type != "auth" {
|
||||
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(4403, "authentication required"), time.Now().Add(wsWriteTimeout))
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
u, e := s.wsUser(auth.Token)
|
||||
if e != nil {
|
||||
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(4403, "invalid staff session"), time.Now().Add(wsWriteTimeout))
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
client := &taskSocket{conn: conn, user: u, token: auth.Token, changed: make(chan struct{}, 1), pong: make(chan string, 4), stop: make(chan wsClose, 1)}
|
||||
if !s.wsHub.add(client) {
|
||||
conn.Close()
|
||||
return
|
||||
}
|
||||
defer s.wsHub.remove(client)
|
||||
readerDone := make(chan struct{})
|
||||
writerDone := make(chan struct{})
|
||||
conn.SetReadDeadline(time.Now().Add(wsReadTimeout))
|
||||
conn.SetPongHandler(func(string) error { return conn.SetReadDeadline(time.Now().Add(wsReadTimeout)) })
|
||||
conn.SetPingHandler(func(data string) error {
|
||||
select {
|
||||
case client.pong <- data:
|
||||
return nil
|
||||
default:
|
||||
return errors.New("too many pings")
|
||||
}
|
||||
})
|
||||
go func() { defer close(writerDone); s.writeSocket(client, readerDone) }()
|
||||
defer func() { close(readerDone); <-writerDone }()
|
||||
for {
|
||||
kind, payload, e = conn.ReadMessage()
|
||||
if e != nil {
|
||||
return
|
||||
}
|
||||
var message struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if kind != websocket.TextMessage || json.Unmarshal(payload, &message) != nil || message.Type != "pong" {
|
||||
client.close(4400, "unexpected message")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
func (s *Server) writeSocket(client *taskSocket, readerDone <-chan struct{}) {
|
||||
defer client.conn.Close()
|
||||
write := func(kind string) error {
|
||||
client.conn.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
|
||||
return client.conn.WriteJSON(map[string]string{"type": kind})
|
||||
}
|
||||
closeFrame := func(reason wsClose) {
|
||||
client.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(reason.code, reason.reason), time.Now().Add(wsWriteTimeout))
|
||||
}
|
||||
if e := write("ready"); e != nil {
|
||||
return
|
||||
}
|
||||
ticker := time.NewTicker(wsHeartbeat)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case reason := <-client.stop:
|
||||
closeFrame(reason)
|
||||
return
|
||||
case <-readerDone:
|
||||
select {
|
||||
case reason := <-client.stop:
|
||||
closeFrame(reason)
|
||||
default:
|
||||
closeFrame(wsClose{1000, "connection closed"})
|
||||
}
|
||||
return
|
||||
case data := <-client.pong:
|
||||
if e := client.conn.WriteControl(websocket.PongMessage, []byte(data), time.Now().Add(wsWriteTimeout)); e != nil {
|
||||
return
|
||||
}
|
||||
case <-client.changed:
|
||||
if e := write("tasks.changed"); e != nil {
|
||||
return
|
||||
}
|
||||
case <-ticker.C:
|
||||
u, e := s.wsUser(client.token)
|
||||
if e != nil || u.ID != client.user.ID || u.Role != client.user.Role || u.StoreID != client.user.StoreID {
|
||||
closeFrame(wsClose{4403, "session expired"})
|
||||
return
|
||||
}
|
||||
if e := client.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(wsWriteTimeout)); e != nil {
|
||||
return
|
||||
}
|
||||
if e := write("heartbeat"); e != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user