Merge 9da7db2ac2 into 6f33c0a515
chao
1 day ago
committed by
GitHub
commit
fbc3942250
@ -0,0 +1,229 @@
|
||||
package msggateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
|
||||
"github.com/openimsdk/tools/log"
|
||||
)
|
||||
|
||||
var ErrWriteFull = fmt.Errorf("websocket write buffer full,close connection")
|
||||
|
||||
type ClientConn interface {
|
||||
ReadMessage() ([]byte, error)
|
||||
WriteMessage(message []byte) error
|
||||
Close() error
|
||||
}
|
||||
|
||||
type websocketMessage struct {
|
||||
MessageType int
|
||||
Data []byte
|
||||
}
|
||||
|
||||
func NewWebSocketClientConn(conn *websocket.Conn, readLimit int64, readTimeout time.Duration, pingInterval time.Duration) ClientConn {
|
||||
c := &websocketClientConn{
|
||||
readTimeout: readTimeout,
|
||||
conn: conn,
|
||||
writer: make(chan *websocketMessage, 256),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
if readLimit > 0 {
|
||||
c.conn.SetReadLimit(readLimit)
|
||||
}
|
||||
c.conn.SetPingHandler(c.pingHandler)
|
||||
c.conn.SetPongHandler(c.pongHandler)
|
||||
|
||||
go c.loopSend()
|
||||
if pingInterval > 0 {
|
||||
go c.doPing(pingInterval)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
type websocketClientConn struct {
|
||||
readTimeout time.Duration
|
||||
conn *websocket.Conn
|
||||
writer chan *websocketMessage
|
||||
done chan struct{}
|
||||
err atomic.Pointer[error]
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) ReadMessage() ([]byte, error) {
|
||||
buf, err := c.readMessage()
|
||||
if err != nil {
|
||||
return nil, c.closeBy(fmt.Errorf("read message %w", err))
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) WriteMessage(message []byte) error {
|
||||
return c.writeMessage(websocket.BinaryMessage, message)
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) Close() error {
|
||||
return c.closeBy(fmt.Errorf("websocket connection closed"))
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) closeBy(err error) error {
|
||||
if !c.err.CompareAndSwap(nil, &err) {
|
||||
return *c.err.Load()
|
||||
}
|
||||
close(c.done)
|
||||
log.ZWarn(context.Background(), "websocket connection closed", err, "remoteAddr", c.conn.RemoteAddr(),
|
||||
"chan length", len(c.writer))
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) writeMessage(messageType int, data []byte) error {
|
||||
if errPtr := c.err.Load(); errPtr != nil {
|
||||
return *errPtr
|
||||
}
|
||||
select {
|
||||
case c.writer <- &websocketMessage{MessageType: messageType, Data: data}:
|
||||
return nil
|
||||
default:
|
||||
return c.closeBy(ErrWriteFull)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) loopSend() {
|
||||
defer func() {
|
||||
_ = c.conn.Close()
|
||||
}()
|
||||
var err error
|
||||
for {
|
||||
select {
|
||||
case <-c.done:
|
||||
for {
|
||||
select {
|
||||
case msg := <-c.writer:
|
||||
switch msg.MessageType {
|
||||
case websocket.TextMessage, websocket.BinaryMessage:
|
||||
err = c.conn.WriteMessage(msg.MessageType, msg.Data)
|
||||
default:
|
||||
err = c.conn.WriteControl(msg.MessageType, msg.Data, time.Time{})
|
||||
}
|
||||
if err != nil {
|
||||
_ = c.closeBy(err)
|
||||
return
|
||||
}
|
||||
default:
|
||||
return
|
||||
}
|
||||
}
|
||||
case msg := <-c.writer:
|
||||
switch msg.MessageType {
|
||||
case websocket.TextMessage, websocket.BinaryMessage:
|
||||
err = c.conn.WriteMessage(msg.MessageType, msg.Data)
|
||||
default:
|
||||
err = c.conn.WriteControl(msg.MessageType, msg.Data, time.Time{})
|
||||
}
|
||||
if err != nil {
|
||||
_ = c.closeBy(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) setReadDeadline() error {
|
||||
deadline := time.Now().Add(c.readTimeout)
|
||||
return c.conn.SetReadDeadline(deadline)
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) readMessage() ([]byte, error) {
|
||||
for {
|
||||
if err := c.setReadDeadline(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
messageType, buf, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch messageType {
|
||||
case websocket.BinaryMessage:
|
||||
return buf, nil
|
||||
case websocket.TextMessage:
|
||||
if err := c.onReadTextMessage(buf); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
case websocket.PingMessage:
|
||||
if err := c.pingHandler(string(buf)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
case websocket.PongMessage:
|
||||
if err := c.pongHandler(string(buf)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
case websocket.CloseMessage:
|
||||
if len(buf) == 0 {
|
||||
return nil, errors.New("websocket connection closed by peer")
|
||||
}
|
||||
return nil, fmt.Errorf("websocket connection closed by peer, data %s", string(buf))
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown websocket message type %d", messageType)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) onReadTextMessage(buf []byte) error {
|
||||
var msg struct {
|
||||
Type string `json:"type"`
|
||||
Body json.RawMessage `json:"body"`
|
||||
}
|
||||
if err := json.Unmarshal(buf, &msg); err != nil {
|
||||
return err
|
||||
}
|
||||
switch msg.Type {
|
||||
case TextPong:
|
||||
return nil
|
||||
case TextPing:
|
||||
msg.Type = TextPong
|
||||
msgData, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.writeMessage(websocket.TextMessage, msgData)
|
||||
default:
|
||||
return fmt.Errorf("not support text message type %s", msg.Type)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) pingHandler(appData string) error {
|
||||
log.ZDebug(context.Background(), "ping handler recv ping", "remoteAddr", c.conn.RemoteAddr(), "appData", appData)
|
||||
if err := c.setReadDeadline(); err != nil {
|
||||
return err
|
||||
}
|
||||
err := c.conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second*1))
|
||||
if err != nil {
|
||||
log.ZWarn(context.Background(), "ping handler write pong error", err, "remoteAddr", c.conn.RemoteAddr(), "appData", appData)
|
||||
}
|
||||
log.ZDebug(context.Background(), "ping handler write pong success", "remoteAddr", c.conn.RemoteAddr(), "appData", appData)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) pongHandler(string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *websocketClientConn) doPing(d time.Duration) {
|
||||
ticker := time.NewTicker(d)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-c.done:
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := c.writeMessage(websocket.PingMessage, nil); err != nil {
|
||||
_ = c.closeBy(fmt.Errorf("send ping %w", err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,179 +0,0 @@
|
||||
// Copyright © 2023 OpenIM. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package msggateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/openimsdk/tools/apiresp"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/openimsdk/tools/errs"
|
||||
)
|
||||
|
||||
type LongConn interface {
|
||||
// Close this connection
|
||||
Close() error
|
||||
// WriteMessage Write message to connection,messageType means data type,can be set binary(2) and text(1).
|
||||
WriteMessage(messageType int, message []byte) error
|
||||
// ReadMessage Read message from connection.
|
||||
ReadMessage() (int, []byte, error)
|
||||
// SetReadDeadline sets the read deadline on the underlying network connection,
|
||||
// after a read has timed out, will return an error.
|
||||
SetReadDeadline(timeout time.Duration) error
|
||||
// SetWriteDeadline sets to write deadline when send message,when read has timed out,will return error.
|
||||
SetWriteDeadline(timeout time.Duration) error
|
||||
// Dial Try to dial a connection,url must set auth args,header can control compress data
|
||||
Dial(urlStr string, requestHeader http.Header) (*http.Response, error)
|
||||
// IsNil Whether the connection of the current long connection is nil
|
||||
IsNil() bool
|
||||
// SetConnNil Set the connection of the current long connection to nil
|
||||
SetConnNil()
|
||||
// SetReadLimit sets the maximum size for a message read from the peer.bytes
|
||||
SetReadLimit(limit int64)
|
||||
SetPongHandler(handler PingPongHandler)
|
||||
SetPingHandler(handler PingPongHandler)
|
||||
// GenerateLongConn Check the connection of the current and when it was sent are the same
|
||||
GenerateLongConn(w http.ResponseWriter, r *http.Request) error
|
||||
}
|
||||
type GWebSocket struct {
|
||||
protocolType int
|
||||
conn *websocket.Conn
|
||||
handshakeTimeout time.Duration
|
||||
writeBufferSize int
|
||||
}
|
||||
|
||||
func newGWebSocket(protocolType int, handshakeTimeout time.Duration, wbs int) *GWebSocket {
|
||||
return &GWebSocket{protocolType: protocolType, handshakeTimeout: handshakeTimeout, writeBufferSize: wbs}
|
||||
}
|
||||
|
||||
func (d *GWebSocket) Close() error {
|
||||
return d.conn.Close()
|
||||
}
|
||||
|
||||
func (d *GWebSocket) GenerateLongConn(w http.ResponseWriter, r *http.Request) error {
|
||||
upgrader := &websocket.Upgrader{
|
||||
HandshakeTimeout: d.handshakeTimeout,
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
}
|
||||
if d.writeBufferSize > 0 { // default is 4kb.
|
||||
upgrader.WriteBufferSize = d.writeBufferSize
|
||||
}
|
||||
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
// The upgrader.Upgrade method usually returns enough error messages to diagnose problems that may occur during the upgrade
|
||||
return errs.WrapMsg(err, "GenerateLongConn: WebSocket upgrade failed")
|
||||
}
|
||||
d.conn = conn
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *GWebSocket) WriteMessage(messageType int, message []byte) error {
|
||||
// d.setSendConn(d.conn)
|
||||
return d.conn.WriteMessage(messageType, message)
|
||||
}
|
||||
|
||||
// func (d *GWebSocket) setSendConn(sendConn *websocket.Conn) {
|
||||
// d.sendConn = sendConn
|
||||
//}
|
||||
|
||||
func (d *GWebSocket) ReadMessage() (int, []byte, error) {
|
||||
return d.conn.ReadMessage()
|
||||
}
|
||||
|
||||
func (d *GWebSocket) SetReadDeadline(timeout time.Duration) error {
|
||||
return d.conn.SetReadDeadline(time.Now().Add(timeout))
|
||||
}
|
||||
|
||||
func (d *GWebSocket) SetWriteDeadline(timeout time.Duration) error {
|
||||
if timeout <= 0 {
|
||||
return errs.New("timeout must be greater than 0")
|
||||
}
|
||||
|
||||
// TODO SetWriteDeadline Future add error handling
|
||||
if err := d.conn.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
|
||||
return errs.WrapMsg(err, "GWebSocket.SetWriteDeadline failed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *GWebSocket) Dial(urlStr string, requestHeader http.Header) (*http.Response, error) {
|
||||
conn, httpResp, err := websocket.DefaultDialer.Dial(urlStr, requestHeader)
|
||||
if err != nil {
|
||||
return httpResp, errs.WrapMsg(err, "GWebSocket.Dial failed", "url", urlStr)
|
||||
}
|
||||
d.conn = conn
|
||||
return httpResp, nil
|
||||
}
|
||||
|
||||
func (d *GWebSocket) IsNil() bool {
|
||||
return d.conn == nil
|
||||
//
|
||||
// if d.conn != nil {
|
||||
// return false
|
||||
// }
|
||||
// return true
|
||||
}
|
||||
|
||||
func (d *GWebSocket) SetConnNil() {
|
||||
d.conn = nil
|
||||
}
|
||||
|
||||
func (d *GWebSocket) SetReadLimit(limit int64) {
|
||||
d.conn.SetReadLimit(limit)
|
||||
}
|
||||
|
||||
func (d *GWebSocket) SetPongHandler(handler PingPongHandler) {
|
||||
d.conn.SetPongHandler(handler)
|
||||
}
|
||||
|
||||
func (d *GWebSocket) SetPingHandler(handler PingPongHandler) {
|
||||
d.conn.SetPingHandler(handler)
|
||||
}
|
||||
|
||||
func (d *GWebSocket) RespondWithError(err error, w http.ResponseWriter, r *http.Request) error {
|
||||
if err := d.GenerateLongConn(w, r); err != nil {
|
||||
return err
|
||||
}
|
||||
data, err := json.Marshal(apiresp.ParseError(err))
|
||||
if err != nil {
|
||||
_ = d.Close()
|
||||
return errs.WrapMsg(err, "json marshal failed")
|
||||
}
|
||||
|
||||
if err := d.WriteMessage(MessageText, data); err != nil {
|
||||
_ = d.Close()
|
||||
return errs.WrapMsg(err, "WriteMessage failed")
|
||||
}
|
||||
_ = d.Close()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *GWebSocket) RespondWithSuccess() error {
|
||||
data, err := json.Marshal(apiresp.ParseError(nil))
|
||||
if err != nil {
|
||||
_ = d.Close()
|
||||
return errs.WrapMsg(err, "json marshal failed")
|
||||
}
|
||||
|
||||
if err := d.WriteMessage(MessageText, data); err != nil {
|
||||
_ = d.Close()
|
||||
return errs.WrapMsg(err, "WriteMessage failed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Loading…
Reference in new issue