Merge pull request #36 from kubescape/feature/timeouts

Adding timeouts from connections
kubescape · Jan 25, 2024 · d2aa659 · d2aa659
2 parents 22233cf + ec281df
commit d2aa659
Show file tree

Hide file tree

Showing 7 changed files with 22 additions and 8 deletions.
diff --git a/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_cassandra.go b/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_cassandra.go
@@ -1,6 +1,8 @@
 package applicationlayerdiscovery
 
 import (
+	"time"
+
 	"github.com/gocql/gocql"
 	"github.com/kubescape/kubescape-network-scanner/internal/pkg/networkscanner/servicediscovery"
 )
@@ -52,6 +54,7 @@ func (d *CassandraDiscovery) Discover(sessionHandler servicediscovery.ISessionHa
 		Username: Username,
 		Password: Password,
 	}
+	cluster.Timeout = time.Second * 3
 
 	// Create a session
 	session, err := cluster.CreateSession()

diff --git a/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_etcd.go b/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_etcd.go
@@ -43,7 +43,7 @@ func (d *EtcdDiscovery) Discover(sessionHandler servicediscovery.ISessionHandler
 	endpoints := []string{fmt.Sprintf("%s:%d", sessionHandler.GetHost(), sessionHandler.GetPort())}
 	config := clientv3.Config{
 		Endpoints:   endpoints,
-		DialTimeout: 5 * time.Second,
+		DialTimeout: 3 * time.Second,
 	}
 
 	client, err := clientv3.New(config)

diff --git a/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_kubeapi.go b/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_kubeapi.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/kubescape/kubescape-network-scanner/internal/pkg/networkscanner/servicediscovery"
 )
@@ -49,7 +50,7 @@ func (d *KubeApiServerDiscovery) Discover(sessionHandler servicediscovery.ISessi
 	}
 
 	// Create an http.Client with the custom transport
-	client := &http.Client{Transport: tr}
+	client := &http.Client{Transport: tr, Timeout: time.Second * 3}
 
 	// Send a GET request to the Kubernetes API server
 	resp, err := client.Get(url)

diff --git a/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_mongodb.go b/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_mongodb.go
@@ -3,6 +3,7 @@ package applicationlayerdiscovery
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -42,6 +43,8 @@ func (d *MongoDBDiscovery) Protocol() string {
 
 func (d *MongoDBDiscovery) Discover(sessionHandler servicediscovery.ISessionHandler, presentationLayerDiscoveryResult servicediscovery.IPresentationDiscoveryResult) (servicediscovery.IApplicationDiscoveryResult, error) {
 	clientOptions := options.Client().ApplyURI(fmt.Sprintf("mongodb://%s:%d", sessionHandler.GetHost(), sessionHandler.GetPort()))
+	connectionTimeout := 3 * time.Second
+	clientOptions.Timeout = &connectionTimeout
 	ctx := context.Background()
 	client, err := mongo.Connect(ctx, clientOptions)
 	defer client.Disconnect(ctx)

diff --git a/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_postgresql.go b/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_postgresql.go
@@ -40,8 +40,8 @@ func (d *PostgresDiscovery) Protocol() string {
 }
 
 func (d *PostgresDiscovery) Discover(sessionHandler servicediscovery.ISessionHandler, presentationLayerDiscoveryResult servicediscovery.IPresentationDiscoveryResult) (servicediscovery.IApplicationDiscoveryResult, error) {
-	// Set a timeout of 50 ms
-	db, err := sql.Open("postgres", fmt.Sprintf("host=%s port=%d user=admin password=\"123456\" sslmode=disable connect_timeout=5", sessionHandler.GetHost(), sessionHandler.GetPort()))
+	// Set a timeout of 30 ms
+	db, err := sql.Open("postgres", fmt.Sprintf("host=%s port=%d user=admin password=\"123456\" sslmode=disable connect_timeout=3", sessionHandler.GetHost(), sessionHandler.GetPort()))
 	if err != nil {
 		fmt.Printf("failed to connect to PostgreSQL server: %v\n", err)
 		return &PostgresDiscoveryResult{

diff --git a/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_rabbitmq.go b/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_rabbitmq.go
@@ -2,6 +2,7 @@ package applicationlayerdiscovery
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/kubescape/kubescape-network-scanner/internal/pkg/networkscanner/servicediscovery"
 	"github.com/streadway/amqp"
@@ -42,7 +43,10 @@ func (d *RabbitMQDiscovery) Protocol() string {
 
 func (d *RabbitMQDiscovery) Discover(sessionHandler servicediscovery.ISessionHandler, presentationLayerDiscoveryResult servicediscovery.IPresentationDiscoveryResult) (servicediscovery.IApplicationDiscoveryResult, error) {
 	connectionString := fmt.Sprintf("amqp://%s:%d", sessionHandler.GetHost(), sessionHandler.GetPort())
-	conn, err := amqp.Dial(connectionString)
+	config := amqp.Config{
+		Dial: amqp.DefaultDial(time.Second * 3),
+	}
+	conn, err := amqp.DialConfig(connectionString, config)
 	if err != nil {
 		return &RabbitMQDiscoveryResult{
 			isDetected:      false,

diff --git a/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_redis.go b/internal/pkg/networkscanner/servicediscovery/applicationlayerdiscovery/al_redis.go
@@ -3,6 +3,7 @@ package applicationlayerdiscovery
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/go-redis/redis/v8"
 	"github.com/kubescape/kubescape-network-scanner/internal/pkg/networkscanner/servicediscovery"
@@ -48,9 +49,11 @@ func (e *RedisPingError) Error() string {
 func (d *RedisDiscovery) Discover(sessionHandler servicediscovery.ISessionHandler, presentationLayerDiscoveryResult servicediscovery.IPresentationDiscoveryResult) (servicediscovery.IApplicationDiscoveryResult, error) {
 
 	redisClient := redis.NewClient(&redis.Options{
-		Addr:     fmt.Sprintf("%s:%d", sessionHandler.GetHost(), sessionHandler.GetPort()),
-		Password: "", // No password for now, modify as needed
-		DB:       0,  // Use default DB
+		Addr:        fmt.Sprintf("%s:%d", sessionHandler.GetHost(), sessionHandler.GetPort()),
+		Password:    "", // No password for now, modify as needed
+		DB:          0,  // Use default DB
+		DialTimeout: 3 * time.Second,
+		MaxRetries:  1,
 	})
 
 	pong, err := redisClient.Ping(context.TODO()).Result()