target based upstream options

internal/out/out.go

@@ -0,0 +1,143 @@
+package out
+
+import (
+	"context"
+	"log"
+	"net"
+	"net/url"
+	"sync"
+
+	"golang.org/x/net/proxy"
+
+	"edgaru089.ink/go/regolith/internal/util"
+)
+
+// Configures upstream TCP connections.
+//
+// Does not support HTTP proxies as upstream. They're
+// just a pain to use besides CONNECT commands.
+//
+// Upstreams are specified as 'direct' for direct dialing
+// by the Go net package, or 'socks5://user:pass@host:port'
+// for SOCKS5 upstreams.
+//
+// Numeric addresses are always dialed directly.
+type Config struct {
+	Default string              // Default URI
+	Proxied map[string][]string // Proxy URI to hostname array mapping
+}
+
+// Dialer implements x/net/proxy.Dialer from the Config given to it.
+type Dialer struct {
+	lock sync.RWMutex
+
+	hostnames map[string]string
+	uris      map[string]interface {
+		proxy.Dialer
+		proxy.ContextDialer
+	}
+
+	def interface {
+		proxy.Dialer
+		proxy.ContextDialer
+	}
+}
+
+var _ proxy.ContextDialer = &Dialer{} // *Dialer implements x/net/proxy.ContextDialer
+var _ proxy.Dialer = &Dialer{}        // *Dialer implements x/net/proxy.Dialer
+
+// New creates a new Dialer from a Config.
+//
+// You can also &out.Dialer{} and then call Load on it.
+func New(c Config) (d *Dialer) {
+	d = &Dialer{}
+	d.Load(c)
+	return
+}
+
+// Load loads/reloads the Dialer struct.
+func (d *Dialer) Load(c Config) {
+	d.lock.Lock()
+	defer d.lock.Unlock()
+
+	if len(c.Default) == 0 {
+		c.Default = "direct"
+	}
+	log.Print("default dialer ", c.Default)
+
+	if d.hostnames == nil {
+		d.hostnames = make(map[string]string)
+		d.uris = make(map[string]interface {
+			proxy.Dialer
+			proxy.ContextDialer
+		})
+	} else {
+		clear(d.hostnames)
+		clear(d.uris)
+	}
+
+	for uri, hosts := range c.Proxied {
+		d.uris[uri] = nil
+		for _, host := range hosts {
+			d.hostnames[host] = uri
+			log.Printf("dialer( %s ) => %s", host, uri)
+		}
+	}
+	d.uris[c.Default] = nil
+
+	// construct the Dialers
+	for uri := range d.uris {
+		if uri == "direct" {
+			d.uris[uri] = proxy.Direct
+		} else {
+			url, err := url.Parse(uri)
+			if err != nil {
+				log.Print("URI parse error: ", err)
+				d.uris[uri] = proxy.Direct
+				continue
+			}
+			dx, err := proxy.FromURL(url, proxy.Direct)
+			if err != nil {
+				log.Print("Proxy.FromURL error: ", err)
+				d.uris[uri] = proxy.Direct
+				continue
+			}
+
+			var ok bool
+			if d.uris[uri], ok = dx.(interface {
+				proxy.ContextDialer
+				proxy.Dialer
+			}); !ok {
+				log.Print("Proxy.FromURL unable to cast to proxy.Dialer+ContextDialer")
+				d.uris[uri] = proxy.Direct
+				continue
+			}
+		}
+	}
+
+	d.def = d.uris[c.Default]
+}
+
+func (d *Dialer) findDialer(host string) interface {
+	proxy.Dialer
+	proxy.ContextDialer
+} {
+	d.lock.RLock()
+	defer d.lock.RUnlock()
+
+	if dx, ok := d.uris[d.hostnames[host]]; ok {
+		return dx
+	} else {
+		return d.def
+	}
+}
+
+func (d *Dialer) Dial(network, address string) (net.Conn, error) {
+	host, _ := util.SplitHostPort(address)
+	return d.findDialer(host).Dial(network, address)
+}
+
+func (d *Dialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	host, _ := util.SplitHostPort(address)
+	return d.findDialer(host).DialContext(ctx, network, address)
+}