Example #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 := xfer.DialWS(&c.wsDialer, url, headers)
	if err != nil {
		return false, err
	}
	defer conn.Close()

	doControl := func(req xfer.Request) xfer.Response {
		req.AppID = c.appID
		var res xfer.Response
		c.control.Handle(req, &res)
		return res
	}

	codec := xfer.NewJSONWebsocketCodec(conn)
	server := rpc.NewServer()
	if err := server.RegisterName("control", xfer.ControlHandlerFunc(doControl)); 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
}
Example #2
0
// handleProbeWS accepts websocket connections from the probe and registers
// them in the control router, such that HandleControl calls can find them.
func handleProbeWS(cr ControlRouter) CtxHandlerFunc {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		probeID := r.Header.Get(xfer.ScopeProbeIDHeader)
		if probeID == "" {
			respondWith(w, http.StatusBadRequest, xfer.ScopeProbeIDHeader)
			return
		}

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

		codec := xfer.NewJSONWebsocketCodec(conn)
		client := rpc.NewClientWithCodec(codec)
		defer client.Close()

		id, err := cr.Register(ctx, probeID, func(req xfer.Request) xfer.Response {
			var res xfer.Response
			if err := client.Call("control.Handle", req, &res); err != nil {
				return xfer.ResponseError(err)
			}
			return res
		})
		if err != nil {
			respondWith(w, http.StatusBadRequest, err.Error())
			return
		}
		defer cr.Deregister(ctx, probeID, id)
		if err := codec.WaitForReadError(); err != nil && err != io.EOF && !xfer.IsExpectedWSCloseError(err) {
			log.Errorf("Error on websocket: %v", err)
		}
	}
}