(cherry picked from commit a1dd79a459)
pull/3642/head
withchao
9 months ago
parent
dd5ff6f0e4
commit
efa709c267
@ -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