x

core/websocket.go

@@ -1,9 +1,16 @@
 package core
 
 import (
+	"encoding/json"
+	"fmt"
 	"net/http"
 	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
 
+	"github.com/buger/jsonparser"
+	"github.com/cdle/sillyplus/utils"
 	"github.com/dop251/goja"
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
@@ -19,6 +26,62 @@ func IsWebSocketRequest(req *http.Request) bool {
 	return true
 }
 
+type WsPattern struct {
+	Value map[string]interface{}
+	Chan  chan map[string]interface{}
+}
+
+type WsConn struct {
+	conn     *websocket.Conn
+	patterns sync.Map
+	Key      int64
+	sync.RWMutex
+}
+
+func (wc *WsConn) Close() error {
+	return wc.conn.Close()
+}
+
+func (wc *WsConn) WriteMessage(messageType int, data []byte, pattern map[string]interface{}) (error, map[string]interface{}) {
+	var res map[string]interface{}
+	wp := &WsPattern{}
+	var timeout int
+	if pattern != nil {
+		if v, ok := pattern["$timeout"]; ok {
+			timeout = utils.Int(v)
+			delete(pattern, "$timeout")
+		} //
+		wp.Value = pattern
+		key := atomic.AddInt64(&wc.Key, 1)
+		wp.Chan = make(chan map[string]interface{}, 1)
+		defer func() {
+			close(wp.Chan)
+			wc.patterns.Delete(key)
+		}()
+		wc.patterns.Store(key, wp)
+	}
+	var err error
+	func() {
+		wc.Lock()
+		defer wc.Unlock()
+		err = wc.conn.WriteMessage(messageType, data)
+	}()
+	if pattern != nil {
+		if timeout == 0 {
+			timeout = 5000
+		}
+		select {
+		case res = <-wp.Chan:
+		case <-time.After(time.Millisecond * time.Duration(timeout)):
+		}
+	}
+	return err, res
+}
+
+// func (wc *WsConn) WriteMessage(messageType int, data []byte) error {
+// 	return wc.conn.WriteMessage(messageType, data)
+// }
+
 func handleWebsocket(c *gin.Context) {
 	for _, function := range Functions {
 		if len(function.Https) != 0 {
@@ -34,14 +97,6 @@ func handleWebsocket(c *gin.Context) {
 					var res = &Response{
 						c: c,
 					}
-					req._event = "connect"
-					function.Handle(&Faker{
-						Type: "websocket",
-					}, func(vm *goja.Runtime) {
-						vm.Set("res", res)
-						vm.Set("req", req)
-					})
-					// message
 					var upGrader = websocket.Upgrader{
 						CheckOrigin: func(r *http.Request) bool {
 							return true
@@ -52,9 +107,63 @@ func handleWebsocket(c *gin.Context) {
 						c.Writer.Write([]byte(err.Error()))
 						return
 					}
-					res.conn = ws
+					wc := &WsConn{}
+					res.conn = wc
+					req._event = "connect"
+					wc.conn = ws
+					go function.Handle(&Faker{
+						Type: "websocket",
+					}, func(vm *goja.Runtime) {
+						vm.Set("res", res)
+						vm.Set("req", req)
+					})
+					// message
 					for {
 						_, data, err := ws.ReadMessage()
+						wc.patterns.Range(func(key, value any) bool {
+							wp := value.(*WsPattern)
+							matched := false
+							// fmt.Println("wp.Value", wp.Value)
+							for k, v := range wp.Value {
+								value, _, _, err := jsonparser.Get(data, strings.Split(k, ".")...)
+								// fmt.Println("k, v", k, v, "key path", strings.Split(k, "."), "data:", string(data), err)
+								if err != nil {
+									// fmt.Println("err1", err)
+									return true
+								}
+								if string(value) != fmt.Sprint(v) {
+									return true
+								}
+								matched = true
+								// v2 := v
+								// fmt.Println("v2 := v", v2, string(value))
+								// err = json.Unmarshal(value, &v2)
+								// if err != nil {
+								// 	fmt.Println("err2", err)
+								// 	return true
+								// }
+								// if v != v2 {
+								// 	fmt.Println("err3", err)
+								// 	return true
+								// } else {
+								// 	matched = true
+								// }
+							}
+							// fmt.Println("matched", matched)
+							if matched {
+								var result = map[string]interface{}{}
+								err := json.Unmarshal(data, &result)
+								if err == nil {
+									select {
+									case wp.Chan <- result:
+									default:
+									}
+								} else {
+									// fmt.Println("err3", err)
+								}
+							}
+							return true
+						})
 						if err != nil { // disconnect
 							req._event = "disconnect"
 							for _, f2 := range Functions {
@@ -78,7 +187,7 @@ func handleWebsocket(c *gin.Context) {
 								function = f2
 							}
 						}
-						function.Handle(&Faker{
+						go function.Handle(&Faker{
 							Type: "websocket",
 						}, func(vm *goja.Runtime) {
 							vm.Set("res", res)