checker.go

package main

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"h12.io/socks"
)

type checker struct {
	cfg         *config
	st          *stats
	exp         *exporter
	input       chan string
	wg          sync.WaitGroup
	defaultPort int
}

type proxy struct {
	host string
	port int
}

func newChecker(cfg *config, st *stats, exp *exporter, defaultPort int) *checker {
	return &checker{
		cfg:         cfg,
		st:          st,
		exp:         exp,
		input:       make(chan string, cfg.Threads*2),
		defaultPort: defaultPort,
	}
}

func (c *checker) start(ctx context.Context) {
	for i := 0; i < c.cfg.Threads; i++ {
		c.wg.Add(1)
		go c.worker(ctx)
	}
}

func (c *checker) wait() {
	close(c.input)
	c.wg.Wait()
}

func (c *checker) worker(ctx context.Context) {
	defer c.wg.Done()

	for {
		select {
		case <-ctx.Done():
			return
		case p, ok := <-c.input:
			if !ok {
				return
			}
			c.check(ctx, p)
		}
	}
}

func (c *checker) check(ctx context.Context, proxyStr string) {
	defer c.st.incChecked()

	p, err := c.parse(proxyStr)
	if err != nil {
		c.st.incErr("parse")
		return
	}

	client := c.makeClient(p)

	req, _ := http.NewRequestWithContext(ctx, "GET", c.cfg.Website, nil)
	req.Header.Set("User-Agent", c.cfg.Headers.UserAgent)
	req.Header.Set("Accept", c.cfg.Headers.Accept)
	req.Header.Set("Connection", "close")

	resp, err := client.Do(req)
	if err != nil {
		if strings.Contains(err.Error(), "timeout") {
			c.st.incErr("timeout")
		} else {
			c.st.incErr("proxy")
		}
		return
	}
	resp.Body.Close()

	if resp.StatusCode != 200 {
		c.st.incErr("status")
		return
	}

	addr := fmt.Sprintf("%s:%d", p.host, p.port)
	c.st.incSuccess()
	c.exp.add(addr)

	if c.cfg.Print.Enabled {
		go c.printHit(p)
	}
}

func (c *checker) parse(s string) (proxy, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return proxy{}, fmt.Errorf("empty")
	}

	parts := strings.Split(s, ":")
	host := parts[0]
	port := c.defaultPort

	if len(parts) > 1 {
		p, err := strconv.Atoi(strings.TrimSpace(parts[1]))
		if err != nil || p < 1 || p > 65535 {
			return proxy{}, fmt.Errorf("bad port")
		}
		port = p
	} else if port == 0 {
		return proxy{}, fmt.Errorf("no port")
	}

	return proxy{host: host, port: port}, nil
}

func (c *checker) makeClient(p proxy) *http.Client {
	timeout := c.cfg.getTimeout()
	addr := fmt.Sprintf("%s:%d", p.host, p.port)

	var tr *http.Transport

	switch c.cfg.Proxy {
	case "socks4":
		tr = &http.Transport{
			Dial:              socks.Dial(fmt.Sprintf("socks4://%s?timeout=%ds", addr, int(timeout.Seconds()))),
			DisableKeepAlives: true,
		}
	case "socks5":
		tr = &http.Transport{
			Dial:              socks.Dial(fmt.Sprintf("socks5://%s?timeout=%ds", addr, int(timeout.Seconds()))),
			DisableKeepAlives: true,
		}
	default:
		proxyURL, _ := url.Parse("http://" + addr)
		tr = &http.Transport{
			Proxy: http.ProxyURL(proxyURL),
			DialContext: (&net.Dialer{
				Timeout:   timeout,
				KeepAlive: 30 * time.Second,
			}).DialContext,
			DisableKeepAlives:   true,
			TLSHandshakeTimeout: timeout,
		}
	}

	return &http.Client{
		Transport: tr,
		Timeout:   timeout,
	}
}

func (c *checker) printHit(p proxy) {
	addr := fmt.Sprintf("%s:%d", p.host, p.port)

	if c.cfg.Print.GeoInfo {
		info := fetchGeo(p.host)
		if info != nil {
			fmt.Printf("\n\033[32m>>>\033[0m %s \033[90m│\033[0m \033[36m%s\033[0m \033[90m│\033[0m %s\n", addr, info.Country, info.AS)
			return
		}
	}

	fmt.Printf("\n\033[32m>>>\033[0m %s\n", addr)
}