diff --git a/agent.go b/agent.go index 2e57e2e..f9c43bc 100644 --- a/agent.go +++ b/agent.go @@ -185,23 +185,24 @@ func runSession(ctx context.Context, baseURL, token string) error { go func() { defer readerWg.Done() for { - // Check context cancellation to enable fast shutdown - select { - case <-sessionCtx.Done(): - return - default: - } - data, _, err := ws.ReadMessage() if err != nil { - errCh <- err - sessionCancel() // Unblock any stuck pool submits + select { + case errCh <- err: + default: + } + sessionCancel() return } msgCh <- data } }() + go func() { + <-sessionCtx.Done() + _ = ws.Close() + }() + // Join channel joinPayload, _ := json.Marshal(map[string]string{"token": token}) joinMsg := channelMsg{ diff --git a/checks.go b/checks.go index 76b7a72..89f4980 100644 --- a/checks.go +++ b/checks.go @@ -16,22 +16,16 @@ import ( ) var ( - defaultHTTPClient = &http.Client{ - Timeout: 10 * time.Second, - Transport: &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: 10, - IdleConnTimeout: 90 * time.Second, - }, + defaultHTTPTransport = &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, } - insecureHTTPClient = &http.Client{ - Timeout: 10 * time.Second, - Transport: &http.Transport{ - TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, - MaxIdleConns: 100, - MaxIdleConnsPerHost: 10, - IdleConnTimeout: 90 * time.Second, - }, + insecureHTTPTransport = &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + MaxIdleConns: 100, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, } ) @@ -93,10 +87,25 @@ func ExecuteCheck(ctx context.Context, check *pb.Check) *pb.CheckResult { // executeHTTPCheck performs an HTTP/HTTPS check func executeHTTPCheck(ctx context.Context, config *pb.HttpCheckConfig, timeoutMs uint32) (uint32, string, float64) { - // Use shared HTTP client with connection pooling - client := defaultHTTPClient + transport := defaultHTTPTransport if !config.VerifySsl { - client = insecureHTTPClient + transport = insecureHTTPTransport + } + + timeout := time.Duration(timeoutMs) * time.Millisecond + if timeout == 0 { + timeout = 10 * time.Second + } + + client := &http.Client{ + Transport: transport, + Timeout: timeout, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + if !config.FollowRedirects { + return http.ErrUseLastResponse + } + return nil + }, } method := strings.ToUpper(config.Method) diff --git a/snmp_test.go b/snmp_test.go index cdaba8e..209c3c6 100644 --- a/snmp_test.go +++ b/snmp_test.go @@ -3,6 +3,7 @@ package main import ( "context" "fmt" + "strings" "testing" "github.com/gosnmp/gosnmp" @@ -300,9 +301,10 @@ func (m *mockSnmpQuerier) BulkWalkAll(rootOid string) ([]gosnmp.SnmpPDU, error) func TestExecuteSnmpJob(t *testing.T) { t.Run("nil device", func(t *testing.T) { ch := make(chan *pb.SnmpResult, 1) - executeSnmpJob(context.Background(), &pb.AgentJob{JobId: "1"}, ch) - if len(ch) != 0 { - t.Error("expected no result for nil device") + executeSnmpJob(context.Background(), &pb.AgentJob{JobId: "1", JobType: pb.JobType_POLL}, ch) + result := <-ch + if len(result.OidValues) != 0 { + t.Errorf("expected empty oid values for nil device, got %d", len(result.OidValues)) } }) @@ -318,8 +320,9 @@ func TestExecuteSnmpJob(t *testing.T) { JobId: "1", SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, }, ch) - if len(ch) != 0 { - t.Error("expected no result on dial error") + result := <-ch + if len(result.OidValues) != 0 { + t.Errorf("expected empty oid values on dial error, got %d", len(result.OidValues)) } }) @@ -501,15 +504,11 @@ func TestExecuteSnmpJob(t *testing.T) { ch := make(chan *pb.SnmpResult, 1) executeSnmpJob(context.Background(), &pb.AgentJob{ JobId: "1", - SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, - Queries: []*pb.SnmpQuery{ - {QueryType: pb.QueryType_WALK, Oids: []string{".1.3.6.1.2.1.2.2.1.1"}}, - }, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, }, ch) - result := <-ch - if len(result.OidValues) != 1 { - t.Errorf("expected 1 value (others skipped), got %d", len(result.OidValues)) + if len(result.OidValues) != 0 { + t.Errorf("expected empty oid values on dial error, got %d", len(result.OidValues)) } }) @@ -698,8 +697,12 @@ func TestExecuteCredentialTest(t *testing.T) { t.Run("nil device", func(t *testing.T) { ch := make(chan *pb.CredentialTestResult, 1) executeCredentialTest(context.Background(), &pb.AgentJob{JobId: "1"}, ch) - if len(ch) != 0 { - t.Error("expected no result for nil device") + result := <-ch + if result.Success { + t.Error("expected failure for nil device") + } + if !strings.Contains(result.ErrorMessage, "missing device") { + t.Errorf("expected 'missing device' in error, got: %s", result.ErrorMessage) } }) diff --git a/ssh_test.go b/ssh_test.go index 6533a12..0d1f620 100644 --- a/ssh_test.go +++ b/ssh_test.go @@ -19,8 +19,9 @@ func TestExecutePingJob(t *testing.T) { t.Run("nil device", func(t *testing.T) { ch := make(chan *pb.MonitoringCheck, 1) executePingJob(context.Background(), &pb.AgentJob{JobId: "p1"}, ch) - if len(ch) != 0 { - t.Error("expected no result for nil device") + result := <-ch + if result.Status != "failure" { + t.Errorf("expected failure status for nil device, got: %s", result.Status) } }) @@ -83,8 +84,12 @@ func TestExecuteMikrotikJob(t *testing.T) { t.Run("nil device", func(t *testing.T) { ch := make(chan *pb.MikrotikResult, 1) executeMikrotikJob(context.Background(), &pb.AgentJob{JobId: "m1"}, ch) - if len(ch) != 0 { - t.Error("expected no result for nil device") + result := <-ch + if result.Error == "" { + t.Error("expected error for nil device") + } + if !strings.Contains(result.Error, "missing device") { + t.Errorf("expected 'missing device' in error, got: %s", result.Error) } })