Merge branch 'main' into embedded-vnc

This commit is contained in:
Viktor Liu
2026-09-09 09:30:51 +02:00
115 changed files with 4562 additions and 1324 deletions
+1 -1
View File
@@ -12,7 +12,7 @@ import (
// Client is the interface for the management service client.
type Client interface {
io.Closer
Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error
Register(setupKey string, jwtToken string, sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
Login(sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
+72 -1
View File
@@ -2,9 +2,11 @@ package client
import (
"context"
"fmt"
"net"
"os"
"sync"
"sync/atomic"
"testing"
"time"
@@ -305,7 +307,7 @@ func TestClient_Sync(t *testing.T) {
defer cancel()
go func() {
err = client.Sync(ctx, info, func(msg *mgmtProto.SyncResponse) error {
err = client.Sync(ctx, func(context.Context) *system.Info { return info }, func(msg *mgmtProto.SyncResponse) error {
ch <- msg
return nil
})
@@ -397,6 +399,75 @@ func wgKeyFromBytes(raw []byte) string {
return k.String()
}
func TestClient_SyncGathersInfoOnEveryConnect(t *testing.T) {
s, lis, mgmtMockServer, serverKey := startMockManagement(t)
defer s.GracefulStop()
testKey, err := wgtypes.GenerateKey()
require.NoError(t, err)
hostnames := make(chan string, 2)
mgmtMockServer.SyncFunc = func(msg *mgmtProto.EncryptedMessage, _ mgmtProto.ManagementService_SyncServer) error {
peerKey, err := wgtypes.ParseKey(msg.GetWgPubKey())
if err != nil {
t.Errorf("invalid peer key: %v", err)
return status.Error(codes.InvalidArgument, err.Error())
}
syncReq := &mgmtProto.SyncRequest{}
if err := encryption.DecryptMessage(peerKey, serverKey, msg.Body, syncReq); err != nil {
t.Errorf("decrypt sync request: %v", err)
return status.Error(codes.InvalidArgument, err.Error())
}
select {
case hostnames <- syncReq.GetMeta().GetHostname():
default:
}
// Returning closes the stream, so the client reconnects and gathers again.
return nil
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
client, err := NewClient(ctx, lis.Addr().String(), testKey, false)
require.NoError(t, err)
var gathers atomic.Int32
done := make(chan struct{})
go func() {
defer close(done)
_ = client.Sync(ctx, func(ctx context.Context) *system.Info {
info := system.GetInfo(ctx)
info.Hostname = fmt.Sprintf("host-%d", gathers.Add(1))
return info
}, func(*mgmtProto.SyncResponse) error { return nil })
}()
// A connect attempt can fail before it reaches the server, so the sequence
// numbers seen here may skip. What matters is that the reconnect carries a
// newly gathered info instead of the one sent on the previous stream.
var seen []int
for len(seen) < 2 {
select {
case got := <-hostnames:
var n int
_, err := fmt.Sscanf(got, "host-%d", &n)
require.NoError(t, err, "hostname should carry the gather sequence number")
seen = append(seen, n)
case <-time.After(10 * time.Second):
t.Fatalf("timeout waiting for the second sync request, got %v", seen)
}
}
assert.Greater(t, seen[1], seen[0], "the reconnect should carry a newly gathered info")
cancel()
select {
case <-done:
case <-time.After(10 * time.Second):
t.Fatal("timeout waiting for Sync to return after cancel")
}
}
func Test_SystemMetaDataFromClient(t *testing.T) {
s, lis, mgmtMockServer, serverKey := startMockManagement(t)
defer s.GracefulStop()
+4 -4
View File
@@ -205,9 +205,9 @@ func (c *GrpcClient) ready() bool {
// Sync wraps the real client's Sync endpoint call and takes care of retries and encryption/decryption of messages
// Blocking request. The result will be sent via msgHandler callback function
func (c *GrpcClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
func (c *GrpcClient) Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error {
return c.handleSyncStream(ctx, serverPubKey, sysInfo, msgHandler, backOff)
return c.handleSyncStream(ctx, serverPubKey, getInfo, msgHandler, backOff)
})
}
@@ -424,11 +424,11 @@ func (c *GrpcClient) sendJobResponse(
return nil
}
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error {
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error {
ctx, cancelStream := context.WithCancel(ctx)
defer cancelStream()
stream, err := c.connectToSyncStream(ctx, serverPubKey, sysInfo)
stream, err := c.connectToSyncStream(ctx, serverPubKey, getInfo(ctx))
if err != nil {
log.Debugf("failed to open Management Service stream: %s", err)
c.notifyDisconnected(err)
+3 -3
View File
@@ -11,7 +11,7 @@ import (
// MockClient is a mock implementation of the Client interface for testing.
type MockClient struct {
CloseFunc func() error
SyncFunc func(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
SyncFunc func(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
RegisterFunc func(setupKey string, jwtToken string, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
LoginFunc func(info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
ExtendAuthSessionFunc func(info *system.Info, jwtToken string) (*proto.ExtendAuthSessionResponse, error)
@@ -38,11 +38,11 @@ func (m *MockClient) Close() error {
return m.CloseFunc()
}
func (m *MockClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
func (m *MockClient) Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
if m.SyncFunc == nil {
return nil
}
return m.SyncFunc(ctx, sysInfo, msgHandler)
return m.SyncFunc(ctx, getInfo, msgHandler)
}
func (m *MockClient) Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error {