-
Notifications
You must be signed in to change notification settings - Fork 2.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
997c5d1
commit 49ee80c
Showing
2 changed files
with
144 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,108 @@ | ||
// ports implements a open-port cache to avoid sending redundant requests to same port | ||
package ports | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
|
||
"github.com/Mzack9999/gcache" | ||
"github.com/projectdiscovery/fastdialer/fastdialer" | ||
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs" | ||
singleflight "github.com/projectdiscovery/utils/memoize/simpleflight" | ||
) | ||
|
||
var ( | ||
portsCacher *PortsCache | ||
ErrPortClosed = errors.New("port closed or filtered") | ||
) | ||
|
||
type PortStatus uint8 | ||
|
||
const ( | ||
Unknown PortStatus = iota | ||
Open | ||
Closed | ||
) | ||
|
||
// PortsCache is a cache for open ports | ||
type PortsCache struct { | ||
cache gcache.Cache[string, PortStatus] | ||
group singleflight.Group[string] | ||
dialer *fastdialer.Dialer | ||
} | ||
|
||
// NewPortsCache creates a new cache for open ports | ||
func NewPortsCache(dialer *fastdialer.Dialer, size int) *PortsCache { | ||
p := &PortsCache{group: singleflight.Group[string]{}, dialer: dialer} | ||
cache := gcache.New[string, PortStatus](size). | ||
LRU(). | ||
EvictedFunc(func(key string, value PortStatus) { | ||
p.group.Forget(key) | ||
}). | ||
Build() | ||
p.cache = cache | ||
return p | ||
} | ||
|
||
// Do performs a check for open ports | ||
func (p *PortsCache) Do(ctx context.Context, input *contextargs.Context) error { | ||
address := input.MetaInput.Address() | ||
if address == "" { | ||
// assume port is open is given info is not present/enough | ||
return nil | ||
} | ||
// check if it exists in cache | ||
if value, err := p.cache.GetIFPresent(address); !errors.Is(err, gcache.KeyNotFoundError) { | ||
switch value { | ||
case Closed: | ||
return ErrPortClosed | ||
default: | ||
return nil | ||
} | ||
} | ||
|
||
// if not in cache then check if it is open | ||
code, _, _ := p.group.Do(address, func() (interface{}, error) { | ||
conn, err := p.dialer.Dial(ctx, "tcp", address) | ||
if err != nil { | ||
p.cache.Set(address, Closed) | ||
return Closed, nil | ||
} | ||
_ = conn.Close() | ||
p.cache.Set(address, Open) | ||
return Open, nil | ||
}) | ||
|
||
if status, ok := code.(PortStatus); ok { | ||
if status == Closed { | ||
return ErrPortClosed | ||
} | ||
} | ||
return nil | ||
} | ||
|
||
// Close closes the ports cache and releases any allocated resources | ||
func (p *PortsCache) Close() { | ||
p.cache = nil | ||
p.group = singleflight.Group[string]{} | ||
} | ||
|
||
// Init initializes the ports package | ||
func Init(dialer *fastdialer.Dialer, size int) { | ||
portsCacher = NewPortsCache(dialer, size) | ||
} | ||
|
||
// Close closes the ports package | ||
func Close() { | ||
if portsCacher != nil { | ||
portsCacher.Close() | ||
} | ||
} | ||
|
||
// IsPortOpen checks if a port is open or not | ||
func IsPortOpen(input *contextargs.Context) error { | ||
if portsCacher == nil { | ||
return nil | ||
} | ||
return portsCacher.Do(input.Context(), input) | ||
} |