live-gateway/live/live.go

package live

import (
	"github.com/pkg/errors"
	"live-gateway/pkg/logger"
	"live-gateway/ws"
)

// InitFunc 初始化方法
type InitFunc func(wsConn *ws.NWebsocket) (err error)

// PreConnectFunc 连接前hook
type PreConnectFunc func() (url string, err error)

// MessageHandlerFunc 消息处理器顶层handler
type MessageHandlerFunc func(v interface{})

type Handler interface {
	HandlerMessage(v interface{})
}

type hook struct {
	init       InitFunc
	preConnect PreConnectFunc
	handler    Handler
}

type Live struct {
	Nws   *ws.NWebsocket // 底层 websocket
	done  chan bool      // done channel
	*hook                // hooks
}

type Options struct {
	WsOpts []ws.ConnectionOption // ws
}

type Option func(*Options)

func applyOpts(opts ...Option) *Options {
	result := &Options{}
	for _, opt := range opts {
		opt(result)
	}
	return result
}

func WithWsOptions(opts ...ws.ConnectionOption) Option {
	return func(options *Options) {
		options.WsOpts = opts
	}
}

func NewLive(opts ...Option) *Live {
	opt := applyOpts(opts...)
	return &Live{
		Nws:  ws.NewWsConnection(opt.WsOpts...),
		done: make(chan bool),
		hook: &hook{},
	}
}

func (l *Live) PreConnect(f func() (url string, err error)) {
	l.preConnect = f
}

func (l *Live) Init(f func(wsConn *ws.NWebsocket) (err error)) {
	l.init = f
}

func (l *Live) Handler(h Handler) {
	l.handler = h
}

var ErrPreConnectNotFound = errors.New("PreConnectFunc not found")

func (l *Live) Serve() error {
	l.Nws.OnConnected(func() {
		logger.SLog.Info("连接成功")
		if l.init != nil {
			if err := l.init(l.Nws); err != nil {
				l.Nws.Close()
			}
		}
	})
	l.Nws.OnDisconnected(func(err error) {
		logger.SLog.Info("连接关闭", err)
	})
	l.Nws.OnClose(func(code int, text string) {
		l.done <- true
	})
	l.Nws.OnReceiveError(func(err error) {
		logger.SLog.Error("接收数据错误", err)
	})
	l.Nws.OnConnectError(func(err error) {
		logger.SLog.Error("连接失败", err)
	})
	l.Nws.OnBinaryMessageReceived(func(v interface{}) {
		if l.handler != nil {
			l.handler.HandlerMessage(v)
		}
	})

	if l.preConnect == nil {
		return ErrPreConnectNotFound
	}

	url, err := l.preConnect()
	if err != nil {
		return err
	}
	go l.Nws.Connect(url)

	for {
		select {
		case <-l.done:
			return nil
		}
	}
}