Socks5RoundTripper for net/http

socks5.go

package socks5

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

type Socks5RoundTripper struct {
	Proxy        string
	Resp         *http.Response
	conn         net.Conn
	timeout      time.Duration
	ModifyHeader func(*http.Header)
}

func (Socks5 *Socks5RoundTripper) generateRequest(req *http.Request) (message []byte, err error) {
	host := ""
	var port uint32
	var msg = make([]byte, 128)
	if len(req.URL.Host) > 120 {
		return nil, errors.New("req.URL.Host too long")
	}
	copy(msg, []byte("\x05\x01\x00\x03"))
	if !strings.Contains(req.URL.Host, ":") {
		host = req.URL.Host
		port = uint32(80)
	} else {
		hostAndPort := strings.Split(host, ":")
		host = hostAndPort[0]
		p, err := strconv.Atoi(hostAndPort[1])
		if err != nil {
			return nil, errors.New("Wired req.URL.Host")
		}
		port = uint32(p)
	}
	msg[4] = byte(len(host))
	copy(msg[5:], []byte(host))
	portBuf := make([]byte, 2)
	binary.BigEndian.PutUint16(portBuf, uint16(port))
	copy(msg[len(host)+5:], portBuf)
	return msg[:len(host)+7], nil
}

func (Socks5 *Socks5RoundTripper) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	if Socks5.ModifyHeader != nil {
		Socks5.ModifyHeader(&req.Header)
	}
	reschan := make(chan *http.Response, 1)
	errchan := make(chan error, 1)
	go func() {
		conn, err := net.Dial("tcp", Socks5.Proxy)
		if err != nil {
			errchan <- err
			return
		}
		Socks5.conn = conn
		if tcpconn, ok := Socks5.conn.(*net.TCPConn); ok {
			tcpconn.SetNoDelay(true)
		}
		_, err = Socks5.conn.Write([]byte("\x05\x01\x00"))
		if err != nil {
			errchan <- err
			return
		}
		msg, err := Socks5.generateRequest(req)
		if err != nil {
			errchan <- err
			return
		}
		_, err = Socks5.conn.Write(msg)
		if err != nil {
			errchan <- err
			return
		}
		msg = make([]byte, 32)
		_, err = Socks5.conn.Read(msg)
		if err != nil {
			errchan <- err
			return
		}
		if !bytes.HasPrefix(msg, []byte{'\x05', '\x00'}) {
			errchan <- errors.New("socks5: request failed")
			return
		}
		err = req.Write(Socks5.conn)
		if err != nil {
			errchan <- errors.New("socks5: request failed")
			return
		}
		reader := bufio.NewReader(Socks5.conn)
		res, err := http.ReadResponse(reader, req)
		if err != nil {
			errchan <- errors.New("socks5: request failed")
			return
		}
		reschan <- res
	}()
	ticker := time.After(Socks5.timeout)
	for {
		select {
		case err := <-errchan:
			{
				return nil, err
			}
		case <-ticker:
			{
				return nil, errors.New("timeout")
			}
		case res := <-reschan:
			{
				return res, nil
			}
		}
	}
}

func SetProxy(url string) (err error) {
	_, err = net.ResolveTCPAddr("tcp", url)
	if err != nil {
		return err
	}
	http.DefaultClient = &http.Client{}
	socks5 := &Socks5RoundTripper{}
	socks5.Proxy = url
	socks5.timeout = 30 * time.Second
	http.DefaultClient.Transport = socks5
	return
}

func ModifyHeader(mh func(*http.Header)) {
	if st, ok := http.DefaultClient.Transport.(*Socks5RoundTripper); ok {
		st.ModifyHeader = mh
	}
}

func SetTimeout(timeout int64) {
	if st, ok := http.DefaultClient.Transport.(*Socks5RoundTripper); ok {
		if timeout < 0 {
			timeout = 30
		}
		st.timeout = time.Duration(timeout)
	}
}

func UnsetProxy() {
	http.DefaultClient = &http.Client{}
}

socks5_example.go

func testGet(url string) {
  socks5.SetProxy("127.0.0.1:9988")
  socks5.ModifyHeader(func(header *http.Header) {
    header.Set("User-Agent", "Mozilla/5.0 (X11; Linux x86_64; rv:25.0) Gecko/20100101 Firefox/25.0")
  })
  resp, err := http.Get(url)
  if err != nil {
    panic(err)
  }
  fmt.Println(resp.Status)
  io.Copy(os.Stdout, resp.Body)
  fmt.Println("")
}

func main() {
  testGet(os.Args[1])
}