[#72940] linux-client: Extract websocket methods into separate package

This extracts the RDFM-specific WebSocket connection into a separate
package. This is in preparation for re-using existing RDFM WS connection
code for shell output WS.

The existing WebSocket connection management struct was renamed to
DeviceManagementConnection to avoid confusion between the two.

Signed-off-by: Łukasz Kędziora <lkedziora@antmicro.com>

Author: lkedziora

devices/linux-client/daemon/device.go

@@ -26,11 +26,10 @@ import (
 	"github.com/antmicro/rdfm/actions"
 	"github.com/antmicro/rdfm/app"
 	"github.com/antmicro/rdfm/conf"
+	"github.com/antmicro/rdfm/serverws"
 	"github.com/antmicro/rdfm/telemetry"
-	"github.com/gorilla/websocket"
 
 	netUtils "github.com/antmicro/rdfm/daemon/net_utils"
-	requests "github.com/antmicro/rdfm/daemon/requests"
 )
 
 const RSA_DEVICE_KEY_SIZE = 4096
@@ -47,11 +46,11 @@ type Device struct {
 	tokenMutex          sync.Mutex
 	httpTransport       *http.Transport
 	logManager          *telemetry.LogManager
-	conn                *DeviceConnection
+	conn                *serverws.DeviceManagementConnection
 	actionRunner        *actions.ActionRunner
 }
 
-func (d *Device) handleRequest(msg []byte) (requests.Request, error) {
+func (d *Device) handleRequest(msg []byte) (serverws.Request, error) {
 	var msgMap map[string]interface{}
 
 	err := json.Unmarshal(msg, &msgMap)
@@ -63,18 +62,18 @@ func (d *Device) handleRequest(msg []byte) (requests.Request, error) {
 
 	log.Infof("Handling '%s' request...", requestName)
 
-	request, err := requests.Parse(string(msg[:]))
+	request, err := serverws.Parse(string(msg[:]))
 	if err != nil {
 		return nil, err
 	}
 
 	switch r := request.(type) {
-	case requests.Alert:
+	case serverws.Alert:
 		for key, val := range r.Alert {
 			log.Printf("Server sent %s: %s", key, val)
 		}
-	case requests.ActionExec:
-		response := requests.ActionExecControl{
+	case serverws.ActionExec:
+		response := serverws.ActionExecControl{
 			Method:      "action_exec_control",
 			ExecutionId: r.ExecutionId,
 			Status:      "ok",
@@ -89,11 +88,11 @@ func (d *Device) handleRequest(msg []byte) (requests.Request, error) {
 		}
 
 		return response, nil
-	case requests.ActionListQuery:
+	case serverws.ActionListQuery:
 		actions := d.actionRunner.List()
-		var reqActions []requests.Action
+		var reqActions []serverws.Action
 		for _, action := range actions {
-			reqAction := requests.Action{
+			reqAction := serverws.Action{
 				ActionId:    action.Id,
 				ActionName:  action.Name,
 				Description: action.Description,
@@ -103,22 +102,22 @@ func (d *Device) handleRequest(msg []byte) (requests.Request, error) {
 
 			reqActions = append(reqActions, reqAction)
 		}
-		response := requests.ActionListUpdate{
+		response := serverws.ActionListUpdate{
 			Method:  "action_list_update",
 			Actions: reqActions,
 		}
 		return response, nil
-	//case requests.DeviceAttachToManager:
+	//case serverws.DeviceAttachToManager:
 	// TODO: Handle shell_attach
 	default:
 		log.Warnf("Request '%s' is unsupported", requestName)
-		response := requests.CantHandleRequest()
+		response := serverws.CantHandleRequest()
 		return response, nil
 	}
 	return nil, nil
 }
 
-func (d *Device) marshalSendRetry(req requests.Request, cancelCtx context.Context) error {
+func (d *Device) marshalSendRetry(req serverws.Request, cancelCtx context.Context) error {
 	msg, err := json.Marshal(req)
 	if err != nil {
 		return err
@@ -175,14 +174,6 @@ func (d *Device) prepareHttpTransport(tlsConf *tls.Config) *http.Transport {
 	}
 }
 
-func (d *Device) prepareWsDialer(tlsConf *tls.Config) *websocket.Dialer {
-	if tlsConf != nil {
-		return &websocket.Dialer{TLSClientConfig: tlsConf}
-	} else {
-		return websocket.DefaultDialer
-	}
-}
-
 func (d *Device) setupConnection() error {
 	// Get MAC address
 	mac, err := netUtils.GetMacAddr()
@@ -210,8 +201,7 @@ func (d *Device) setupConnection() error {
 	}
 
 	d.httpTransport = d.prepareHttpTransport(tlsConf)
-	wsDialer := d.prepareWsDialer(tlsConf)
-	d.conn = NewDeviceConnection(serverUrl, *wsDialer, 1024)
+	d.conn = serverws.NewDeviceConnection(serverUrl, tlsConf, 1024)
 	// TODO: Hardcode capabilities right now. This should be read from the
 	// config file instead, but we don't have any configuration options that
 	// determine whether action/shell should be enabled or not.
@@ -230,7 +220,7 @@ func (d *Device) setupActionRunner() error {
 }
 
 func (d *Device) actionResultCallback(result actions.ActionResult, cancelCtx context.Context) bool {
-	res := requests.ActionExecResult{
+	res := serverws.ActionExecResult{
 		Method:      "action_exec_result",
 		ExecutionId: result.ExecId,
 		StatusCode:  result.StatusCode,

devices/linux-client/daemon/net_utils/net_utils.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"errors"
 	"net"
-	"net/url"
 	"regexp"
 	"time"
 
@@ -13,19 +12,6 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-func HostWithOrWithoutPort(addr string, withPort bool) (string, error) {
-	host, err := url.Parse(addr)
-	if err != nil {
-		return "", err
-	}
-	if !withPort {
-		re := regexp.MustCompile(`:\d+`)
-		result := re.Split(host.Host, -1)
-		return result[0], nil
-	}
-	return host.Host, nil
-}
-
 func ShouldEncryptProxy(addr string) (bool, error) {
 	pattern := `http`
 	re := regexp.MustCompile(pattern)

…linux-client/daemon/device_connection.go …linux-client/serverws/mgmt_connection.godevices/linux-client/daemon/device_connection.go renamed to devices/linux-client/serverws/mgmt_connection.go devices/linux-client/daemon/device_connection.go renamed to devices/linux-client/serverws/mgmt_connection.go

@@ -1,13 +1,11 @@
-package daemon
+package serverws
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/json"
-	"net/http"
-	"net/url"
 	"sync"
 
-	netUtils "github.com/antmicro/rdfm/daemon/net_utils"
 	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
 )
@@ -19,7 +17,7 @@ const (
 	stateConnected
 )
 
-type DeviceConnection struct {
+type DeviceManagementConnection struct {
 	ws           *websocket.Conn
 	rx           chan []byte
 	txMut        sync.Mutex
@@ -39,7 +37,7 @@ func (e *ConnClosedError) Error() string {
 	return e.message
 }
 
-func (d *DeviceConnection) startRecvLoop(cancelCtx context.Context) error {
+func (d *DeviceManagementConnection) startRecvLoop(cancelCtx context.Context) error {
 	for {
 		d.wsMut.RLock()
 		if d.ws == nil {
@@ -58,7 +56,7 @@ func (d *DeviceConnection) startRecvLoop(cancelCtx context.Context) error {
 	}
 }
 
-func (d *DeviceConnection) Close() error {
+func (d *DeviceManagementConnection) Close() error {
 	d.stateCnd.L.Lock()
 	defer d.stateCnd.L.Unlock()
 	err := d.tryClose()
@@ -68,7 +66,7 @@ func (d *DeviceConnection) Close() error {
 	return err
 }
 
-func (d *DeviceConnection) tryClose() error {
+func (d *DeviceManagementConnection) tryClose() error {
 	d.wsMut.RLock()
 	defer d.wsMut.RUnlock()
 	if d.ws != nil {
@@ -77,18 +75,18 @@ func (d *DeviceConnection) tryClose() error {
 	return nil
 }
 
-func (d *DeviceConnection) setState(state deviceConnectionState) {
+func (d *DeviceManagementConnection) setState(state deviceConnectionState) {
 	d.stateCnd.L.Lock()
 	d.state = state
 	d.stateCnd.Broadcast()
 	d.stateCnd.L.Unlock()
 }
 
-func (d *DeviceConnection) announceCapabilities() error {
+func (d *DeviceManagementConnection) announceCapabilities() error {
 	res := map[string]interface{}{
 		"method":       "capability_report",
 		"capabilities": d.capabilities,
-	} 
+	}
 
 	msg, err := json.Marshal(res)
 	if err != nil {
@@ -99,15 +97,19 @@ func (d *DeviceConnection) announceCapabilities() error {
 	return nil
 }
 
-func (d *DeviceConnection) CreateConnection(deviceToken string, cancelCtx context.Context) error {
+func (d *DeviceManagementConnection) CreateConnection(deviceToken string, cancelCtx context.Context) error {
 	var wg sync.WaitGroup
 
 	defer func() {
 		d.tryClose()
 		d.setState(stateDisconnected)
 	}()
 
-	err := d.prepareWs(deviceToken)
+	endpoint, err := formatDeviceWsUrl(d.serverUrl)
+	if err != nil {
+		return err
+	}
+	d.ws, err = ConnectToRdfmWs(d.dialer, endpoint, deviceToken)
 	if err != nil {
 		return err
 	}
@@ -156,18 +158,18 @@ func (d *DeviceConnection) CreateConnection(deviceToken string, cancelCtx contex
 	return err
 }
 
-func NewDeviceConnection(serverUrl string, dialer websocket.Dialer, buffer_size int) *DeviceConnection {
-	dc := new(DeviceConnection)
+func NewDeviceConnection(serverUrl string, tlsConf *tls.Config, buffer_size int) *DeviceManagementConnection {
+	dc := new(DeviceManagementConnection)
 	dc.rx = make(chan []byte, buffer_size)
 	dc.serverUrl = serverUrl
-	dc.dialer = dialer
+	dc.dialer = *prepareWsDialer(tlsConf)
 	dc.state = stateDisconnected
 	dc.stateCnd = sync.NewCond(&sync.Mutex{})
 	dc.capabilities = make(map[string]bool)
 	return dc
 }
 
-func (d *DeviceConnection) Recv(cancelCtx context.Context) []byte {
+func (d *DeviceManagementConnection) Recv(cancelCtx context.Context) []byte {
 	select {
 	case msg, ok := <-d.rx:
 		if ok {
@@ -178,7 +180,7 @@ func (d *DeviceConnection) Recv(cancelCtx context.Context) []byte {
 	return nil
 }
 
-func (d *DeviceConnection) Send(msg []byte) error {
+func (d *DeviceManagementConnection) Send(msg []byte) error {
 	d.txMut.Lock()
 	d.wsMut.RLock()
 	defer d.txMut.Unlock()
@@ -195,47 +197,14 @@ func (d *DeviceConnection) Send(msg []byte) error {
 	return nil
 }
 
-func (d *DeviceConnection) EnsureReady() {
+func (d *DeviceManagementConnection) EnsureReady() {
 	d.stateCnd.L.Lock()
 	defer d.stateCnd.L.Unlock()
 	for d.state != stateConnected {
 		d.stateCnd.Wait()
 	}
 }
 
-func (d *DeviceConnection) prepareWs(deviceToken string) error {
-	d.wsMut.Lock()
-	defer d.wsMut.Unlock()
-
-	// Get the endpoint URL
-	addr, err := netUtils.HostWithOrWithoutPort(d.serverUrl, true)
-	if err != nil {
-		return err
-	}
-
-	scheme := "ws"
-	if d.dialer.TLSClientConfig != nil {
-		scheme = "wss"
-	}
-
-	u := url.URL{
-		Scheme: scheme,
-		Host:   addr,
-		Path:   "/api/v1/devices/ws",
-	}
-
-	// Connect to the endpoint
-	log.Infoln("Connecting to", u.String())
-
-	authHeader := http.Header{
-		"Authorization": []string{"Bearer token=" + deviceToken},
-	}
-
-	ws, _, err := d.dialer.Dial(u.String(), authHeader)
-	d.ws = ws
-	return err
-}
-
-func (d *DeviceConnection) SetCapability(cap string, value bool) {
+func (d *DeviceManagementConnection) SetCapability(cap string, value bool) {
 	d.capabilities[cap] = value
 }

…linux-client/daemon/requests/requests.go devices/linux-client/serverws/requests.godevices/linux-client/daemon/requests/requests.go renamed to devices/linux-client/serverws/requests.go devices/linux-client/daemon/requests/requests.go renamed to devices/linux-client/serverws/requests.go

@@ -1,4 +1,4 @@
-package requests
+package serverws
 
 import (
 	"encoding/json"

devices/linux-client/serverws/ws.go

@@ -0,0 +1,42 @@
+package serverws
+
+import (
+	"crypto/tls"
+	"net/http"
+	"net/url"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
+func formatDeviceWsUrl(serverUrl string) (string, error) {
+	return url.JoinPath(serverUrl, "/api/v1/devices/ws")
+}
+
+func prepareWsDialer(tlsConf *tls.Config) *websocket.Dialer {
+	if tlsConf != nil {
+		return &websocket.Dialer{TLSClientConfig: tlsConf}
+	} else {
+		return websocket.DefaultDialer
+	}
+}
+
+func ConnectToRdfmWs(dialer websocket.Dialer, endpoint string, deviceToken string) (*websocket.Conn, error) {
+	scheme := "ws"
+	if dialer.TLSClientConfig != nil {
+		scheme = "wss"
+	}
+	url, err := url.Parse(endpoint)
+	if err != nil {
+		return nil, err
+	}
+	url.Scheme = scheme
+	wsEndpoint := url.String()
+
+	log.Infoln("Connecting to", wsEndpoint)
+	authHeader := http.Header{
+		"Authorization": []string{"Bearer token=" + deviceToken},
+	}
+	ws, _, err := dialer.Dial(wsEndpoint, authHeader)
+	return ws, err
+}