graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #22 from towerops-app/security-updates

Browse source 
Security fixes
This commit is contained in:
Graham McIntire 2026-02-12 11:05:43 -06:00 • committed by GitHub
commit cd6fb4e79e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 495 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
agent.go
View file

@ -4,6 +4,7 @@ import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"log/slog"
"math/rand/v2"
@ -22,6 +23,11 @@ import (
var osExit = os.Exit
var doSelfUpdate = selfUpdate
var errRestartRequested = fmt.Errorf("restart requested")
var joinTimeout = 10 * time.Second
const maxJobPayloadBytes = 4 << 20 // 4 MB — well above any legitimate job list
// channelMsg is the WebSocket channel message format (JSON wrapper around binary protobuf).
type channelMsg struct {
Topic string `json:"topic"`
@ -47,6 +53,10 @@ func runAgent(ctx context.Context, wsURL, token string) {
if ctx.Err() != nil {
return
}
if errors.Is(err, errRestartRequested) {
osExit(0)
return
}
if err != nil {
slog.Error("agent disconnected", "error", err)
}
@ -135,12 +145,28 @@ func runSession(ctx context.Context, baseURL, token string) error {
return
}
encoded := base64.StdEncoding.EncodeToString(bin)
*bp = bin[:0]
full := (*bp)[:cap(*bp)]
zeroBytes(full)
*bp = full[:0]
bufPool.Put(bp)
payload, _ := json.Marshal(map[string]string{"binary": encoded})
sendMsg(event, payload)
}
// Reader goroutine — must start before join so we can receive the reply
msgCh := make(chan []byte, 100)
errCh := make(chan error, 1)
go func() {
for {
data, _, err := ws.ReadMessage()
if err != nil {
errCh <- err
return
}
msgCh <- data
}
}()
// Join channel
joinPayload, _ := json.Marshal(map[string]string{"token": token})
joinMsg := channelMsg{
@ -155,33 +181,47 @@ func runSession(ctx context.Context, baseURL, token string) error {
}
slog.Debug("sent channel join request")
// Wait for join reply before entering main loop
select {
case data := <-msgCh:
var reply channelMsg
if err := json.Unmarshal(data, &reply); err != nil {
return fmt.Errorf("join reply unmarshal: %w", err)
}
if reply.Event == "phx_reply" {
var status struct {
Status string `json:"status"`
Response any `json:"response"`
}
if err := json.Unmarshal(reply.Payload, &status); err == nil && status.Status != "ok" {
return fmt.Errorf("join rejected: %s", status.Status)
}
}
slog.Info("channel joined")
case err := <-errCh:
return fmt.Errorf("read during join: %w", err)
case <-time.After(joinTimeout):
return fmt.Errorf("join timeout")
}
// Writer goroutine - serializes all writes to the WebSocket
var writerWg sync.WaitGroup
writeErrCh := make(chan error, 1)
writerWg.Add(1)
go func() {
defer writerWg.Done()
for data := range writeCh {
if err := ws.WriteText(data); err != nil {
slog.Error("websocket write", "error", err)
select {
case writeErrCh <- err:
default:
}
return
}
}
}()
// Reader goroutine - reads messages and dispatches
msgCh := make(chan []byte, 100)
errCh := make(chan error, 1)
go func() {
for {
data, _, err := ws.ReadMessage()
if err != nil {
errCh <- err
return
}
msgCh <- data
}
}()
heartbeatTicker := time.NewTicker(60 * time.Second)
defer heartbeatTicker.Stop()
channelHeartbeatTicker := time.NewTicker(25 * time.Second)
@ -222,13 +262,20 @@ func runSession(ctx context.Context, baseURL, token string) error {
flushSnmpBatch()
return fmt.Errorf("read: %w", err)
case err := <-writeErrCh:
flushSnmpBatch()
return fmt.Errorf("write: %w", err)
case data := <-msgCh:
var msg channelMsg
if err := json.Unmarshal(data, &msg); err != nil {
slog.Warn("invalid message", "error", err)
continue
}
handleMessage(ctx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
if handleMessage(ctx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh) {
flushSnmpBatch()
return errRestartRequested
}
case result := <-snmpResultCh:
snmpBatch = append(snmpBatch, result)
@ -279,6 +326,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
}
// handleMessage dispatches incoming channel messages.
// Returns true if the session should end (e.g. restart requested).
func handleMessage(
ctx context.Context,
msg channelMsg,
@ -287,7 +335,7 @@ func handleMessage(
mikrotikResultCh chan<- *pb.MikrotikResult,
credTestResultCh chan<- *pb.CredentialTestResult,
monitoringCheckCh chan<- *pb.MonitoringCheck,
) {
) bool {
switch msg.Event {
case "phx_reply":
slog.Debug("channel reply", "topic", msg.Topic)
@ -298,17 +346,21 @@ func handleMessage(
}
if err := json.Unmarshal(msg.Payload, &payload); err != nil {
slog.Error("decode job payload", "error", err)
return
return false
}
if len(payload.Binary) > maxJobPayloadBytes {
slog.Error("job payload too large", "size", len(payload.Binary), "max", maxJobPayloadBytes)
return false
}
bin, err := base64.StdEncoding.DecodeString(payload.Binary)
if err != nil {
slog.Error("decode base64", "error", err)
return
return false
}
var jobList pb.AgentJobList
if err := proto.Unmarshal(bin, &jobList); err != nil {
slog.Error("unmarshal job list", "error", err)
return
return false
}
slog.Info("received jobs", "count", len(jobList.Jobs))
for _, job := range jobList.Jobs {
@ -316,8 +368,8 @@ func handleMessage(
}
case "restart":
slog.Info("restart requested by server, exiting")
osExit(0)
slog.Info("restart requested by server")
return true
case "update":
var payload struct {
@ -326,7 +378,7 @@ func handleMessage(
}
if err := json.Unmarshal(msg.Payload, &payload); err != nil || payload.URL == "" || payload.Checksum == "" {
slog.Error("invalid update payload")
return
return false
}
slog.Info("update requested", "url", payload.URL)
if err := doSelfUpdate(payload.URL, payload.Checksum); err != nil {
@ -336,6 +388,7 @@ func handleMessage(
default:
slog.Debug("ignoring event", "event", msg.Event)
}
return false
}
// jobPools holds the worker pools for each job type.
@ -357,15 +410,19 @@ func dispatchJob(
) {
slog.Info("starting job", "job_id", job.JobId, "type", job.JobType)
var ok bool
switch job.JobType {
case pb.JobType_MIKROTIK:
pools.mikrotik.submit(func() { executeMikrotikJob(ctx, job, mikrotikResultCh) })
ok = pools.mikrotik.submit(ctx, func() { executeMikrotikJob(ctx, job, mikrotikResultCh) })
case pb.JobType_TEST_CREDENTIALS:
pools.snmp.submit(func() { executeCredentialTest(ctx, job, credTestResultCh) })
ok = pools.snmp.submit(ctx, func() { executeCredentialTest(ctx, job, credTestResultCh) })
case pb.JobType_PING:
pools.ping.submit(func() { executePingJob(ctx, job, monitoringCheckCh) })
ok = pools.ping.submit(ctx, func() { executePingJob(ctx, job, monitoringCheckCh) })
default:
pools.snmp.submit(func() { executeSnmpJob(ctx, job, snmpResultCh) })
ok = pools.snmp.submit(ctx, func() { executeSnmpJob(ctx, job, snmpResultCh) })
}
if !ok {
slog.Warn("job dropped, pool full", "job_id", job.JobId)
}
}

201
agent_test.go
View file

@ -5,6 +5,9 @@ import (
"encoding/base64"
"encoding/json"
"fmt"
"net"
"strings"
"sync"
"testing"
"time"
@ -165,20 +168,14 @@ func TestHandleMessage(t *testing.T) {
})
t.Run("restart", func(t *testing.T) {
origExit := osExit
defer func() { osExit = origExit }()
var exitCode int
osExit = func(code int) { exitCode = code }
snmpCh := make(chan *pb.SnmpResult, 1)
mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
shouldEnd := handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
if exitCode != 0 {
t.Errorf("expected exit code 0, got %d", exitCode)
if !shouldEnd {
t.Error("expected handleMessage to return true for restart")
}
})
@ -336,6 +333,57 @@ func TestHandleMessage(t *testing.T) {
})
}
func TestHandleMessageRejectsOversizedPayload(t *testing.T) {
snmpCh := make(chan *pb.SnmpResult, 1)
mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
// Create a binary payload larger than maxJobPayloadBytes
oversized := make([]byte, maxJobPayloadBytes+1)
encoded := base64.StdEncoding.EncodeToString(oversized)
payload, _ := json.Marshal(map[string]string{"binary": encoded})
handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Verify no jobs were dispatched
select {
case <-snmpCh:
t.Error("expected no SNMP result for oversized payload")
case <-time.After(100 * time.Millisecond):
// Good — nothing dispatched
}
}
func TestBufferPoolZeroesOnReturn(t *testing.T) {
pool := &sync.Pool{
New: func() any {
b := make([]byte, 0, 64)
return &b
},
}
// Get a buffer, write some data, return it
bp := pool.Get().(*[]byte)
*bp = append((*bp)[:0], []byte("sensitive credentials data here")...)
full := (*bp)[:cap(*bp)]
// Simulate the zeroing that sendBinaryResult should do
zeroBytes(full)
*bp = full[:0]
pool.Put(bp)
// Get the buffer back and verify it's zeroed
bp2 := pool.Get().(*[]byte)
full2 := (*bp2)[:cap(*bp2)]
for i, b := range full2 {
if b != 0 {
t.Errorf("byte[%d] = %d, expected 0 — pool buffer not zeroed", i, b)
break
}
}
}
func TestZeroBytes(t *testing.T) {
b := []byte{1, 2, 3, 4, 5}
zeroBytes(b)
@ -496,3 +544,138 @@ func TestDispatchJob(t *testing.T) {
}
})
}
func TestRunSessionRejectsFailedJoin(t *testing.T) {
origTimeout := joinTimeout
defer func() { joinTimeout = origTimeout }()
joinTimeout = 2 * time.Second
// Start a fake WebSocket server that accepts the upgrade then sends a phx_error join reply
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer func() { _ = ln.Close() }()
go func() {
conn, err := ln.Accept()
if err != nil {
return
}
defer func() { _ = conn.Close() }()
// Read HTTP upgrade request
buf := make([]byte, 4096)
n, _ := conn.Read(buf)
reqStr := string(buf[:n])
// Extract key and compute accept
key := extractWSKey(reqStr)
accept := computeAcceptKey(key)
// Send valid 101 upgrade
resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
_, _ = conn.Write([]byte(resp))
// Read the join message (masked WebSocket frame) — just consume it
frameBuf := make([]byte, 4096)
_, _ = conn.Read(frameBuf)
// Send phx_error reply as an unmasked text frame
reply, _ := json.Marshal(channelMsg{
Topic: "agent:agent-0",
Event: "phx_reply",
Payload: json.RawMessage(`{"status":"error","response":{"reason":"invalid token"}}`),
Ref: strPtr("1"),
})
frame := makeTextFrame(reply)
_, _ = conn.Write(frame)
// Keep connection open for a bit
time.Sleep(time.Second)
}()
addr := ln.Addr().String()
err = runSession(context.Background(), "ws://"+addr, "bad-token")
if err == nil {
t.Fatal("expected error from rejected join")
}
if !strings.Contains(err.Error(), "join rejected") {
t.Errorf("expected 'join rejected' in error, got: %v", err)
}
}
func TestRunSessionJoinTimeout(t *testing.T) {
origTimeout := joinTimeout
defer func() { joinTimeout = origTimeout }()
joinTimeout = 500 * time.Millisecond
// Server that upgrades but never sends a join reply
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer func() { _ = ln.Close() }()
go func() {
conn, err := ln.Accept()
if err != nil {
return
}
defer func() { _ = conn.Close() }()
buf := make([]byte, 4096)
n, _ := conn.Read(buf)
reqStr := string(buf[:n])
key := extractWSKey(reqStr)
accept := computeAcceptKey(key)
resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
_, _ = conn.Write([]byte(resp))
// Read join frame but never reply
frameBuf := make([]byte, 4096)
_, _ = conn.Read(frameBuf)
time.Sleep(5 * time.Second)
}()
addr := ln.Addr().String()
start := time.Now()
err = runSession(context.Background(), "ws://"+addr, "token")
elapsed := time.Since(start)
if err == nil {
t.Fatal("expected timeout error")
}
if !strings.Contains(err.Error(), "join timeout") {
t.Errorf("expected 'join timeout' in error, got: %v", err)
}
if elapsed > 3*time.Second {
t.Errorf("took too long (%v), timeout didn't trigger", elapsed)
}
}
// extractWSKey extracts the Sec-WebSocket-Key from a raw HTTP request.
func extractWSKey(req string) string {
for _, line := range strings.Split(req, "\r\n") {
lower := strings.ToLower(line)
if strings.HasPrefix(lower, "sec-websocket-key: ") {
return strings.TrimSpace(line[len("Sec-WebSocket-Key: "):])
}
}
return ""
}
// makeTextFrame creates an unmasked WebSocket text frame (server→client).
func makeTextFrame(payload []byte) []byte {
length := len(payload)
var frame []byte
frame = append(frame, 0x81) // FIN + text
if length <= 125 {
frame = append(frame, byte(length))
} else if length <= 65535 {
frame = append(frame, 126, byte(length>>8), byte(length))
}
frame = append(frame, payload...)
return frame
}

7
ssh.go
View file

@ -4,6 +4,8 @@ import (
"context"
"fmt"
"log/slog"
"net"
"strconv"
"time"
"github.com/towerops-app/towerops-agent/pb"
@ -11,6 +13,7 @@ import (
)
var sshBackup = executeMikrotikBackup
var sshDial = ssh.Dial
var doPing = pingDevice
// executeMikrotikBackup connects via SSH and runs /export compact.
@ -26,8 +29,8 @@ func executeMikrotikBackup(ip string, port uint16, username, password string) (s
Timeout: 30 * time.Second,
}
addr := fmt.Sprintf("%s:%d", ip, port)
conn, err := ssh.Dial("tcp", addr, config)
addr := net.JoinHostPort(ip, strconv.Itoa(int(port)))
conn, err := sshDial("tcp", addr, config)
if err != nil {
return "", fmt.Errorf("ssh dial %s: %w", addr, err)
}

34
ssh_test.go
View file

@ -270,6 +270,40 @@ func TestExecuteMikrotikBackupViaSSH(t *testing.T) {
})
}
func TestSSHBackupIPv6Address(t *testing.T) {
origDial := sshDial
defer func() { sshDial = origDial }()
var capturedAddr string
sshDial = func(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
capturedAddr = addr
return nil, fmt.Errorf("mock dial")
}
_, _ = executeMikrotikBackup("::1", 22, "admin", "pass")
if capturedAddr != "[::1]:22" {
t.Errorf("expected [::1]:22, got %q", capturedAddr)
}
}
func TestSSHBackupIPv4Address(t *testing.T) {
origDial := sshDial
defer func() { sshDial = origDial }()
var capturedAddr string
sshDial = func(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
capturedAddr = addr
return nil, fmt.Errorf("mock dial")
}
_, _ = executeMikrotikBackup("10.0.0.1", 22, "admin", "pass")
if capturedAddr != "10.0.0.1:22" {
t.Errorf("expected 10.0.0.1:22, got %q", capturedAddr)
}
}
func TestExecuteMikrotikBackupDialError(t *testing.T) {
_, err := executeMikrotikBackup("127.0.0.1", 1, "admin", "pass")
if err == nil {

17
update.go
View file

@ -9,6 +9,7 @@ import (
"net/http"
"net/url"
"os"
"strings"
"syscall"
)
@ -78,6 +79,20 @@ func selfUpdate(downloadURL, expectedChecksum string) error {
slog.Info("binary replaced", "path", currentExe)
// Re-exec with same arguments
slog.Info("re-executing", "args", os.Args)
slog.Info("re-executing", "args", sanitizeArgs(os.Args))
return syscall.Exec(currentExe, os.Args, os.Environ())
}
// sanitizeArgs returns a copy of args with token values masked.
func sanitizeArgs(args []string) []string {
out := make([]string, len(args))
copy(out, args)
for i, a := range out {
if (a == "--token" || a == "-token") && i+1 < len(out) {
out[i+1] = "***"
} else if strings.HasPrefix(a, "--token=") || strings.HasPrefix(a, "-token=") {
out[i] = a[:strings.Index(a, "=")+1] + "***"
}
}
return out
}

39
update_test.go
View file

@ -283,6 +283,45 @@ func TestSelfUpdateTooLarge(t *testing.T) {
}
}
func TestSanitizeArgs(t *testing.T) {
tests := []struct {
name string
in []string
want []string
}{
{"flag with separate value", []string{"agent", "--token", "secret"}, []string{"agent", "--token", "***"}},
{"no sensitive flags", []string{"agent", "--api-url", "wss://x"}, []string{"agent", "--api-url", "wss://x"}},
{"equals syntax", []string{"agent", "--token=secret"}, []string{"agent", "--token=***"}},
{"empty args", []string{"agent"}, []string{"agent"}},
{"trailing flag no value", []string{"agent", "--token"}, []string{"agent", "--token"}},
{"short flag with value", []string{"agent", "-token", "secret"}, []string{"agent", "-token", "***"}},
{"short equals syntax", []string{"agent", "-token=secret"}, []string{"agent", "-token=***"}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := sanitizeArgs(tt.in)
if len(got) != len(tt.want) {
t.Fatalf("length: got %d, want %d", len(got), len(tt.want))
}
for i := range got {
if got[i] != tt.want[i] {
t.Errorf("arg[%d]: got %q, want %q", i, got[i], tt.want[i])
}
}
})
}
// Verify original slice is not mutated
t.Run("does not mutate input", func(t *testing.T) {
orig := []string{"agent", "--token", "secret"}
_ = sanitizeArgs(orig)
if orig[2] != "secret" {
t.Error("sanitizeArgs mutated the input slice")
}
})
}
// rewriteToHTTPS converts an httptest TLS server URL to use the https scheme.
// httptest.NewTLSServer returns URLs with https:// already, but this ensures consistency.
func rewriteToHTTPS(rawURL string) string {

9
websocket.go
View file

@ -13,6 +13,7 @@ import (
"net/url"
"strings"
"sync"
"time"
)
// websocketGUID is the magic GUID from RFC 6455 Section 4.2.2.
@ -42,6 +43,8 @@ type WSConn struct {
mu sync.Mutex // serializes writes
}
var wsHandshakeTimeout = 30 * time.Second
var randRead = rand.Read
var netDial = net.Dial
var tlsDial = func(network, addr string) (net.Conn, error) {
@ -75,6 +78,9 @@ func WSDial(rawURL string) (*WSConn, error) {
return nil, fmt.Errorf("dial %s: %w", host, err)
}
// Set handshake deadline — prevents a slow/malicious server from blocking indefinitely
_ = conn.SetDeadline(time.Now().Add(wsHandshakeTimeout))
// Generate random key for Sec-WebSocket-Key
keyBytes := make([]byte, 16)
if _, err := randRead(keyBytes); err != nil {
@ -124,6 +130,9 @@ func WSDial(rawURL string) (*WSConn, error) {
return nil, fmt.Errorf("missing Sec-WebSocket-Accept header")
}
// Clear handshake deadline — normal operation uses no deadline
_ = conn.SetDeadline(time.Time{})
return &WSConn{conn: conn, reader: bufio.NewReaderSize(conn, 8192)}, nil
}

39
websocket_test.go
View file

@ -10,6 +10,7 @@ import (
"net/http"
"strings"
"testing"
"time"
)
func TestWriteFrameMasked(t *testing.T) {
@ -558,6 +559,44 @@ func TestWSDialBadURL(t *testing.T) {
}
}
func TestWSDialHandshakeTimeout(t *testing.T) {
// Server accepts connection but never sends data — should trigger timeout
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer func() { _ = ln.Close() }()
go func() {
conn, err := ln.Accept()
if err != nil {
return
}
// Read the request but never respond
buf := make([]byte, 4096)
_, _ = conn.Read(buf)
// Hold connection open until test ends
<-make(chan struct{})
_ = conn.Close()
}()
origTimeout := wsHandshakeTimeout
defer func() { wsHandshakeTimeout = origTimeout }()
wsHandshakeTimeout = 500 * time.Millisecond // Short timeout for test
addr := ln.Addr().String()
start := time.Now()
_, err = WSDial("ws://" + addr + "/socket")
elapsed := time.Since(start)
if err == nil {
t.Error("expected timeout error")
}
if elapsed > 5*time.Second {
t.Errorf("took too long (%v), timeout didn't work", elapsed)
}
}
func TestWSDialDefaultPorts(t *testing.T) {
// Test that ws:// defaults to port 80 — will fail to connect but verifies URL parsing
_, err := WSDial("ws://127.0.0.1/path")

26
workerpool.go
View file

@ -1,6 +1,10 @@
package main
import "sync"
import (
"context"
"log/slog"
"sync"
)
// workerPool is a fixed-size goroutine pool for executing tasks.
type workerPool struct {
@ -19,16 +23,28 @@ func newWorkerPool(n int) *workerPool {
go func() {
defer p.wg.Done()
for fn := range p.tasks {
fn()
func() {
defer func() {
if r := recover(); r != nil {
slog.Error("worker panic recovered", "error", r)
}
}()
fn()
}()
}
}()
}
return p
}
// submit enqueues a task. Blocks if all workers are busy and the queue is full.
func (p *workerPool) submit(fn func()) {
p.tasks <- fn
// submit enqueues a task. Returns false if the context is cancelled before the task can be queued.
func (p *workerPool) submit(ctx context.Context, fn func()) bool {
select {
case p.tasks <- fn:
return true
case <-ctx.Done():
return false
}
}
// stop closes the task channel and waits for all workers to finish.

58
workerpool_test.go
View file

@ -1,6 +1,7 @@
package main
import (
"context"
"sync/atomic"
"testing"
"time"
@ -13,7 +14,7 @@ func TestWorkerPool(t *testing.T) {
var count atomic.Int32
for i := 0; i < 100; i++ {
pool.submit(func() {
pool.submit(context.Background(), func() {
count.Add(1)
})
}
@ -32,7 +33,7 @@ func TestWorkerPool(t *testing.T) {
var maxConcurrent atomic.Int32
for i := 0; i < 20; i++ {
pool.submit(func() {
pool.submit(context.Background(), func() {
cur := concurrent.Add(1)
for {
old := maxConcurrent.Load()
@ -57,3 +58,56 @@ func TestWorkerPool(t *testing.T) {
pool.stop() // should not panic
})
}
func TestWorkerPoolRecoversPanic(t *testing.T) {
pool := newWorkerPool(1)
defer pool.stop()
// Submit a function that panics
pool.submit(context.Background(), func() { panic("boom") })
// Give the panic time to be processed
time.Sleep(50 * time.Millisecond)
// Submit a normal function — the worker should still be alive
done := make(chan struct{})
ok := pool.submit(context.Background(), func() { close(done) })
if !ok {
t.Fatal("expected submit to succeed after panic recovery")
}
select {
case <-done:
// Worker survived the panic
case <-time.After(2 * time.Second):
t.Error("timed out — worker did not survive panic")
}
}
func TestWorkerPoolSubmitRespectsContext(t *testing.T) {
pool := newWorkerPool(1) // 1 worker, queue capacity 4
defer pool.stop()
blocker := make(chan struct{})
// Occupy the single worker
pool.submit(context.Background(), func() { <-blocker })
// Fill the buffered queue (capacity = n*4 = 4)
for i := 0; i < 4; i++ {
pool.submit(context.Background(), func() { <-blocker })
}
// Now the queue is full and the worker is busy.
// Submit with a cancelled context should return false immediately.
ctx, cancel := context.WithCancel(context.Background())
cancel()
ok := pool.submit(ctx, func() { t.Error("should not execute") })
if ok {
t.Error("expected submit to return false with cancelled context")
}
// Unblock everything for cleanup
close(blocker)
}