nuclei/pkg/js/utils/pgwrap/pgwrap.go
HD Moore f26996cb89
Remove singletons from Nuclei engine (continuation of #6210) (#6296)
* introducing execution id

* wip

* .

* adding separate execution context id

* lint

* vet

* fixing pg dialers

* test ignore

* fixing loader FD limit

* test

* fd fix

* wip: remove CloseProcesses() from dev merge

* wip: fix merge issue

* protocolstate: stop memguarding on last dialer delete

* avoid data race in dialers.RawHTTPClient

* use shared logger and avoid race conditions

* use shared logger and avoid race conditions

* go mod

* patch executionId into compiled template cache

* clean up comment in Parse

* go mod update

* bump echarts

* address merge issues

* fix use of gologger

* switch cmd/nuclei to options.Logger

* address merge issues with go.mod

* go vet: address copy of lock with new Copy function

* fixing tests

* disable speed control

* fix nil ExecuterOptions

* removing deprecated code

* fixing result print

* default logger

* cli default logger

* filter warning from results

* fix performance test

* hardcoding path

* disable upload

* refactor(runner): uses `Warning` instead of `Print` for `pdcpUploadErrMsg`

Signed-off-by: Dwi Siswanto <git@dw1.io>

* Revert "disable upload"

This reverts commit 114fbe6663361bf41cf8b2645fd2d57083d53682.

* Revert "hardcoding path"

This reverts commit cf12ca800e0a0e974bd9fd4826a24e51547f7c00.

---------

Signed-off-by: Dwi Siswanto <git@dw1.io>
Co-authored-by: Mzack9999 <mzack9999@protonmail.com>
Co-authored-by: Dwi Siswanto <git@dw1.io>
Co-authored-by: Dwi Siswanto <25837540+dwisiswant0@users.noreply.github.com>
2025-07-10 01:17:26 +05:30

79 lines
2.2 KiB
Go

package pgwrap
import (
"context"
"database/sql"
"database/sql/driver"
"fmt"
"net"
"net/url"
"time"
"github.com/lib/pq"
"github.com/projectdiscovery/fastdialer/fastdialer"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
)
const (
PGWrapDriver = "pgwrap"
)
type pgDial struct {
executionId string
}
func (p *pgDial) Dial(network, address string) (net.Conn, error) {
dialers := protocolstate.GetDialersWithId(p.executionId)
if dialers == nil {
return nil, fmt.Errorf("dialers not initialized for %s", p.executionId)
}
return dialers.Fastdialer.Dial(context.TODO(), network, address)
}
func (p *pgDial) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
dialers := protocolstate.GetDialersWithId(p.executionId)
if dialers == nil {
return nil, fmt.Errorf("dialers not initialized for %s", p.executionId)
}
ctx, cancel := context.WithTimeoutCause(context.Background(), timeout, fastdialer.ErrDialTimeout)
defer cancel()
return dialers.Fastdialer.Dial(ctx, network, address)
}
func (p *pgDial) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
dialers := protocolstate.GetDialersWithId(p.executionId)
if dialers == nil {
return nil, fmt.Errorf("dialers not initialized for %s", p.executionId)
}
return dialers.Fastdialer.Dial(ctx, network, address)
}
// Unfortunately lib/pq does not provide easy to customize or
// replace dialer so we need to hijack it by wrapping it in our own
// driver and register it as postgres driver
// PgDriver is the Postgres database driver.
type PgDriver struct{}
// Open opens a new connection to the database. name is a connection string.
// Most users should only use it through database/sql package from the standard
// library.
func (d PgDriver) Open(name string) (driver.Conn, error) {
// Parse the connection string to get executionId
u, err := url.Parse(name)
if err != nil {
return nil, fmt.Errorf("invalid connection string: %v", err)
}
values := u.Query()
executionId := values.Get("executionId")
// Remove executionId from the connection string
values.Del("executionId")
u.RawQuery = values.Encode()
return pq.DialOpen(&pgDial{executionId: executionId}, u.String())
}
func init() {
sql.Register(PGWrapDriver, &PgDriver{})
}