mirror of
https://github.com/yusing/godoxy.git
synced 2026-04-23 09:18:51 +02:00
refactor(api): restructured API for type safety, maintainability and docs generation
- These changes makes the API incombatible with previous versions - Added new types for error handling, success responses, and health checks. - Updated health check logic to utilize the new types for better clarity and structure. - Refactored existing handlers to improve response consistency and error handling. - Updated Makefile to include a new target for generating API types from Swagger. - Updated "new agent" API to respond an encrypted cert pair
This commit is contained in:
yusing
252
internal/net/gphttp/websocket/manager.go
Normal file
252
internal/net/gphttp/websocket/manager.go
Normal file
@@ -0,0 +1,252 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/yusing/go-proxy/internal/common"
|
||||
)
|
||||
|
||||
// Manager handles WebSocket connection state and ping-pong
|
||||
type Manager struct {
|
||||
conn *websocket.Conn
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
pongWriteTimeout time.Duration
|
||||
pingCheckTicker *time.Ticker
|
||||
lastPingTime atomic.Value
|
||||
readCh chan []byte
|
||||
err error
|
||||
}
|
||||
|
||||
var defaultUpgrader = websocket.Upgrader{
|
||||
ReadBufferSize: 4096,
|
||||
WriteBufferSize: 4096,
|
||||
// TODO: add CORS
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true
|
||||
},
|
||||
}
|
||||
|
||||
var (
|
||||
ErrReadTimeout = errors.New("read timeout")
|
||||
ErrWriteTimeout = errors.New("write timeout")
|
||||
)
|
||||
|
||||
const (
|
||||
TextMessage = websocket.TextMessage
|
||||
BinaryMessage = websocket.BinaryMessage
|
||||
)
|
||||
|
||||
// NewManagerWithUpgrade upgrades the HTTP connection to a WebSocket connection and returns a Manager.
|
||||
// If the upgrade fails, the error is returned.
|
||||
// If the upgrade succeeds, the Manager is returned.
|
||||
func NewManagerWithUpgrade(c *gin.Context, upgrader ...websocket.Upgrader) (*Manager, error) {
|
||||
var actualUpgrader websocket.Upgrader
|
||||
if len(upgrader) == 0 {
|
||||
actualUpgrader = defaultUpgrader
|
||||
} else {
|
||||
actualUpgrader = upgrader[0]
|
||||
}
|
||||
|
||||
conn, err := actualUpgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(c.Request.Context())
|
||||
cm := &Manager{
|
||||
conn: conn,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
pongWriteTimeout: 2 * time.Second,
|
||||
pingCheckTicker: time.NewTicker(3 * time.Second),
|
||||
readCh: make(chan []byte, 1),
|
||||
}
|
||||
cm.lastPingTime.Store(time.Now())
|
||||
|
||||
conn.SetCloseHandler(func(code int, text string) error {
|
||||
if common.IsDebug {
|
||||
cm.err = fmt.Errorf("connection closed: code=%d, text=%s", code, text)
|
||||
}
|
||||
cm.Close()
|
||||
return nil
|
||||
})
|
||||
|
||||
go cm.pingCheckRoutine()
|
||||
go cm.readRoutine()
|
||||
|
||||
return cm, nil
|
||||
}
|
||||
|
||||
// Periodic writes data to the connection periodically.
|
||||
// If the connection is closed, the error is returned.
|
||||
// If the write timeout is reached, ErrWriteTimeout is returned.
|
||||
func (cm *Manager) PeriodicWrite(interval time.Duration, getData func() (any, error)) error {
|
||||
write := func() {
|
||||
data, err := getData()
|
||||
if err != nil {
|
||||
cm.err = err
|
||||
cm.Close()
|
||||
return
|
||||
}
|
||||
if err := cm.WriteJSON(data, interval); err != nil {
|
||||
cm.err = err
|
||||
cm.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// initial write before the ticker starts
|
||||
write()
|
||||
if cm.err != nil {
|
||||
return cm.err
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-cm.ctx.Done():
|
||||
return cm.err
|
||||
case <-ticker.C:
|
||||
write()
|
||||
if cm.err != nil {
|
||||
return cm.err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WriteJSON writes a JSON message to the connection with json.
|
||||
// If the connection is closed, the error is returned.
|
||||
// If the write timeout is reached, ErrWriteTimeout is returned.
|
||||
func (cm *Manager) WriteJSON(data any, timeout time.Duration) error {
|
||||
bytes, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return cm.WriteData(websocket.TextMessage, bytes, timeout)
|
||||
}
|
||||
|
||||
// WriteData writes a message to the connection with sonic.
|
||||
// If the connection is closed, the error is returned.
|
||||
// If the write timeout is reached, ErrWriteTimeout is returned.
|
||||
func (cm *Manager) WriteData(typ int, data []byte, timeout time.Duration) error {
|
||||
select {
|
||||
case <-cm.ctx.Done():
|
||||
return cm.err
|
||||
default:
|
||||
if err := cm.conn.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
|
||||
return err
|
||||
}
|
||||
err := cm.conn.WriteMessage(typ, data)
|
||||
if err != nil {
|
||||
if errors.Is(err, websocket.ErrCloseSent) {
|
||||
return cm.err
|
||||
}
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return ErrWriteTimeout
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// ReadJSON reads a JSON message from the connection and unmarshals it into the provided struct with sonic
|
||||
// If the connection is closed, the error is returned.
|
||||
// If the message fails to unmarshal, the error is returned.
|
||||
// If the read timeout is reached, ErrReadTimeout is returned.
|
||||
func (cm *Manager) ReadJSON(out any, timeout time.Duration) error {
|
||||
select {
|
||||
case <-cm.ctx.Done():
|
||||
return cm.err
|
||||
case data := <-cm.readCh:
|
||||
return json.Unmarshal(data, out)
|
||||
case <-time.After(timeout):
|
||||
return ErrReadTimeout
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes the connection and cancels the context
|
||||
func (cm *Manager) Close() {
|
||||
cm.cancel()
|
||||
cm.pingCheckTicker.Stop()
|
||||
cm.conn.Close()
|
||||
}
|
||||
|
||||
func (cm *Manager) GracefulClose() {
|
||||
_ = cm.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
|
||||
_ = cm.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
|
||||
cm.Close()
|
||||
}
|
||||
|
||||
// Done returns a channel that is closed when the context is done or the connection is closed
|
||||
func (cm *Manager) Done() <-chan struct{} {
|
||||
return cm.ctx.Done()
|
||||
}
|
||||
|
||||
func (cm *Manager) pingCheckRoutine() {
|
||||
for {
|
||||
select {
|
||||
case <-cm.ctx.Done():
|
||||
return
|
||||
case <-cm.pingCheckTicker.C:
|
||||
if time.Since(cm.lastPingTime.Load().(time.Time)) > 5*time.Second {
|
||||
if common.IsDebug {
|
||||
cm.err = errors.New("no ping received in 5 seconds, closing connection")
|
||||
}
|
||||
cm.Close()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (cm *Manager) readRoutine() {
|
||||
for {
|
||||
select {
|
||||
case <-cm.ctx.Done():
|
||||
return
|
||||
default:
|
||||
typ, data, err := cm.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if cm.ctx.Err() == nil { // connection is not closed
|
||||
cm.err = fmt.Errorf("failed to read message: %w", err)
|
||||
cm.Close()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if typ == websocket.TextMessage && string(data) == "ping" {
|
||||
cm.lastPingTime.Store(time.Now())
|
||||
if err := cm.conn.SetWriteDeadline(time.Now().Add(cm.pongWriteTimeout)); err != nil {
|
||||
cm.err = fmt.Errorf("failed to set write deadline: %w", err)
|
||||
cm.Close()
|
||||
return
|
||||
}
|
||||
if err := cm.conn.WriteMessage(websocket.TextMessage, []byte("pong")); err != nil {
|
||||
cm.err = fmt.Errorf("failed to write pong message: %w", err)
|
||||
cm.Close()
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if typ == websocket.TextMessage || typ == websocket.BinaryMessage {
|
||||
select {
|
||||
case <-cm.ctx.Done():
|
||||
return
|
||||
case cm.readCh <- data:
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user