golang tcp conn forward, support proxy

gistfile1.txt

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"

	"crypto/tls"
	"net/http"
	"net/url"

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

var (
	local     string
	remote    string
	proxyAddr string
	dialer    proxy.Dialer
)

type direct struct{}

// Direct is a direct proxy: one that makes network connections directly.
var Direct = direct{}

func (direct) Dial(network, addr string) (net.Conn, error) {
	return net.Dial(network, addr)
}

// httpsDialer
type httpsDialer struct{}

// HTTPSDialer is a https proxy: one that makes network connections on tls.
var HttpsDialer = httpsDialer{}
var TlsConfig = &tls.Config{}

func (d httpsDialer) Dial(network, addr string) (c net.Conn, err error) {
	c, err = tls.Dial("tcp", addr, TlsConfig)
	if err != nil {
		fmt.Println(err)
	}
	return
}

// httpProxy is a HTTP/HTTPS connect proxy.
type httpProxy struct {
	host     string
	haveAuth bool
	username string
	password string
	forward  proxy.Dialer
	uri      *url.URL
}

func newHTTPProxy(uri *url.URL, forward proxy.Dialer) (proxy.Dialer, error) {
	s := new(httpProxy)
	s.host = uri.Host
	s.forward = forward
	s.uri = uri
	if uri.User != nil {
		s.haveAuth = true
		s.username = uri.User.Username()
		s.password, _ = uri.User.Password()
	}

	return s, nil
}

func (s *httpProxy) Dial(network, addr string) (net.Conn, error) {

	// Dial and create the https client connection.
	c, err := s.forward.Dial("tcp", s.host)
	if err != nil {
		return nil, err
	}

	// HACK. http.ReadRequest also does this.
	reqURL, err := url.Parse("http://" + addr)
	if err != nil {
		c.Close()
		return nil, err
	}
	reqURL.Scheme = ""

	req, err := http.NewRequest("CONNECT", reqURL.String(), nil)
	if err != nil {
		c.Close()
		return nil, err
	}
	req.Close = false
	if s.haveAuth {
		req.SetBasicAuth(s.username, s.password)
	}
	// 代理验证头部
	req.Header.Set("Proxy-Authorization", req.Header.Get("Authorization"))
	req.Header.Set("User-Agent", "Powerby Gota")

	err = req.Write(c)
	if err != nil {
		c.Close()
		return nil, err
	}
	resp, err := http.ReadResponse(bufio.NewReader(c), req)
	if err != nil {
		// TODO close resp body ?
		resp.Body.Close()
		c.Close()
		return nil, err
	}
	resp.Body.Close()
	if resp.StatusCode > 300 {
		c.Close()
		err = fmt.Errorf("connect server using proxy error, StatusCode [%d]", resp.StatusCode)
		return nil, err
	}

	return c, nil
}

func init() {
	proxy.RegisterDialerType("http", newHTTPProxy)
	proxy.RegisterDialerType("https", newHTTPProxy)
}

func main() {

	flag.StringVar(&local, "l", "", "local bind")
	flag.StringVar(&remote, "r", "", "remote")
	flag.StringVar(&proxyAddr, "p", "", "proxy")
	flag.Parse()
	listener, err := net.Listen("tcp", local)
	if err != nil {
		panic("connection error:" + err.Error())
	}
	dialer = &net.Dialer{}

	if proxyAddr != "" {
		fmt.Printf("proxy: %v\n", proxyAddr)
		httpProxyURI, err := url.Parse(proxyAddr)
		if err != nil {
			panic(err)
		}
		dialer, err = proxy.FromURL(httpProxyURI, &net.Dialer{})
		if err != nil {
			panic(err)
		}
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Accept Error:", err)
			continue
		}
		log.Printf("accept conn from %v\n", conn.RemoteAddr())
		go copyConn(conn)
	}
}

func copyConn(src net.Conn) {
	defer src.Close()
	dst, err := dialer.Dial("tcp", remote)
	if err != nil {
		log.Printf("dail remote[%v] err:%v", remote, err)
		return
	}
	defer dst.Close()
	done := make(chan struct{})

	go func() {
		_, err = io.Copy(dst, src)
		if err != nil {
			log.Printf("copy err:%v", err)
		}
		done <- struct{}{}
	}()
	go func() {
		_, err := io.Copy(src, dst)
		if err != nil {
			log.Printf("copy err:%v", err)
		}
		done <- struct{}{}
	}()
	<-done
	<-done
}