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 closeWithErr(log lager.Logger, conn *websocket.Conn, code int, reason string) {
err := conn.WriteControl(
websocket.CloseMessage,
websocket.FormatCloseMessage(code, reason),
time.Time{},
)
if err != nil {
log.Error("failed-to-close-websocket-connection", err)
}
}
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
}
// Pick requests off the RemoteServer queue and send them into the tunnel
func wsWriter(rs *remoteServer, ws *websocket.Conn, ch chan int) {
var req *remoteRequest
var err error
for {
// fetch a request
select {
case req = <-rs.requestQueue:
// awesome...
case _ = <-ch:
// time to close shop
rs.log.Info("WS closing on signal", "ws", wsp(ws))
ws.Close()
return
}
//log.Printf("WS->%s#%d start %s", req.token, req.id, req.info)
// See whether the request has already expired
if req.deadline.Before(time.Now()) {
req.replyChan <- responseBuffer{
err: errors.New("Timeout before forwarding the request"),
}
req.log.Info("WS SND timeout before sending", "ago",
time.Now().Sub(req.deadline).Seconds())
continue
}
// write the request into the tunnel
ws.SetWriteDeadline(time.Now().Add(time.Minute))
var w io.WriteCloser
w, err = ws.NextWriter(websocket.BinaryMessage)
// got an error, reply with a "hey, retry" to the request handler
if err != nil {
break
}
// write the request Id
_, err = fmt.Fprintf(w, "%04x", req.id)
if err != nil {
break
}
// write the request itself
_, err = req.buffer.WriteTo(w)
if err != nil {
break
}
// done
err = w.Close()
if err != nil {
break
}
req.log.Info("WS SND", "info", req.info)
}
// tell the sender to retry the request
req.replyChan <- responseBuffer{err: RetryError}
req.log.Info("WS error causes retry")
// close up shop
ws.WriteControl(websocket.CloseMessage, nil, time.Now().Add(5*time.Second))
time.Sleep(2 * time.Second)
ws.Close()
}
func (m *multiplexer) routeMessages(ws *websocket.Conn) {
stopSignal := make(chan bool, 1)
// Read messages from backend
go func(stop chan<- bool) {
for {
msgType, msg, err := ws.ReadMessage()
if err != nil {
m.shutdown(stop)
return
}
if msgType != websocket.TextMessage {
continue
}
message := common.ParseMessage(string(msg))
m.frontendMu.RLock()
frontendChan, ok := m.frontendChans[message.Key]
if ok {
frontendChan <- message
}
m.frontendMu.RUnlock()
if !ok && message.Type != common.Close {
m.sendClose(message.Key)
}
}
}(stopSignal)
// Write messages to backend
go func(stop <-chan bool) {
ticker := time.NewTicker(time.Second * 5)
defer ticker.Stop()
for {
select {
case message, ok := <-m.messagesToBackend:
if !ok {
return
}
err := ws.WriteMessage(websocket.TextMessage, []byte(message))
if err != nil {
log.WithFields(log.Fields{"error": err, "msg": message}).Error("Could not write message.")
}
case <-ticker.C:
ws.WriteControl(websocket.PingMessage, []byte(""), time.Now().Add(time.Second))
case <-stop:
return
}
}
}(stopSignal)
}
// ping will loop infinitely and send ping messages over ws. If a write takes
// longer than writeWait, it will remove ws from the connection pool.
func ping(ws *websocket.Conn, id id) {
b := make([]byte, 2)
for range time.Tick(pingPeriod) {
rand.Read(b)
deadline := time.Now().Add(writeWait)
if err := ws.WriteControl(websocket.PingMessage, b, deadline); err != nil {
die(id, err)
return
}
log.Debugf(appengine.BackgroundContext(), "ws: ping %x for %s(%d %q) sent", b, id.kind, id.intID, id.stringID)
}
}
func (w *WebsocketServer) streamLogs(appId string, ws *websocket.Conn) {
defer ws.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(100*time.Millisecond))
websocketSink := sinks.NewWebsocketSink(
appId,
w.logger,
ws,
w.keepAliveInterval,
w.bufferSize,
)
w.logger.Debugf("WebsocketServer: Requesting a wss sink for app %s", websocketSink.AppId())
w.sinkManager.RegisterSink(websocketSink, false)
}
func stopSend(con *websocket.Conn, done chan int) {
//send close message
fmt.Println("wsSend: got stop mssage")
//fmt.Println("wsSend: closing send")
//its okay if we continue recieving a bit before the close message is read.
//channel was closed, send close frame and exit
fmt.Println("wsSend: sending close frame")
err := con.WriteControl(websocket.CloseMessage, nil, time.Now().Add(3*time.Second))
if err != nil { //if theres an error sending, print err, exit
fmt.Println("wsSend: control frame send err:", err)
}
close(done)
}
func ping(ws *websocket.Conn, done chan struct{}) {
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
for {
select {
case <-ticker.C:
if err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
log.Println("ping:", err)
}
case <-done:
return
}
}
}
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
}
func pingDispatcher(s *websocket.Conn) {
for {
if s == nil {
log.Info("can't ping as connection closed")
return
}
if err := s.WriteControl(websocket.PingMessage, []byte(fmt.Sprintf("server at %s", time.Now().String())), time.Now().Add(WS_TIMEOUT)); err == nil {
log.Tracef("ping sent")
} else {
log.Error("ping failed")
return
}
time.Sleep(WS_TIMEOUT)
}
}
// pinger simulates the websocket to keep it alive with ping messages.
func pinger(conn *websocket.Conn, done chan struct{}) {
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
for { // blocking loop with select to wait for stimulation.
select {
case <-ticker.C:
if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, err.Error()), time.Time{})
return
}
case <-done:
return // clean up this routine.
}
}
}
// pumpStdout handles reading data from stdout of the process and writing
// it to websocket connection.
func pumpStdout(conn *websocket.Conn, stdout io.Reader, done chan struct{}) {
go pinger(conn, done)
defer func() {
conn.Close()
close(done) // make sure to close the pinger when we are done.
}()
s := bufio.NewScanner(stdout)
for s.Scan() {
conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := conn.WriteMessage(websocket.TextMessage, bytes.TrimSpace(s.Bytes())); err != nil {
break
}
}
if s.Err() != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, s.Err().Error()), time.Time{})
}
}
// reader is the guts of this package. It takes the stdin and stdout pipes
// of the cmd we created in ServeWS and pipes them between the client and server
// over websockets.
func reader(conn *websocket.Conn, stdout io.ReadCloser, stdin io.WriteCloser) {
// Setup our connection's websocket ping/pong handlers from our const values.
conn.SetReadLimit(maxMessageSize)
conn.SetReadDeadline(time.Now().Add(pongWait))
conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
tickerChan := make(chan bool)
defer close(tickerChan) // make sure to close the ticker when we are done.
go ticker(conn, tickerChan)
for {
msgType, r, err := conn.NextReader()
if err != nil {
if msgType == -1 {
return // we got a disconnect from the client. We are good to close.
}
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
w, err := conn.NextWriter(msgType)
if err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
if _, err := io.Copy(stdin, r); err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
go func() {
if _, err := io.Copy(w, stdout); err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
if err := w.Close(); err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
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)
}
func WSHandler(c *websocket.Conn) {
closed := false
ret := websocket.CloseError{
Code: websocket.CloseNormalClosure,
}
const writeWait = 10 * time.Second
defer func() {
log.Infof("Close websocket [%s]", ret.Error())
c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(ret.Code, ret.Text), time.Now().Add(writeWait))
c.Close()
}()
go func() {
for {
_, r, err := c.NextReader()
if err != nil {
break
}
io.Copy(ioutil.Discard, r)
}
log.Infof("Pipe closed")
closed = true
}()
for !closed {
signal := M{
"pm2.5": atomic.LoadInt32(¤tPM25),
"pm10": atomic.LoadInt32(¤tPM100),
}
if err := c.WriteJSON(signal); err != nil && err != websocket.ErrCloseSent {
if _, ok := err.(net.Error); !ok {
log.Warnf("Error writing to client: %s", err.Error())
}
closed = true
}
time.Sleep(2 * time.Second)
}
}
func (h *Hijacker) handleInput(conn *websocket.Conn, inputs <-chan atc.HijackInput, finished chan struct{}) {
ticker := time.NewTicker(h.interval)
defer ticker.Stop()
for {
select {
case input := <-inputs:
err := conn.WriteJSON(input)
if err != nil {
fmt.Fprintf(os.Stderr, "failed to send input: %s", err.Error())
return
}
case t := <-ticker.C:
err := conn.WriteControl(websocket.PingMessage, []byte(t.String()), time.Now().Add(time.Second))
if err != nil {
fmt.Fprintf(os.Stderr, "failed to send heartbeat: %s", err.Error())
}
case <-finished:
return
}
}
}
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()
}
}
func closeWebsocket(closeCode int, ws *websocket.Conn) bool {
closeMessage := websocket.FormatCloseMessage(closeCode, "")
err := ws.WriteControl(websocket.CloseMessage, closeMessage, time.Now().Add(1*time.Second))
return err == nil
}
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),
)
}
}
}
func done(conn *websocket.Conn, cc int, msg string) {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(cc, msg), time.Time{})
}
// closeWebsocketNormally is a shortcut for sending a 'close' control message
// with 'normal closure' and timeout given in app config
func closeWebsocketNormally(conn *websocket.Conn, msg string) (err error) {
err = conn.WriteControl(websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, msg),
time.Now().Add(time.Duration(appVars.config.UploadMaxIdleDurationS)*time.Second))
return
}
func closeWS(ws *websocket.Conn, code int, reason string) {
ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(code, reason), time.Now().Add(time.Second))
ws.Close()
}
func (w *WebsocketServer) recentLogs(appId string, ws *websocket.Conn) {
defer ws.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(100*time.Millisecond))
logMessages := w.sinkManager.RecentLogsFor(appId)
sendMessagesToWebsocket(logMessages, ws, w.logger)
}
func (h *FrontendHandler) closeConnection(ws *websocket.Conn) {
ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
ws.Close()
}
