Example #1
1
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
}
Example #2
1
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)
	}
}
Example #3
0
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
}
Example #4
0
// 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()
}
Example #5
0
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)
}
Example #6
0
// 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)
	}
}
Example #7
0
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)
}
Example #8
0
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)
}
Example #9
0
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
		}
	}
}
Example #10
0
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
}
Example #11
0
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)
	}
}
Example #12
0
// 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.
		}
	}
}
Example #13
0
// 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{})
	}
}
Example #14
0
// 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
			}
		}()
	}
}
Example #15
0
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)
}
Example #16
0
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(&currentPM25),
			"pm10":  atomic.LoadInt32(&currentPM100),
		}
		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)
	}
}
Example #17
0
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()
	}
}
Example #19
0
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
}
Example #20
0
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),
			)
		}
	}
}
Example #21
0
func done(conn *websocket.Conn, cc int, msg string) {
	conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(cc, msg), time.Time{})
}
Example #22
0
// 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
}
Example #23
0
func closeWS(ws *websocket.Conn, code int, reason string) {
	ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(code, reason), time.Now().Add(time.Second))
	ws.Close()
}
Example #24
0
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)
}
Example #25
0
func (h *FrontendHandler) closeConnection(ws *websocket.Conn) {
	ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
	ws.Close()
}