service/ws.go

package service

import (
	"encoding/json"

	"apigo.cc/gojs"
	"apigo.cc/gojs/goja"
	"github.com/gorilla/websocket"
	"github.com/ssgo/discover"
	"github.com/ssgo/log"
	"github.com/ssgo/s"
	"github.com/ssgo/u"
)

func MakeWSClient(client *websocket.Conn, id string) gojs.Map {
	return gojs.Map{
		"id": id,
		"read": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			typ, data, err := readWSMessage(client)
			if err != nil {
				panic(vm.NewGoError(err))
			}
			return vm.ToValue(gojs.Map{
				"type": typ,
				"data": data,
			})
		},
		"write": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := gojs.MakeArgs(&argsIn, vm).Check(1)
			var err error
			switch args.Any(0).(type) {
			case []byte:
				err = client.WriteMessage(websocket.BinaryMessage, args.Bytes(0))
			case string:
				err = client.WriteMessage(websocket.TextMessage, args.Bytes(0))
			default:
				err = client.WriteMessage(websocket.TextMessage, u.JsonBytes(args.Any(0)))
			}
			if err != nil {
				panic(vm.NewGoError(err))
			}
			return nil
		},
		"ping": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			err := client.WriteMessage(websocket.PingMessage, nil)
			if err != nil {
				panic(vm.NewGoError(err))
			}
			return nil
		},
		"writeMessage": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			args := gojs.MakeArgs(&argsIn, vm).Check(2)
			var err error
			switch args.Str(0) {
			case "text":
				err = client.WriteMessage(websocket.TextMessage, args.Bytes(1))
			case "binary":
				err = client.WriteMessage(websocket.BinaryMessage, args.Bytes(1))
			case "ping":
				err = client.WriteMessage(websocket.PingMessage, args.Bytes(1))
			case "pong":
				err = client.WriteMessage(websocket.PongMessage, args.Bytes(1))
			case "close":
				err = client.WriteMessage(websocket.CloseMessage, args.Bytes(1))
			case "json":
				err = client.WriteMessage(websocket.TextMessage, u.JsonBytes(args.Any(1)))
			default:
				err = client.WriteMessage(websocket.TextMessage, args.Bytes(1))
			}
			if err != nil {
				panic(vm.NewGoError(err))
			}
			return nil
		},
		"close": func(argsIn goja.FunctionCall, vm *goja.Runtime) goja.Value {
			err := client.Close()
			if err != nil {
				panic(vm.NewGoError(err))
			}
			return nil
		},
	}
}

func readWSMessage(client *websocket.Conn) (string, any, error) {
	msgType := "close"
	var msgData any
	typ, data, err := client.ReadMessage()
	if err == nil && typ != websocket.CloseMessage {
		switch typ {
		case websocket.TextMessage:
			if err2 := json.Unmarshal(data, &msgData); err2 != nil {
				msgData = string(data)
				msgType = "text"
			} else {
				msgType = "json"
			}
		case websocket.BinaryMessage:
			if err2 := json.Unmarshal(data, &msgData); err2 != nil {
				msgData = data
				msgType = "binary"
			} else {
				msgType = "json"
			}
		case websocket.PingMessage:
			msgData = string(data)
			msgType = "ping"
		case websocket.PongMessage:
			msgData = string(data)
			msgType = "pong"
		}
	}
	return msgType, msgData, err
}

func makeWSAction(startFile string, actionKey string) any {
	// onOpen
	return func(args map[string]any, headers map[string]string, request *s.Request, client *websocket.Conn, caller *discover.Caller, session *Session, logger *log.Logger) any {
		if pool := getPool(startFile); pool != nil {
			// 调用 onOpen
			hasOnMessage := false
			hasOnClose := false
			func() {
				rt := pool.Get()
				defer pool.Put(rt)
				if action, ok := rt.GetGoData(actionKey).(goja.Callable); ok {
					var thisArg goja.Value
					if thisArgV, ok := rt.GetGoData(actionKey + "This").(goja.Value); ok {
						thisArg = thisArgV
					}
					_, _ = rt.RunVM(func(vm *goja.Runtime) (any, error) {
						r2, err2 := runAction(action, vm, thisArg, args, headers, request, nil, client, caller, session, logger)
						return r2, err2

					})
					_, hasOnMessage = rt.GetGoData(actionKey + "onMessage").(goja.Callable)
					_, hasOnClose = rt.GetGoData(actionKey + "onClose").(goja.Callable)
				}
			}()
			// 循环读取消息，回调js的onMessage
			if hasOnMessage {
				for {
					typ, data, err := readWSMessage(client)
					if err != nil || typ == "close" {
						break
					}
					func() {
						rt := pool.Get()
						defer pool.Put(rt)
						if action, ok := rt.GetGoData(actionKey + "onMessage").(goja.Callable); ok {
							var thisArg goja.Value
							if thisArgV, ok := rt.GetGoData(actionKey + "This").(goja.Value); ok {
								thisArg = thisArgV
							}
							_, _ = rt.RunVM(func(vm *goja.Runtime) (any, error) {
								requestParams, _ := makeRequestParams(args, headers, request, nil, client, caller, session, logger)
								requestParams["type"] = typ
								requestParams["data"] = data
								_, err := action(thisArg, vm.ToValue(requestParams))
								if err != nil {
									logger.Error(err.Error())
								}
								return nil, nil
							})
						}
					}()
				}
			}

			if hasOnClose {
				func() {
					rt := pool.Get()
					defer pool.Put(rt)
					if action, ok := rt.GetGoData(actionKey + "onClose").(goja.Callable); ok {
						var thisArg goja.Value
						if thisArgV, ok := rt.GetGoData(actionKey + "This").(goja.Value); ok {
							thisArg = thisArgV
						}
						_, _ = rt.RunVM(func(vm *goja.Runtime) (any, error) {
							requestParams, _ := makeRequestParams(args, headers, request, nil, client, caller, session, logger)
							_, err := action(thisArg, vm.ToValue(requestParams))
							if err != nil {
								logger.Error(err.Error())
							}
							return nil, nil
						})
					}
				}()
			}
		}
		return nil
	}
}