probe/discovery/protocol/tcp/connection/connection.go

package connection

import (
	"context"
	"net"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	omd "git.loafle.net/overflow/model/discovery"
	omm "git.loafle.net/overflow/model/meta"
	omu "git.loafle.net/overflow/model/util"
	"git.loafle.net/overflow_scanner/probe/discovery/session"
	"golang.org/x/sync/semaphore"
)

const (
	defaultUlimit = 1024
)

var (
	discovererType = omm.ToMetaDiscovererType(omm.MetaDiscovererTypeEnumTCPConnection)
)

func Scan(discoverySession session.DiscoverySession, targetHost *omd.Host) error {
	if nil == targetHost || nil == discoverySession.DiscoverPort() || !discoverySession.DiscoverPort().IncludeTCP {
		return nil
	}

	if !discoverySession.HostNeedScan(targetHost, discovererType) {
		return nil
	}

	dp := discoverySession.DiscoverPort()

	lock := semaphore.NewWeighted(Ulimit())
	var wg sync.WaitGroup
	stopChan := make(chan struct{})
	timeout := 500 * time.Millisecond
	ports := make(map[int]*omd.Port)

Loop:
	for portNumber := dp.FirstScanRange; portNumber < dp.LastScanRange; portNumber++ {
		if nil != dp.ExcludePorts {
			for _, exPortNumber := range dp.ExcludePorts {
				if portNumber == exPortNumber {
					continue Loop
				}
			}
		}

		lock.Acquire(context.TODO(), 1)
		wg.Add(1)

		go func(port int) {
			defer func() {
				lock.Release(1)
				wg.Done()
			}()

			tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
		}(portNumber)

		select {
		case <-time.After(time.Microsecond * 100):
		case <-discoverySession.StopChan():
			close(stopChan)
			return nil
		}
	}

	wg.Wait()

	return nil
}

func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Port, targetHost *omd.Host, port int, timeout time.Duration, stopChan <-chan struct{}) {
	addr := net.JoinHostPort(targetHost.Address, strconv.Itoa(port))
	conn, err := net.DialTimeout("tcp", addr, timeout)
	dp := discoverySession.DiscoverPort()

	if err != nil {
		if strings.Contains(err.Error(), "too many open files") {
			select {
			case <-stopChan:
				return
			default:
			}
			time.Sleep(timeout)
			tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
		}
		return
	}

	conn.Close()

	if _, ok := ports[port]; ok || !dp.Contains(port) {
		return
	}

	p := omd.NewPort(
		targetHost,
		omm.ToMetaPortType(omm.MetaPortTypeEnumTCP),
		port,
	)
	p.DiscoveredDate = omu.NowPtr()

	ports[port] = p

	go discoverySession.AddPort(discovererType, p, nil)
}

func Ulimit() int64 {
	out, err := exec.Command("ulimit", "-n").Output()
	if nil != err {
		return defaultUlimit
	}

	s := strings.TrimSpace(string(out))

	i, err := strconv.ParseInt(s, 10, 64)
	if nil != err {
		return defaultUlimit
	}

	return i
}