func connectToAgent(timeout int, onReady func(*websocket.Conn)) (*websocket.Conn, error) {
var ws *websocket.Conn
var err error
t := 0
for {
if t > timeout {
return nil, errors.New("Connection to Agent : timeout reached")
}
ws, err = newClient()
if err == nil {
break
}
time.Sleep(1 * time.Second)
t++
}
ready := false
h := func(message string) error {
err := ws.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(time.Second))
if err != nil {
return err
}
if !ready {
ready = true
onReady(ws)
}
return nil
}
ws.SetPingHandler(h)
return ws, nil
}
func AddWSSink(receivedChan chan []byte, port string, path string) (*websocket.Conn, <-chan struct{}) {
connectionDroppedChannel := make(chan struct{}, 1)
var ws *websocket.Conn
ip, _ := localip.LocalIP()
fullURL := "ws://" + ip + ":" + port + path
Eventually(func() error {
var err error
ws, _, err = websocket.DefaultDialer.Dial(fullURL, http.Header{})
return err
}, 5, 1).ShouldNot(HaveOccurred(), fmt.Sprintf("Unable to connect to server at %s.", fullURL))
ws.SetPingHandler(func(message string) error {
ws.WriteControl(websocket.PongMessage, []byte(message), time.Time{})
return nil
})
go func() {
for {
_, data, err := ws.ReadMessage()
if err != nil {
close(connectionDroppedChannel)
close(receivedChan)
return
}
receivedChan <- data
}
}()
return ws, connectionDroppedChannel
}
func AddWSSink(receivedChan chan []byte, port string, path string) (*websocket.Conn, chan bool, <-chan bool) {
dontKeepAliveChan := make(chan bool, 1)
connectionDroppedChannel := make(chan bool, 1)
var ws *websocket.Conn
i := 0
for {
var err error
ws, _, err = websocket.DefaultDialer.Dial("ws://localhost:"+port+path, http.Header{})
if err != nil {
i++
if i > 10 {
fmt.Printf("Unable to connect to Server in 100ms, giving up.\n")
return nil, nil, nil
}
<-time.After(10 * time.Millisecond)
continue
} else {
break
}
}
ws.SetPingHandler(func(message string) error {
select {
case <-dontKeepAliveChan:
// do nothing
default:
ws.WriteControl(websocket.PongMessage, []byte(message), time.Time{})
}
return nil
})
go func() {
for {
_, data, err := ws.ReadMessage()
if err != nil {
close(connectionDroppedChannel)
close(receivedChan)
return
}
receivedChan <- data
}
}()
return ws, dontKeepAliveChan, connectionDroppedChannel
}
/* reader reads input from the web socket connection and sends it to the master.
* reader runs in its own goroutine.
* If the connection closes, the reader goroutine terminates.
*/
func reader(conn *websocket.Conn, ch chan readerResult) {
conn.SetPongHandler(func(appData string) error {
ch <- readerResult{conn, 10, nil, nil}
return nil
})
conn.SetPingHandler(func(appData string) error {
ch <- readerResult{conn, 9, nil, nil}
return nil
})
for {
messageType, data, err := conn.ReadMessage()
readerResult := readerResult{conn, messageType, data, err}
ch <- readerResult
if err != nil {
return
}
}
}
func wsSetPingHandler(t *WSTunnelServer, ws *websocket.Conn, rs *remoteServer) {
// timeout handler sends a close message, waits a few seconds, then kills the socket
timeout := func() {
ws.WriteControl(websocket.CloseMessage, nil, time.Now().Add(1*time.Second))
time.Sleep(5 * time.Second)
rs.log.Info("WS closing due to ping timeout", "ws", wsp(ws))
ws.Close()
}
// timeout timer
timer := time.AfterFunc(t.WSTimeout, timeout)
// ping handler resets last ping time
ph := func(message string) error {
timer.Reset(t.WSTimeout)
ws.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(t.WSTimeout/3))
// update lastActivity
rs.lastActivity = time.Now()
return nil
}
ws.SetPingHandler(ph)
}
// Reads incoming data from the websocket and forwards it to stdout.
func readWS(ws *websocket.Conn, query *regexp.Regexp, done chan struct{}) {
defer func() {
done <- struct{}{}
}()
ws.SetPingHandler(func(string) error {
ws.SetWriteDeadline(time.Now().Add(writeTimeout))
return ws.WriteMessage(websocket.PongMessage, []byte{})
})
for {
_, msg, err := ws.ReadMessage()
if err != nil {
return
}
var log LogMessage
err = json.Unmarshal(msg, &log)
if err == nil {
if query == nil || query.MatchString(log.Message) {
logrus.Printf("%s - %s", log.Timestamp, log.Message)
}
} else {
logrus.StandardLogger().Out.Write(msg)
}
}
}
func addWSSink(receivedChan chan []byte, url string) (chan struct{}, <-chan struct{}, func()) {
stopKeepAlive := make(chan struct{})
connectionDropped := make(chan struct{})
var ws *websocket.Conn
tryToConnect := func() error {
var err error
ws, _, err = websocket.DefaultDialer.Dial(url, http.Header{})
return err
}
Eventually(tryToConnect, 5).ShouldNot(HaveOccurred())
ws.SetPingHandler(func(message string) error {
select {
case <-stopKeepAlive:
return nil
default:
return ws.WriteControl(websocket.PongMessage, []byte(message), time.Time{})
}
})
go func() {
defer close(connectionDropped)
for {
_, data, err := ws.ReadMessage()
if err != nil {
return
}
receivedChan <- data
}
}()
return stopKeepAlive, connectionDropped, func() {
ws.Close()
}
}
keepAliveCompleted = make(chan struct{})
testServer = httptest.NewServer(makeTestHandler(keepAliveCompleted))
var err error
wsClient, _, err = websocket.DefaultDialer.Dial(httpToWs(testServer.URL), nil)
Expect(err).NotTo(HaveOccurred())
})
AfterEach(func() {
wsClient.Close()
testServer.Close()
})
It("sends pings to the client", func() {
var pingCount int32
wsClient.SetPingHandler(func(string) error {
atomic.AddInt32(&pingCount, 1)
return nil
})
go wsClient.ReadMessage()
Eventually(func() int32 { return atomic.LoadInt32(&pingCount) }).ShouldNot(BeZero())
})
It("doesn't close the client when it responds with pong frames", func() {
go wsClient.ReadMessage()
// default ping handler responds with pong frames
Consistently(keepAliveCompleted).ShouldNot(BeClosed())
})
It("closes the client when it doesn't respond with pong frames", func() {
wsClient.SetPingHandler(func(string) error { return nil })
go wsClient.ReadMessage()
func CmdLoop() {
const (
// Time allowed to write a message to the peer.
writeWait = 10 * time.Second
// Time allowed to resond to a ping event from peer
pingWait = 55 * time.Second
)
var (
// Connection object
ws *websocket.Conn
)
bigobjectURL := url.URL{
Scheme: "ws",
Host: net.JoinHostPort(host, port),
Path: "exec",
}
var err error
dialer := websocket.DefaultDialer
if ws, _, err = dialer.Dial(bigobjectURL.String(), nil); err != nil {
if !quiet {
log.Fatalln(err)
} else {
os.Exit(1)
}
}
ws.SetPingHandler(nil)
input, next := MakeReadline()
ticker := time.Tick(pingWait)
defer func() {
if err := recover(); err != nil {
ws.Close()
if !quiet {
log.Fatalln(err)
} else {
os.Exit(1)
}
}
}()
for {
select {
case stmt, ok := <-input:
if !ok {
return
}
now := time.Now()
switch parseAction(stmt) {
default:
req := &api.RPCRequest{Stmt: stmt}
if err := ws.WriteJSON(req); err != nil {
if !quiet {
fmt.Println(err)
}
} else {
var resp interface{}
if err := ws.ReadJSON(&resp); err != nil {
panic(err)
}
text, _ := json.MarshalIndent(resp, "", " ")
fmt.Println(string(text))
}
break
case GET_RESULT_NOW:
req := &api.RPCRequest{Stmt: stmt, Opts: &api.RPCOpts{Handle: handle}}
if err := ws.WriteJSON(req); err != nil {
if !quiet {
fmt.Println(err)
}
} else if !handle {
var idx int64 = 1
for idx > 0 {
var resp map[string]interface{}
if err := ws.ReadJSON(&resp); err != nil {
panic(err)
}
if thing, ok := resp["Content"]; ok && thing != nil {
payload := thing.(map[string]interface{})
data := payload["content"].([]interface{})
idx = int64(payload["index"].(float64))
for _, thing := range data {
text, _ := json.MarshalIndent(thing, "", " ")
fmt.Println(string(text))
}
} else {
text, _ := json.MarshalIndent(resp, "", " ")
fmt.Println(string(text))
break
}
}
} else {
var resp interface{}
if err := ws.ReadJSON(&resp); err != nil {
panic(err)
}
text, _ := json.MarshalIndent(resp, "", " ")
fmt.Println(string(text))
}
break
case FETCH_DATA:
req := &api.RPCRequest{Stmt: stmt}
if err := ws.WriteJSON(req); err != nil {
if !quiet {
fmt.Println(err)
}
} else {
var idx int64 = 1
for idx > 0 {
var resp map[string]interface{}
if err := ws.ReadJSON(&resp); err != nil {
panic(err)
}
if thing, ok := resp["Content"]; ok && thing != nil {
payload := thing.(map[string]interface{})
data := payload["content"].([]interface{})
idx = int64(payload["index"].(float64))
for _, thing := range data {
text, _ := json.MarshalIndent(thing, "", " ")
fmt.Println(string(text))
}
} else {
text, _ := json.MarshalIndent(resp, "", " ")
fmt.Println(string(text))
break
}
}
}
}
next <- time.Since(now)
case <-ticker:
ws.WriteControl(
websocket.PongMessage,
[]byte{},
time.Now().Add(writeWait),
)
}
}
}
