Ejemplo n.º 1
0
func (c *appClient) controlConnection() (bool, error) {
	headers := http.Header{}
	c.ProbeConfig.authorizeHeaders(headers)
	url := sanitize.URL("ws://", 0, "/api/control/ws")(c.target)
	conn, _, err := c.wsDialer.Dial(url, headers)
	if err != nil {
		return false, err
	}
	defer func() {
		conn.Close()
	}()

	codec := xfer.NewJSONWebsocketCodec(conn)
	server := rpc.NewServer()
	if err := server.RegisterName("control", c.control); err != nil {
		return false, err
	}

	// Will return false if we are exiting
	if !c.registerConn("control", conn) {
		return true, nil
	}
	defer c.closeConn("control")

	server.ServeCodec(codec)
	return false, nil
}
Ejemplo n.º 2
0
// handleProbeWS accepts websocket connections from the probe and registers
// them in the control router, such that HandleControl calls can find them.
func (cr *controlRouter) handleProbeWS(w http.ResponseWriter, r *http.Request) {
	probeID := r.Header.Get(xfer.ScopeProbeIDHeader)
	if probeID == "" {
		respondWith(w, http.StatusBadRequest, xfer.ScopeProbeIDHeader)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Error upgrading to websocket: %v", err)
		return
	}
	defer conn.Close()

	codec := xfer.NewJSONWebsocketCodec(conn)
	client := rpc.NewClientWithCodec(codec)
	handler := controlHandler{
		id:     rand.Int63(),
		client: client,
	}

	cr.set(probeID, handler)

	codec.WaitForReadError()

	cr.rm(probeID, handler)
	client.Close()
}