/
wsinterface.go
80 lines (71 loc) · 1.9 KB
/
wsinterface.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
package notificationhandler
import (
"fmt"
"net/http"
"strings"
"sync"
"github.com/gorilla/websocket"
"github.com/kubescape/go-logger"
)
// IWebsocketActions -
type IWebsocketActions interface {
ReadMessage() (int, []byte, error)
Close() error
WritePingMessage() error
DefaultDialer(requestHeader http.Header) (*http.Response, error)
}
// WebsocketActions -
type WebsocketActions struct {
host string
conn *websocket.Conn
mutex *sync.Mutex
}
// NewWebsocketActions -
func NewWebsocketActions(host string) *WebsocketActions {
return &WebsocketActions{
host: host,
mutex: &sync.Mutex{},
}
}
// Close -
func (wa *WebsocketActions) Close() error {
wa.mutex.Lock()
defer wa.mutex.Unlock()
return wa.conn.Close()
}
// ReadMessage -
func (wa *WebsocketActions) ReadMessage() (int, []byte, error) {
messageType, p, err := wa.conn.ReadMessage()
return messageType, p, err
}
// DefaultDialer -
func (wa *WebsocketActions) DefaultDialer(requestHeader http.Header) (*http.Response, error) {
logger.L().Info("Connecting websocket to " + wa.host)
wa.mutex.Lock()
defer wa.mutex.Unlock()
conn, res, err := websocket.DefaultDialer.Dial(wa.host, nil)
if err != nil {
if strings.Contains(err.Error(), "bad handshake") {
if strings.HasPrefix(wa.host, "ws://") {
wa.host = strings.Replace(wa.host, "ws://", "wss://", 1)
} else if strings.HasPrefix(wa.host, "wss://") {
wa.host = strings.Replace(wa.host, "wss://", "ws://", 1)
}
conn, res, err = websocket.DefaultDialer.Dial(wa.host, nil)
}
}
if err == nil {
wa.conn = conn
logger.L().Info("Successfully connected websocket to " + wa.host)
} else {
err = fmt.Errorf("failed dialing to: '%s', reason: '%s'", wa.host, err.Error())
}
return res, err
}
// WritePingMessage -
func (wa *WebsocketActions) WritePingMessage() error {
wa.mutex.Lock()
defer wa.mutex.Unlock()
err := wa.conn.WriteMessage(websocket.PingMessage, []byte{})
return err
}