added semaphores to implement limits for max number of network connections

Author: caffix

engine/plugins/service_discovery/http_probes/plugin.go

@@ -105,7 +105,11 @@ func (hp *httpProbing) query(e *et.Event, entity *dbt.Entity, target string, por
 	ctx, cancel := context.WithTimeout(e.Session.Ctx(), 5*time.Second)
 	defer cancel()
 
-	if resp, err := http.RequestWebPage(ctx, &http.Request{URL: target}); err == nil && resp != nil {
+	e.Session.NetSem().Acquire()
+	resp, err := http.RequestWebPage(ctx, &http.Request{URL: target})
+	e.Session.NetSem().Release()
+
+	if err == nil && resp != nil {
 		findings = append(findings, hp.store(e, resp, entity, port)...)
 	}
 	return findings

engine/sessions/manager.go

@@ -41,7 +41,7 @@ func NewManager(l *slog.Logger, reg et.Registry) et.SessionManager {
 }
 
 func (r *manager) NewSession(cfg *config.Config) (et.Session, error) {
-	s, err := CreateSession(r.registry, cfg)
+	s, err := CreateSession(r, r.registry, cfg)
 	if err == nil {
 		err = r.AddSession(s)
 
@@ -120,6 +120,13 @@ func (r *manager) GetSession(id uuid.UUID) et.Session {
 	return nil
 }
 
+func (r *manager) NumOfSessions() int {
+	r.RLock()
+	defer r.RUnlock()
+
+	return len(r.sessions)
+}
+
 // Shutdown: cleans all sessions from a session storage and shutdown the session storage.
 func (r *manager) Shutdown() {
 	var list []uuid.UUID

engine/sessions/session.go

@@ -12,13 +12,15 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/owasp-amass/amass/v5/config"
 	"github.com/owasp-amass/amass/v5/engine/pubsub"
 	"github.com/owasp-amass/amass/v5/engine/sessions/scope"
 	et "github.com/owasp-amass/amass/v5/engine/types"
+	amassnet "github.com/owasp-amass/amass/v5/internal/net"
 	assetdb "github.com/owasp-amass/asset-db"
 	"github.com/owasp-amass/asset-db/repository"
 	"github.com/owasp-amass/asset-db/repository/neo4j"
@@ -29,47 +31,54 @@ import (
 )
 
 type Session struct {
-	id        uuid.UUID
-	ctx       context.Context
-	cancel    context.CancelFunc
-	log       *slog.Logger
-	ps        *pubsub.Logger
-	cfg       *config.Config
-	scope     et.Scope
-	start     time.Time
-	db        repository.Repository
-	backlog   *sessionBacklog
-	pipelines et.SessionPipelines
-	dsn       string
-	dbtype    string
-	ranger    cidranger.Ranger
-	tmpdir    string
-	stats     *et.SessionStats
-	done      chan struct{}
-	finished  bool
+	id           uuid.UUID
+	mgr          *manager
+	ctx          context.Context
+	cancel       context.CancelFunc
+	log          *slog.Logger
+	ps           *pubsub.Logger
+	cfg          *config.Config
+	scope        et.Scope
+	start        time.Time
+	db           repository.Repository
+	backlog      *sessionBacklog
+	pipelines    et.SessionPipelines
+	dsn          string
+	dbtype       string
+	ranger       cidranger.Ranger
+	tmpdir       string
+	stats        *et.SessionStats
+	done         chan struct{}
+	finished     bool
+	numOfSess    int
+	netSemaphore *sessSemaphore
 }
 
 // CreateSession initializes a new Session object based on the provided configuration.
 // The session object represents the state of an active engine enumeration.
-func CreateSession(reg et.Registry, cfg *config.Config) (et.Session, error) {
+func CreateSession(mgr *manager, reg et.Registry, cfg *config.Config) (et.Session, error) {
 	// Use default configuration if none is provided
 	if cfg == nil {
 		cfg = config.NewConfig()
 	}
 
 	startTime := time.Now()
+	numOfSessions := mgr.NumOfSessions() + 1
 	ctx, cancel := context.WithCancel(context.Background())
 	// Create a new session object
 	s := &Session{
-		id:     uuid.New(),
-		ctx:    ctx,
-		cancel: cancel,
-		cfg:    cfg,
-		start:  startTime,
-		ranger: NewAmassRanger(),
-		ps:     pubsub.NewLogger(),
-		stats:  new(et.SessionStats),
-		done:   make(chan struct{}),
+		id:           uuid.New(),
+		mgr:          mgr,
+		ctx:          ctx,
+		cancel:       cancel,
+		cfg:          cfg,
+		start:        startTime,
+		ranger:       NewAmassRanger(),
+		ps:           pubsub.NewLogger(),
+		stats:        new(et.SessionStats),
+		done:         make(chan struct{}),
+		numOfSess:    numOfSessions,
+		netSemaphore: NewSessSemaphore(amassnet.MaxNetworkConns / numOfSessions),
 	}
 	s.scope = scope.CreateFromConfigScope(s)
 	s.log = slog.New(slog.NewJSONHandler(s.ps, nil)).With("session", s.id)
@@ -102,6 +111,7 @@ func CreateSession(reg et.Registry, cfg *config.Config) (et.Session, error) {
 	s.log.Info("Temporary directory created", slog.String("dir", s.tmpdir))
 	s.log.Info("Database connection established", slog.String("dsn", s.dsn))
 	go s.updateStats()
+	go s.updateSessionSemaphore()
 	return s, nil
 }
 
@@ -121,6 +131,10 @@ func (s *Session) PubSub() *pubsub.Logger {
 	return s.ps
 }
 
+func (s *Session) NetSem() et.SessionSemaphone {
+	return s.netSemaphore
+}
+
 func (s *Session) Config() *config.Config {
 	return s.cfg
 }
@@ -298,3 +312,66 @@ func (s *Session) calculateStats() {
 	ss.WorkItemsCompleted = completed
 	ss.Unlock()
 }
+
+type sessSemaphore struct {
+	sync.Mutex
+	sem amassnet.Semaphore
+}
+
+func NewSessSemaphore(limit int) *sessSemaphore {
+	return &sessSemaphore{sem: amassnet.NewSemaphore(limit)}
+}
+
+func (ss *sessSemaphore) Acquire() {
+	ss.Lock()
+	defer ss.Unlock()
+
+	ss.sem.Acquire()
+}
+
+func (ss *sessSemaphore) Release() {
+	ss.Lock()
+	defer ss.Unlock()
+
+	ss.sem.Release()
+}
+
+func (s *Session) updateSessionSemaphore() {
+	tick := time.NewTicker(10 * time.Second)
+	defer tick.Stop()
+
+	for {
+		select {
+		case <-s.done:
+			return
+		case <-tick.C:
+			if num := s.mgr.NumOfSessions(); num != s.numOfSess {
+				s.numOfSess = num
+				s.buildNewSessionSemaphore()
+			}
+		}
+	}
+}
+
+func (s *Session) buildNewSessionSemaphore() {
+	s.netSemaphore.Lock()
+	defer s.netSemaphore.Unlock()
+
+	limit := amassnet.MaxNetworkConns / s.numOfSess
+	sem := amassnet.NewSemaphore(limit)
+loop:
+	for range limit {
+		select {
+		case s.netSemaphore.sem <- struct{}{}:
+			select {
+			case <-sem:
+			default:
+				break loop
+			}
+		default:
+			break loop
+		}
+	}
+
+	s.netSemaphore.sem = sem
+}

engine/types/sessions.go

@@ -26,11 +26,17 @@ import (
 
 type SessionPipelines map[oam.AssetType]*AssetPipeline
 
+type SessionSemaphone interface {
+	Acquire()
+	Release()
+}
+
 type Session interface {
 	ID() uuid.UUID
 	Ctx() context.Context
 	Log() *slog.Logger
 	PubSub() *pubsub.Logger
+	NetSem() SessionSemaphone
 	Config() *config.Config
 	Scope() Scope
 	StartTime() time.Time

internal/net/http/http.go

@@ -201,7 +201,9 @@ func RequestWebPage(ctx context.Context, r *Request) (*Response, error) {
 		}
 	}
 
+	amassnet.MaxNetConnSem.Acquire()
 	resp, err := DefaultClient.Do(req)
+	amassnet.MaxNetConnSem.Release()
 	if err != nil {
 		return nil, err
 	}

internal/net/network.go

@@ -14,6 +14,9 @@ import (
 	"time"
 )
 
+// Maximum number of connections allowed by each process.
+const MaxNetworkConns = 500
+
 // IPv4RE is a regular expression that will match an IPv4 address.
 const IPv4RE = "((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)[.]){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)"
 
@@ -48,12 +51,40 @@ var ReservedCIDRs = []string{
 // The reserved network address ranges
 var reservedAddrRanges []*net.IPNet
 
+type Semaphore chan struct{}
+
+var MaxNetConnSem Semaphore
+
 func init() {
 	for _, cidr := range ReservedCIDRs {
 		if _, ipnet, err := net.ParseCIDR(cidr); err == nil {
 			reservedAddrRanges = append(reservedAddrRanges, ipnet)
 		}
 	}
+
+	// setup the semaphore for limiting max network connections
+	MaxNetConnSem = NewSemaphore(MaxNetworkConns)
+}
+
+func NewSemaphore(cap int) Semaphore {
+	sem := make(Semaphore, cap)
+
+	for range cap {
+		sem <- struct{}{}
+	}
+
+	return sem
+}
+
+func (r Semaphore) Acquire() {
+	<-r
+}
+
+func (r Semaphore) Release() {
+	select {
+	case r <- struct{}{}:
+	default:
+	}
 }
 
 // DialContext performs the dial using global variables (e.g. LocalAddr).