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