ovpn-lb-socks5/go-socks-lb/main.go

package main

import (
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"

	wr "github.com/mroth/weightedrand"
	"golang.org/x/net/proxy"
	"gopkg.in/yaml.v2"
)

const (
	URL = "http://connectivitycheck.gstatic.com/generate_204"
)

func proxyTestStatusCode(proxyURL string, URL string, StatusCode int) bool {
	// create a socks5 dialer
	u, err := url.Parse(proxyURL)
	if err != nil {
		fmt.Println("error parsing")
		return false
	}
	dialer, err := proxy.FromURL(u, proxy.Direct)
	if err != nil {
		fmt.Fprintln(os.Stderr, "can't connect to the proxy:", err)
		return false
	}
	// setup a http client
	httpTransport := &http.Transport{}
	httpClient := &http.Client{
		Transport: httpTransport,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			fmt.Printf("redirect %v", *req)
			return http.ErrUseLastResponse
		},
	}
	// set our socks5 as the dialer
	httpTransport.Dial = dialer.Dial
	// create a request
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, "can't create request:", err)
		return false
	}
	// use the http client to fetch the page
	resp, err := httpClient.Do(req)
	if err != nil {
		fmt.Fprintln(os.Stderr, "can't GET page:", err)
		return false
	}
	// defer resp.Body.Close()
	// b, err := ioutil.ReadAll(resp.Body)
	// if err != nil {
	// 	fmt.Fprintln(os.Stderr, "error reading body:", err)
	// 	return false
	// }
	return resp.StatusCode == StatusCode
}

func testTask(testfn func() bool, interval int, stop chan bool) {
	for {
		go testfn()
		select {
		case <-stop:
			fmt.Println("stopping")
			return
		default:
		}
		time.Sleep(time.Duration(interval) * time.Second)
	}
}

type proxyInst struct {
	URL           string
	Weight        int
	StatusHistory []bool
}
type ProxyManager struct {
	Proxys             []proxyInst
	Cache              map[string]int
	Sticky             bool
	CacheCleanInterval int
	Chooser            wr.Chooser
	mux                sync.Mutex
}
type proxyConf struct {
	Proxys []struct {
		URL    string `yaml:"url"`
		Weight int    `yaml:"weight"`
	} `yaml:"proxy"`
	Sticky             bool `yaml:"sticky"`
	CacheCleanInterval int  `yaml:"cache-clean-interval"`
}

func NewPM(confFp string) (*ProxyManager, error) {
	pm := ProxyManager{}
	f, err := os.Open(confFp)
	if err != nil {
		return &ProxyManager{}, err
	}
	defer f.Close()

	var cfg proxyConf
	decoder := yaml.NewDecoder(f)
	err = decoder.Decode(&cfg)
	if err != nil {
		return &ProxyManager{}, err
	}
	var chooseArr []wr.Choice
	for idx, pc := range cfg.Proxys {
		var pi proxyInst
		pi.URL = pc.URL
		pi.Weight = pc.Weight
		pm.Proxys = append(pm.Proxys, pi)
		chooseArr = append(chooseArr, wr.Choice{Item: idx, Weight: uint(pi.Weight)})
	}
	pm.Sticky = cfg.Sticky
	pm.CacheCleanInterval = cfg.CacheCleanInterval
	pm.Cache = make(map[string]int)
	pm.Chooser = wr.NewChooser(chooseArr...)
	go pm.keepClearingCache()
	return &pm, nil
}

func (pm *ProxyManager) Get(addr string) string {
	pm.mux.Lock()
	defer pm.mux.Unlock()
	// fmt.Println(pm.Cache)
	if pm.Sticky {
		idx, ok := pm.Cache[addr]
		if ok {
			// fmt.Println("match addr", addr, "using:", idx)
			return pm.Proxys[idx].URL
		}
	}
	idx := pm.Chooser.Pick().(int)
	pm.Cache[addr] = idx
	// fmt.Println("addr", addr, "using:", idx)
	return pm.Proxys[idx].URL
}

func (pm *ProxyManager) ClearCache() {
	// fmt.Println("clearing cache")
	pm.mux.Lock()
	pm.Cache = make(map[string]int)
	pm.mux.Unlock()
	// fmt.Println("after:", pm.Cache)
}

func (pm *ProxyManager) keepClearingCache() {
	interval := pm.CacheCleanInterval
	if interval <= 0 {
		return
	}
	for {
		time.Sleep(time.Duration(interval) * time.Second)
		pm.ClearCache()
	}
}

func main() {
	// testFn := func() bool {
	// 	status := proxyTestStatusCode("socks5://127.0.0.1:1080", URL, 204)
	// 	fmt.Println(status)
	// 	return status
	// }
	// stop := make(chan bool)
	// go testTask(testFn, 5, stop)
	// time.Sleep(10 * time.Second)
	// stop <- true
	// time.Sleep(10 * time.Second)

	configPath := flag.String("config", "", "Config file.")
	bindAddr := flag.String("bind", "127.0.0.1:7000", "Bind address and port")
	flag.Parse()
	if *configPath == "" {
		flag.PrintDefaults()
		os.Exit(1)
	}
	rand.Seed(time.Now().UTC().UnixNano()) // always seed random!

	var pm *ProxyManager
	pm, err := NewPM(*configPath)
	if err != nil {
		return
	}
	if pm.Sticky {
		fmt.Printf("Sticky with interval %ds.\n", pm.CacheCleanInterval)
	} else {
		fmt.Println("Randomize every connection.")
	}
	conf := &Config{PM: pm}
	server, err := New(conf)
	if err != nil {
		panic(err)
	}
	// Create SOCKS5 proxy on localhost port 8000
	if err := server.ListenAndServe("tcp", *bindAddr); err != nil {
		panic(err)
	}
}