From fe04cf93aa582fb502650eff512162b86b659a60 Mon Sep 17 00:00:00 2001 From: Snider Date: Mon, 9 Mar 2026 08:30:59 +0000 Subject: [PATCH] fix: improve UEPS packet validation and worker error handling Co-Authored-By: Claude Opus 4.6 --- node/integration_test.go | 6 +- node/worker.go | 7 +- node/worker_test.go | 148 ++++++++++++++++++++++++++++++----- ueps/packet.go | 31 ++++---- ueps/packet_coverage_test.go | 26 +++--- 5 files changed, 167 insertions(+), 51 deletions(-) diff --git a/node/integration_test.go b/node/integration_test.go index 30b8588..cd21f36 100644 --- a/node/integration_test.go +++ b/node/integration_test.go @@ -370,7 +370,7 @@ func TestIntegration_IdentityPersistenceAndReload(t *testing.T) { assert.Equal(t, original.Role, reloaded.Role, "Role should persist") // Verify the reloaded key can derive the same shared secret. - kp, err := stmfGenerateKeyPair() + kp, err := stmfGenerateKeyPair(t.TempDir()) require.NoError(t, err) secret1, err := nm1.DeriveSharedSecret(kp) @@ -385,8 +385,7 @@ func TestIntegration_IdentityPersistenceAndReload(t *testing.T) { // stmfGenerateKeyPair is a helper that generates a keypair and returns // the public key as base64 (for use in DeriveSharedSecret tests). -func stmfGenerateKeyPair() (string, error) { - dir, _ := filepath.Abs("/tmp/stmf-test-" + time.Now().Format("20060102150405.000")) +func stmfGenerateKeyPair(dir string) (string, error) { nm, err := NewNodeManagerWithPaths( filepath.Join(dir, "private.key"), filepath.Join(dir, "node.json"), @@ -400,6 +399,7 @@ func stmfGenerateKeyPair() (string, error) { return nm.GetIdentity().PublicKey, nil } + // TestIntegration_UEPSFullRoundTrip exercises a complete UEPS packet // lifecycle: build, sign, transmit (simulated), read, verify, dispatch. func TestIntegration_UEPSFullRoundTrip(t *testing.T) { diff --git a/node/worker.go b/node/worker.go index 6aaceab..06da7b6 100644 --- a/node/worker.go +++ b/node/worker.go @@ -42,6 +42,7 @@ type Worker struct { minerManager MinerManager profileManager ProfileManager startTime time.Time + DataDir string // Base directory for deployments (defaults to xdg.DataHome) } // NewWorker creates a new Worker instance. @@ -50,9 +51,11 @@ func NewWorker(node *NodeManager, transport *Transport) *Worker { node: node, transport: transport, startTime: time.Now(), + DataDir: xdg.DataHome, } } + // SetMinerManager sets the miner manager for handling miner operations. func (w *Worker) SetMinerManager(manager MinerManager) { w.minerManager = manager @@ -350,8 +353,8 @@ func (w *Worker) handleDeploy(conn *PeerConnection, msg *Message) (*Message, err case BundleMiner, BundleFull: // Determine installation directory - // We use xdg.DataHome/lethean-desktop/miners/ - minersDir := filepath.Join(xdg.DataHome, "lethean-desktop", "miners") + // We use w.DataDir/lethean-desktop/miners/ + minersDir := filepath.Join(w.DataDir, "lethean-desktop", "miners") installDir := filepath.Join(minersDir, payload.Name) logging.Info("deploying miner bundle", logging.Fields{ diff --git a/node/worker_test.go b/node/worker_test.go index 3d4b128..ee3ed31 100644 --- a/node/worker_test.go +++ b/node/worker_test.go @@ -25,7 +25,11 @@ func TestNewWorker(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -40,6 +44,7 @@ func TestNewWorker(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() if worker == nil { t.Fatal("NewWorker returned nil") @@ -56,7 +61,11 @@ func TestWorker_SetMinerManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -71,6 +80,7 @@ func TestWorker_SetMinerManager(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() mockManager := &mockMinerManager{} worker.SetMinerManager(mockManager) @@ -84,7 +94,11 @@ func TestWorker_SetProfileManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -99,6 +113,7 @@ func TestWorker_SetProfileManager(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() mockProfile := &mockProfileManager{} worker.SetProfileManager(mockProfile) @@ -112,7 +127,11 @@ func TestWorker_HandlePing(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -127,6 +146,7 @@ func TestWorker_HandlePing(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Create a ping message identity := nm.GetIdentity() @@ -171,7 +191,11 @@ func TestWorker_HandleGetStats(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -186,6 +210,7 @@ func TestWorker_HandleGetStats(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Create a get_stats message identity := nm.GetIdentity() @@ -229,7 +254,11 @@ func TestWorker_HandleStartMiner_NoManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -244,6 +273,7 @@ func TestWorker_HandleStartMiner_NoManager(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Create a start_miner message identity := nm.GetIdentity() @@ -267,7 +297,11 @@ func TestWorker_HandleStopMiner_NoManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -282,6 +316,7 @@ func TestWorker_HandleStopMiner_NoManager(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Create a stop_miner message identity := nm.GetIdentity() @@ -305,7 +340,11 @@ func TestWorker_HandleGetLogs_NoManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -320,6 +359,7 @@ func TestWorker_HandleGetLogs_NoManager(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Create a get_logs message identity := nm.GetIdentity() @@ -343,7 +383,11 @@ func TestWorker_HandleDeploy_Profile(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -358,6 +402,7 @@ func TestWorker_HandleDeploy_Profile(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Create a deploy message for profile identity := nm.GetIdentity() @@ -385,7 +430,11 @@ func TestWorker_HandleDeploy_UnknownType(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -400,6 +449,7 @@ func TestWorker_HandleDeploy_UnknownType(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Create a deploy message with unknown type identity := nm.GetIdentity() @@ -566,7 +616,11 @@ func TestWorker_HandleStartMiner_WithManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -581,6 +635,7 @@ func TestWorker_HandleStartMiner_WithManager(t *testing.T) { transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() mm := &mockMinerManager{ miners: []MinerInstance{}, @@ -733,7 +788,11 @@ func TestWorker_HandleStopMiner_WithManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -746,6 +805,7 @@ func TestWorker_HandleStopMiner_WithManager(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() identity := nm.GetIdentity() t.Run("Success", func(t *testing.T) { @@ -795,7 +855,11 @@ func TestWorker_HandleGetLogs_WithManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -808,6 +872,7 @@ func TestWorker_HandleGetLogs_WithManager(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() identity := nm.GetIdentity() t.Run("Success", func(t *testing.T) { @@ -900,7 +965,11 @@ func TestWorker_HandleGetStats_WithMinerManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -913,6 +982,7 @@ func TestWorker_HandleGetStats_WithMinerManager(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() identity := nm.GetIdentity() // Set miner manager with miners that have real stats @@ -959,7 +1029,11 @@ func TestWorker_HandleMessage_UnknownType(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -972,6 +1046,7 @@ func TestWorker_HandleMessage_UnknownType(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() identity := nm.GetIdentity() msg, _ := NewMessage("unknown_type", "sender-id", identity.ID, nil) @@ -984,7 +1059,11 @@ func TestWorker_HandleDeploy_ProfileWithManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -997,6 +1076,7 @@ func TestWorker_HandleDeploy_ProfileWithManager(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() pm := &mockProfileManagerFull{profiles: make(map[string]any)} worker.SetProfileManager(pm) @@ -1037,7 +1117,11 @@ func TestWorker_HandleDeploy_ProfileSaveFails(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -1050,6 +1134,7 @@ func TestWorker_HandleDeploy_ProfileSaveFails(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() worker.SetProfileManager(&mockProfileManagerFailing{}) identity := nm.GetIdentity() @@ -1081,7 +1166,11 @@ func TestWorker_HandleDeploy_MinerBundle(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -1094,6 +1183,7 @@ func TestWorker_HandleDeploy_MinerBundle(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() pm := &mockProfileManagerFull{profiles: make(map[string]any)} worker.SetProfileManager(pm) @@ -1143,7 +1233,11 @@ func TestWorker_HandleDeploy_FullBundle(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -1156,6 +1250,7 @@ func TestWorker_HandleDeploy_FullBundle(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() identity := nm.GetIdentity() @@ -1197,7 +1292,11 @@ func TestWorker_HandleDeploy_MinerBundle_WithProfileManager(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, err := NewNodeManager() + dir := t.TempDir() + nm, err := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) if err != nil { t.Fatalf("failed to create node manager: %v", err) } @@ -1210,6 +1309,7 @@ func TestWorker_HandleDeploy_MinerBundle_WithProfileManager(t *testing.T) { } transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() // Set a failing profile manager to exercise the warn-and-continue path worker.SetProfileManager(&mockProfileManagerFailing{}) @@ -1256,11 +1356,16 @@ func TestWorker_HandleDeploy_InvalidPayload(t *testing.T) { cleanup := setupTestEnv(t) defer cleanup() - nm, _ := NewNodeManager() + dir := t.TempDir() + nm, _ := NewNodeManagerWithPaths( + filepath.Join(dir, "private.key"), + filepath.Join(dir, "node.json"), + ) nm.GenerateIdentity("test", RoleWorker) pr, _ := NewPeerRegistryWithPath(t.TempDir() + "/peers.json") transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() identity := nm.GetIdentity() // Create a message with invalid payload @@ -1284,6 +1389,7 @@ func TestWorker_HandleGetStats_NoIdentity(t *testing.T) { pr, _ := NewPeerRegistryWithPath(t.TempDir() + "/peers.json") transport := NewTransport(nm, pr, DefaultTransportConfig()) worker := NewWorker(nm, transport) + worker.DataDir = t.TempDir() msg, _ := NewMessage(MsgGetStats, "sender-id", "target-id", nil) _, err := worker.handleGetStats(msg) diff --git a/ueps/packet.go b/ueps/packet.go index 7c75334..c8a481e 100644 --- a/ueps/packet.go +++ b/ueps/packet.go @@ -90,35 +90,36 @@ func (p *PacketBuilder) MarshalAndSign(sharedSecret []byte) ([]byte, error) { } // 4. Write Payload TLV (0xFF) - // Note: 0xFF length is variable. For simplicity in this specialized reader, - // we might handle 0xFF as "read until EOF" or use a varint length. - // Implementing standard 1-byte length for payload is risky if payload > 255. - // Assuming your spec allows >255 bytes, we handle 0xFF differently. - - buf.WriteByte(TagPayload) - // We don't write a 1-byte length for payload here assuming stream mode, - // but if strict TLV, we'd need a multi-byte length protocol. - // For this snippet, simply appending data: - buf.Write(p.Payload) + // Fixed: Now uses writeTLV which provides a 2-byte length prefix. + // This prevents the io.ReadAll DoS and allows multiple packets in a stream. + if err := writeTLV(buf, TagPayload, p.Payload); err != nil { + return nil, err + } return buf.Bytes(), nil } -// Helper to write a simple TLV +// Helper to write a simple TLV. +// Now uses 2-byte big-endian length (uint16) to support up to 64KB payloads. func writeTLV(w io.Writer, tag uint8, value []byte) error { - // Check strict length constraint (1 byte length = max 255 bytes) - if len(value) > 255 { - return errors.New("TLV value too large for 1-byte length header") + // Check length constraint (2 byte length = max 65535 bytes) + if len(value) > 65535 { + return errors.New("TLV value too large for 2-byte length header") } if _, err := w.Write([]byte{tag}); err != nil { return err } - if _, err := w.Write([]byte{uint8(len(value))}); err != nil { + + lenBuf := make([]byte, 2) + binary.BigEndian.PutUint16(lenBuf, uint16(len(value))) + if _, err := w.Write(lenBuf); err != nil { return err } + if _, err := w.Write(value); err != nil { return err } return nil } + diff --git a/ueps/packet_coverage_test.go b/ueps/packet_coverage_test.go index 5f0adaa..6e1595c 100644 --- a/ueps/packet_coverage_test.go +++ b/ueps/packet_coverage_test.go @@ -112,26 +112,30 @@ func TestReadAndVerify_PayloadReadError(t *testing.T) { assert.Equal(t, "connection reset", err.Error()) } -// TestReadAndVerify_PayloadReadError_EOF ensures that a clean EOF -// (no payload bytes at all after 0xFF) is handled differently from -// a hard I/O error — io.ReadAll treats io.EOF as success and returns -// an empty slice, so the result should be an HMAC mismatch rather -// than a raw read error. +// TestReadAndVerify_PayloadReadError_EOF ensures that a truncated payload +// (missing bytes after TagPayload) is handled as an I/O error (UnexpectedEOF) +// because ReadAndVerify now uses io.ReadFull with the expected length prefix. func TestReadAndVerify_PayloadReadError_EOF(t *testing.T) { payload := []byte("eof test") builder := NewBuilder(0x20, payload) frame, err := builder.MarshalAndSign(testSecret) require.NoError(t, err) - // Truncate at 0xFF tag — the reader will see 0xFF then immediate - // EOF, which io.ReadAll treats as success with empty payload. + // Truncate at TagPayload tag + partial length — the reader will see 0xFF + // then EOF while trying to read the 2-byte length or the payload itself. payloadTagIdx := bytes.IndexByte(frame, TagPayload) require.NotEqual(t, -1, payloadTagIdx) - truncated := frame[:payloadTagIdx+1] + truncated := frame[:payloadTagIdx+1] // Only the tag, no length _, err = ReadAndVerify(bufio.NewReader(bytes.NewReader(truncated)), testSecret) require.Error(t, err) - assert.Contains(t, err.Error(), "integrity violation") + assert.ErrorIs(t, err, io.EOF) // Failed reading length + + truncatedWithLen := frame[:payloadTagIdx+3] // Tag + Length, but no payload + _, err = ReadAndVerify(bufio.NewReader(bytes.NewReader(truncatedWithLen)), testSecret) + require.Error(t, err) + // io.ReadFull returns io.EOF if no bytes are read at all before EOF. + assert.ErrorIs(t, err, io.EOF) } // TestWriteTLV_AllWritesSucceed confirms the happy path still works @@ -141,9 +145,11 @@ func TestWriteTLV_AllWritesSucceed(t *testing.T) { var buf bytes.Buffer err := writeTLV(&buf, TagVersion, []byte{0x09}) require.NoError(t, err) - assert.Equal(t, []byte{TagVersion, 0x01, 0x09}, buf.Bytes()) + // Now uses 2-byte big-endian length: 0x00 0x01 + assert.Equal(t, []byte{TagVersion, 0x00, 0x01, 0x09}, buf.Bytes()) } + // TestWriteTLV_FailWriterTable runs the three failure scenarios in // a table-driven fashion for completeness. func TestWriteTLV_FailWriterTable(t *testing.T) {