package playwright import ( "encoding/base64" "errors" ) type webSocketImpl struct { channelOwner isClosed bool page *pageImpl } func (ws *webSocketImpl) URL() string { return ws.initializer["url"].(string) } func newWebsocket(parent *channelOwner, objectType string, guid string, initializer map[string]interface{}) *webSocketImpl { ws := &webSocketImpl{} ws.createChannelOwner(ws, parent, objectType, guid, initializer) ws.page = fromChannel(parent.channel).(*pageImpl) ws.channel.On("close", func() { ws.Lock() ws.isClosed = true ws.Unlock() ws.Emit("close", ws) }) ws.channel.On( "frameSent", func(params map[string]interface{}) { ws.onFrameSent(params["opcode"].(float64), params["data"].(string)) }, ) ws.channel.On( "frameReceived", func(params map[string]interface{}) { ws.onFrameReceived(params["opcode"].(float64), params["data"].(string)) }, ) ws.channel.On( "socketError", func(params map[string]interface{}) { ws.Emit("socketerror", params["error"]) }, ) return ws } func (ws *webSocketImpl) onFrameSent(opcode float64, data string) { if opcode == 2 { payload, err := base64.StdEncoding.DecodeString(data) if err != nil { logger.Error("could not decode WebSocket.onFrameSent payload", "error", err) return } ws.Emit("framesent", payload) } else { ws.Emit("framesent", []byte(data)) } } func (ws *webSocketImpl) onFrameReceived(opcode float64, data string) { if opcode == 2 { payload, err := base64.StdEncoding.DecodeString(data) if err != nil { logger.Error("could not decode WebSocket.onFrameReceived payload", "error", err) return } ws.Emit("framereceived", payload) } else { ws.Emit("framereceived", []byte(data)) } } func (ws *webSocketImpl) ExpectEvent(event string, cb func() error, options ...WebSocketExpectEventOptions) (interface{}, error) { return ws.expectEvent(event, cb, options...) } func (ws *webSocketImpl) WaitForEvent(event string, options ...WebSocketWaitForEventOptions) (interface{}, error) { if len(options) == 1 { option := WebSocketExpectEventOptions(options[0]) return ws.expectEvent(event, nil, option) } else { return ws.expectEvent(event, nil) } } func (ws *webSocketImpl) expectEvent(event string, cb func() error, options ...WebSocketExpectEventOptions) (interface{}, error) { var predicate interface{} = nil timeout := ws.page.timeoutSettings.Timeout() if len(options) == 1 { if options[0].Timeout != nil { timeout = *options[0].Timeout } if options[0].Predicate != nil { predicate = options[0].Predicate } } waiter := newWaiter().WithTimeout(timeout) if event != "close" { waiter.RejectOnEvent(ws, "close", errors.New("websocket closed")) } if event != "socketerror" { waiter.RejectOnEvent(ws, "socketerror", errors.New("websocket error")) } waiter.RejectOnEvent(ws.page, "close", errors.New("page closed")) if cb == nil { return waiter.WaitForEvent(ws, event, predicate).Wait() } else { return waiter.WaitForEvent(ws, event, predicate).RunAndWait(cb) } } func (ws *webSocketImpl) IsClosed() bool { ws.RLock() defer ws.RUnlock() return ws.isClosed } func (ws *webSocketImpl) OnClose(fn func(WebSocket)) { ws.On("close", fn) } func (ws *webSocketImpl) OnFrameReceived(fn func(payload []byte)) { ws.On("framereceived", fn) } func (ws *webSocketImpl) OnFrameSent(fn func(payload []byte)) { ws.On("framesent", fn) } func (ws *webSocketImpl) OnSocketError(fn func(string)) { ws.On("socketerror", fn) }