mirror of
https://github.com/fankes/beszel.git
synced 2025-10-18 17:29:28 +08:00
rename /src to /internal (sorry i'll fix the prs)
This commit is contained in:
henrygd
182
internal/hub/ws/ws.go
Normal file
182
internal/hub/ws/ws.go
Normal file
@@ -0,0 +1,182 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
"weak"
|
||||
|
||||
"github.com/henrygd/beszel/src/entities/system"
|
||||
|
||||
"github.com/henrygd/beszel/src/common"
|
||||
|
||||
"github.com/fxamacker/cbor/v2"
|
||||
"github.com/lxzan/gws"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
const (
|
||||
deadline = 70 * time.Second
|
||||
)
|
||||
|
||||
// Handler implements the WebSocket event handler for agent connections.
|
||||
type Handler struct {
|
||||
gws.BuiltinEventHandler
|
||||
}
|
||||
|
||||
// WsConn represents a WebSocket connection to an agent.
|
||||
type WsConn struct {
|
||||
conn *gws.Conn
|
||||
responseChan chan *gws.Message
|
||||
DownChan chan struct{}
|
||||
}
|
||||
|
||||
// FingerprintRecord is fingerprints collection record data in the hub
|
||||
type FingerprintRecord struct {
|
||||
Id string `db:"id"`
|
||||
SystemId string `db:"system"`
|
||||
Fingerprint string `db:"fingerprint"`
|
||||
Token string `db:"token"`
|
||||
}
|
||||
|
||||
var upgrader *gws.Upgrader
|
||||
|
||||
// GetUpgrader returns a singleton WebSocket upgrader instance.
|
||||
func GetUpgrader() *gws.Upgrader {
|
||||
if upgrader != nil {
|
||||
return upgrader
|
||||
}
|
||||
handler := &Handler{}
|
||||
upgrader = gws.NewUpgrader(handler, &gws.ServerOption{})
|
||||
return upgrader
|
||||
}
|
||||
|
||||
// NewWsConnection creates a new WebSocket connection wrapper.
|
||||
func NewWsConnection(conn *gws.Conn) *WsConn {
|
||||
return &WsConn{
|
||||
conn: conn,
|
||||
responseChan: make(chan *gws.Message, 1),
|
||||
DownChan: make(chan struct{}, 1),
|
||||
}
|
||||
}
|
||||
|
||||
// OnOpen sets a deadline for the WebSocket connection.
|
||||
func (h *Handler) OnOpen(conn *gws.Conn) {
|
||||
conn.SetDeadline(time.Now().Add(deadline))
|
||||
}
|
||||
|
||||
// OnMessage routes incoming WebSocket messages to the response channel.
|
||||
func (h *Handler) OnMessage(conn *gws.Conn, message *gws.Message) {
|
||||
conn.SetDeadline(time.Now().Add(deadline))
|
||||
if message.Opcode != gws.OpcodeBinary || message.Data.Len() == 0 {
|
||||
return
|
||||
}
|
||||
wsConn, ok := conn.Session().Load("wsConn")
|
||||
if !ok {
|
||||
_ = conn.WriteClose(1000, nil)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case wsConn.(*WsConn).responseChan <- message:
|
||||
default:
|
||||
// close if the connection is not expecting a response
|
||||
wsConn.(*WsConn).Close(nil)
|
||||
}
|
||||
}
|
||||
|
||||
// OnClose handles WebSocket connection closures and triggers system down status after delay.
|
||||
func (h *Handler) OnClose(conn *gws.Conn, err error) {
|
||||
wsConn, ok := conn.Session().Load("wsConn")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
wsConn.(*WsConn).conn = nil
|
||||
// wait 5 seconds to allow reconnection before setting system down
|
||||
// use a weak pointer to avoid keeping references if the system is removed
|
||||
go func(downChan weak.Pointer[chan struct{}]) {
|
||||
time.Sleep(5 * time.Second)
|
||||
downChanValue := downChan.Value()
|
||||
if downChanValue != nil {
|
||||
*downChanValue <- struct{}{}
|
||||
}
|
||||
}(weak.Make(&wsConn.(*WsConn).DownChan))
|
||||
}
|
||||
|
||||
// Close terminates the WebSocket connection gracefully.
|
||||
func (ws *WsConn) Close(msg []byte) {
|
||||
if ws.IsConnected() {
|
||||
ws.conn.WriteClose(1000, msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Ping sends a ping frame to keep the connection alive.
|
||||
func (ws *WsConn) Ping() error {
|
||||
ws.conn.SetDeadline(time.Now().Add(deadline))
|
||||
return ws.conn.WritePing(nil)
|
||||
}
|
||||
|
||||
// sendMessage encodes data to CBOR and sends it as a binary message to the agent.
|
||||
func (ws *WsConn) sendMessage(data common.HubRequest[any]) error {
|
||||
if ws.conn == nil {
|
||||
return gws.ErrConnClosed
|
||||
}
|
||||
bytes, err := cbor.Marshal(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return ws.conn.WriteMessage(gws.OpcodeBinary, bytes)
|
||||
}
|
||||
|
||||
// RequestSystemData requests system metrics from the agent and unmarshals the response.
|
||||
func (ws *WsConn) RequestSystemData(data *system.CombinedData) error {
|
||||
var message *gws.Message
|
||||
|
||||
ws.sendMessage(common.HubRequest[any]{
|
||||
Action: common.GetData,
|
||||
})
|
||||
select {
|
||||
case <-time.After(10 * time.Second):
|
||||
ws.Close(nil)
|
||||
return gws.ErrConnClosed
|
||||
case message = <-ws.responseChan:
|
||||
}
|
||||
defer message.Close()
|
||||
return cbor.Unmarshal(message.Data.Bytes(), data)
|
||||
}
|
||||
|
||||
// GetFingerprint authenticates with the agent using SSH signature and returns the agent's fingerprint.
|
||||
func (ws *WsConn) GetFingerprint(token string, signer ssh.Signer, needSysInfo bool) (common.FingerprintResponse, error) {
|
||||
var clientFingerprint common.FingerprintResponse
|
||||
challenge := []byte(token)
|
||||
|
||||
signature, err := signer.Sign(nil, challenge)
|
||||
if err != nil {
|
||||
return clientFingerprint, err
|
||||
}
|
||||
|
||||
err = ws.sendMessage(common.HubRequest[any]{
|
||||
Action: common.CheckFingerprint,
|
||||
Data: common.FingerprintRequest{
|
||||
Signature: signature.Blob,
|
||||
NeedSysInfo: needSysInfo,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return clientFingerprint, err
|
||||
}
|
||||
|
||||
var message *gws.Message
|
||||
select {
|
||||
case message = <-ws.responseChan:
|
||||
case <-time.After(10 * time.Second):
|
||||
return clientFingerprint, errors.New("request expired")
|
||||
}
|
||||
defer message.Close()
|
||||
|
||||
err = cbor.Unmarshal(message.Data.Bytes(), &clientFingerprint)
|
||||
return clientFingerprint, err
|
||||
}
|
||||
|
||||
// IsConnected returns true if the WebSocket connection is active.
|
||||
func (ws *WsConn) IsConnected() bool {
|
||||
return ws.conn != nil
|
||||
}
|
||||
222
internal/hub/ws/ws_test.go
Normal file
222
internal/hub/ws/ws_test.go
Normal file
@@ -0,0 +1,222 @@
|
||||
//go:build testing
|
||||
// +build testing
|
||||
|
||||
package ws
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/henrygd/beszel/src/common"
|
||||
|
||||
"github.com/fxamacker/cbor/v2"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
// TestGetUpgrader tests the singleton upgrader
|
||||
func TestGetUpgrader(t *testing.T) {
|
||||
// Reset the global upgrader to test singleton behavior
|
||||
upgrader = nil
|
||||
|
||||
// First call should create the upgrader
|
||||
upgrader1 := GetUpgrader()
|
||||
assert.NotNil(t, upgrader1, "Upgrader should not be nil")
|
||||
|
||||
// Second call should return the same instance
|
||||
upgrader2 := GetUpgrader()
|
||||
assert.Same(t, upgrader1, upgrader2, "Should return the same upgrader instance")
|
||||
|
||||
// Verify it's properly configured
|
||||
assert.NotNil(t, upgrader1, "Upgrader should be configured")
|
||||
}
|
||||
|
||||
// TestNewWsConnection tests WebSocket connection creation
|
||||
func TestNewWsConnection(t *testing.T) {
|
||||
// We can't easily mock gws.Conn, so we'll pass nil and test the structure
|
||||
wsConn := NewWsConnection(nil)
|
||||
|
||||
assert.NotNil(t, wsConn, "WebSocket connection should not be nil")
|
||||
assert.Nil(t, wsConn.conn, "Connection should be nil as passed")
|
||||
assert.NotNil(t, wsConn.responseChan, "Response channel should be initialized")
|
||||
assert.NotNil(t, wsConn.DownChan, "Down channel should be initialized")
|
||||
assert.Equal(t, 1, cap(wsConn.responseChan), "Response channel should have capacity of 1")
|
||||
assert.Equal(t, 1, cap(wsConn.DownChan), "Down channel should have capacity of 1")
|
||||
}
|
||||
|
||||
// TestWsConn_IsConnected tests the connection status check
|
||||
func TestWsConn_IsConnected(t *testing.T) {
|
||||
// Test with nil connection
|
||||
wsConn := NewWsConnection(nil)
|
||||
assert.False(t, wsConn.IsConnected(), "Should not be connected when conn is nil")
|
||||
}
|
||||
|
||||
// TestWsConn_Close tests the connection closing with nil connection
|
||||
func TestWsConn_Close(t *testing.T) {
|
||||
wsConn := NewWsConnection(nil)
|
||||
|
||||
// Should handle nil connection gracefully
|
||||
assert.NotPanics(t, func() {
|
||||
wsConn.Close([]byte("test message"))
|
||||
}, "Should not panic when closing nil connection")
|
||||
}
|
||||
|
||||
// TestWsConn_SendMessage_CBOR tests CBOR encoding in sendMessage
|
||||
func TestWsConn_SendMessage_CBOR(t *testing.T) {
|
||||
wsConn := NewWsConnection(nil)
|
||||
|
||||
testData := common.HubRequest[any]{
|
||||
Action: common.GetData,
|
||||
Data: "test data",
|
||||
}
|
||||
|
||||
// This will fail because conn is nil, but we can test the CBOR encoding logic
|
||||
// by checking that the function properly encodes to CBOR before failing
|
||||
err := wsConn.sendMessage(testData)
|
||||
assert.Error(t, err, "Should error with nil connection")
|
||||
|
||||
// Test CBOR encoding separately
|
||||
bytes, err := cbor.Marshal(testData)
|
||||
assert.NoError(t, err, "Should encode to CBOR successfully")
|
||||
|
||||
// Verify we can decode it back
|
||||
var decodedData common.HubRequest[any]
|
||||
err = cbor.Unmarshal(bytes, &decodedData)
|
||||
assert.NoError(t, err, "Should decode from CBOR successfully")
|
||||
assert.Equal(t, testData.Action, decodedData.Action, "Action should match")
|
||||
}
|
||||
|
||||
// TestWsConn_GetFingerprint_SignatureGeneration tests signature creation logic
|
||||
func TestWsConn_GetFingerprint_SignatureGeneration(t *testing.T) {
|
||||
// Generate test key pair
|
||||
_, privKey, err := ed25519.GenerateKey(nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
signer, err := ssh.NewSignerFromKey(privKey)
|
||||
require.NoError(t, err)
|
||||
|
||||
token := "test-token"
|
||||
|
||||
// This will timeout since conn is nil, but we can verify the signature logic
|
||||
// We can't test the full flow, but we can test that the signature is created properly
|
||||
challenge := []byte(token)
|
||||
signature, err := signer.Sign(nil, challenge)
|
||||
assert.NoError(t, err, "Should create signature successfully")
|
||||
assert.NotEmpty(t, signature.Blob, "Signature blob should not be empty")
|
||||
assert.Equal(t, signer.PublicKey().Type(), signature.Format, "Signature format should match key type")
|
||||
|
||||
// Test the fingerprint request structure
|
||||
fpRequest := common.FingerprintRequest{
|
||||
Signature: signature.Blob,
|
||||
NeedSysInfo: true,
|
||||
}
|
||||
|
||||
// Test CBOR encoding of fingerprint request
|
||||
fpData, err := cbor.Marshal(fpRequest)
|
||||
assert.NoError(t, err, "Should encode fingerprint request to CBOR")
|
||||
|
||||
var decodedFpRequest common.FingerprintRequest
|
||||
err = cbor.Unmarshal(fpData, &decodedFpRequest)
|
||||
assert.NoError(t, err, "Should decode fingerprint request from CBOR")
|
||||
assert.Equal(t, fpRequest.Signature, decodedFpRequest.Signature, "Signature should match")
|
||||
assert.Equal(t, fpRequest.NeedSysInfo, decodedFpRequest.NeedSysInfo, "NeedSysInfo should match")
|
||||
|
||||
// Test the full hub request structure
|
||||
hubRequest := common.HubRequest[any]{
|
||||
Action: common.CheckFingerprint,
|
||||
Data: fpRequest,
|
||||
}
|
||||
|
||||
hubData, err := cbor.Marshal(hubRequest)
|
||||
assert.NoError(t, err, "Should encode hub request to CBOR")
|
||||
|
||||
var decodedHubRequest common.HubRequest[cbor.RawMessage]
|
||||
err = cbor.Unmarshal(hubData, &decodedHubRequest)
|
||||
assert.NoError(t, err, "Should decode hub request from CBOR")
|
||||
assert.Equal(t, common.CheckFingerprint, decodedHubRequest.Action, "Action should be CheckFingerprint")
|
||||
}
|
||||
|
||||
// TestWsConn_RequestSystemData_RequestFormat tests system data request format
|
||||
func TestWsConn_RequestSystemData_RequestFormat(t *testing.T) {
|
||||
// Test the request format that would be sent
|
||||
request := common.HubRequest[any]{
|
||||
Action: common.GetData,
|
||||
}
|
||||
|
||||
// Test CBOR encoding
|
||||
data, err := cbor.Marshal(request)
|
||||
assert.NoError(t, err, "Should encode request to CBOR")
|
||||
|
||||
// Test decoding
|
||||
var decodedRequest common.HubRequest[any]
|
||||
err = cbor.Unmarshal(data, &decodedRequest)
|
||||
assert.NoError(t, err, "Should decode request from CBOR")
|
||||
assert.Equal(t, common.GetData, decodedRequest.Action, "Should have GetData action")
|
||||
}
|
||||
|
||||
// TestFingerprintRecord tests the FingerprintRecord struct
|
||||
func TestFingerprintRecord(t *testing.T) {
|
||||
record := FingerprintRecord{
|
||||
Id: "test-id",
|
||||
SystemId: "system-123",
|
||||
Fingerprint: "test-fingerprint",
|
||||
Token: "test-token",
|
||||
}
|
||||
|
||||
assert.Equal(t, "test-id", record.Id)
|
||||
assert.Equal(t, "system-123", record.SystemId)
|
||||
assert.Equal(t, "test-fingerprint", record.Fingerprint)
|
||||
assert.Equal(t, "test-token", record.Token)
|
||||
}
|
||||
|
||||
// TestDeadlineConstant tests that the deadline constant is reasonable
|
||||
func TestDeadlineConstant(t *testing.T) {
|
||||
assert.Equal(t, 70*time.Second, deadline, "Deadline should be 70 seconds")
|
||||
}
|
||||
|
||||
// TestCommonActions tests that the common actions are properly defined
|
||||
func TestCommonActions(t *testing.T) {
|
||||
// Test that the actions we use exist and have expected values
|
||||
assert.Equal(t, common.WebSocketAction(0), common.GetData, "GetData should be action 0")
|
||||
assert.Equal(t, common.WebSocketAction(1), common.CheckFingerprint, "CheckFingerprint should be action 1")
|
||||
}
|
||||
|
||||
// TestHandler tests that we can create a Handler
|
||||
func TestHandler(t *testing.T) {
|
||||
handler := &Handler{}
|
||||
assert.NotNil(t, handler, "Handler should be created successfully")
|
||||
|
||||
// The Handler embeds gws.BuiltinEventHandler, so it should have the embedded type
|
||||
assert.NotNil(t, handler.BuiltinEventHandler, "Should have embedded BuiltinEventHandler")
|
||||
}
|
||||
|
||||
// TestWsConnChannelBehavior tests channel behavior without WebSocket connections
|
||||
func TestWsConnChannelBehavior(t *testing.T) {
|
||||
wsConn := NewWsConnection(nil)
|
||||
|
||||
// Test that channels are properly initialized and can be used
|
||||
select {
|
||||
case wsConn.DownChan <- struct{}{}:
|
||||
// Should be able to write to channel
|
||||
default:
|
||||
t.Error("Should be able to write to DownChan")
|
||||
}
|
||||
|
||||
// Test reading from DownChan
|
||||
select {
|
||||
case <-wsConn.DownChan:
|
||||
// Should be able to read from channel
|
||||
case <-time.After(10 * time.Millisecond):
|
||||
t.Error("Should be able to read from DownChan")
|
||||
}
|
||||
|
||||
// Response channel should be empty initially
|
||||
select {
|
||||
case <-wsConn.responseChan:
|
||||
t.Error("Response channel should be empty initially")
|
||||
default:
|
||||
// Expected - channel should be empty
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user