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
}
// 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)
}
}
}