arblant
firefox-vpn-client / core / internal / tun / relay.go
relay.go67 lines · 1.7 KB
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
package tun

import (
	"context"
	"errors"
	"io"
	"net"
	"sync"
	"time"

	"golang.org/x/net/proxy"
)

// Dialer opens a connection to target ("host:port") through the proxy.
type Dialer func(ctx context.Context, target string) (net.Conn, error)

// Allow reports whether a connection to host:port is sent to the proxy at
// all; one it refuses is reset at once, without being tried.
type Allow func(host, port string) bool

// dialTimeout bounds connecting through the proxy, which may first wake
// and connect to Mozilla.
const dialTimeout = 30 * time.Second

// SOCKS5Dialer dials through the SOCKS5 proxy at addr, with its login if
// user is not empty.
func SOCKS5Dialer(addr, user, password string) (Dialer, error) {
	var auth *proxy.Auth
	if user != "" {
		auth = &proxy.Auth{User: user, Password: password}
	}
	d, err := proxy.SOCKS5("tcp", addr, auth, &net.Dialer{Timeout: dialTimeout})
	if err != nil {
		return nil, err
	}
	cd, ok := d.(proxy.ContextDialer)
	if !ok {
		return nil, errors.New("the SOCKS5 dialer cannot take a context")
	}
	return func(ctx context.Context, target string) (net.Conn, error) {
		return cd.DialContext(ctx, "tcp", target)
	}, nil
}

// closeWriter is a connection that can end its sending half.
type closeWriter interface{ CloseWrite() error }

// Relay copies between a and b until both directions end, passing on each
// end of stream as a half close, then closes both.
func Relay(a, b net.Conn) {
	var wg sync.WaitGroup
	copyHalf := func(dst, src net.Conn) {
		defer wg.Done()
		_, _ = io.Copy(dst, src)
		if cw, ok := dst.(closeWriter); ok {
			_ = cw.CloseWrite()
		} else {
			_ = dst.Close()
		}
	}
	wg.Add(2)
	go copyHalf(a, b)
	go copyHalf(b, a)
	wg.Wait()
	_ = a.Close()
	_ = b.Close()
}