vendor/golang.org/x/net/proxy/proxy_test.go

   1 // Copyright 2011 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 package proxy
   6
   7 import (
   8         "bytes"
   9         "fmt"
  10         "io"
  11         "net"
  12         "net/url"
  13         "os"
  14         "strconv"
  15         "strings"
  16         "sync"
  17         "testing"
  18 )
  19
  20 type proxyFromEnvTest struct {
  21         allProxyEnv string
  22         noProxyEnv  string
  23         wantTypeOf  Dialer
  24 }
  25
  26 func (t proxyFromEnvTest) String() string {
  27         var buf bytes.Buffer
  28         space := func() {
  29                 if buf.Len() > 0 {
  30                         buf.WriteByte(' ')
  31                 }
  32         }
  33         if t.allProxyEnv != "" {
  34                 fmt.Fprintf(&buf, "all_proxy=%q", t.allProxyEnv)
  35         }
  36         if t.noProxyEnv != "" {
  37                 space()
  38                 fmt.Fprintf(&buf, "no_proxy=%q", t.noProxyEnv)
  39         }
  40         return strings.TrimSpace(buf.String())
  41 }
  42
  43 func TestFromEnvironment(t *testing.T) {
  44         ResetProxyEnv()
  45
  46         type dummyDialer struct {
  47                 direct
  48         }
  49
  50         RegisterDialerType("irc", func(_ *url.URL, _ Dialer) (Dialer, error) {
  51                 return dummyDialer{}, nil
  52         })
  53
  54         proxyFromEnvTests := []proxyFromEnvTest{
  55                 {allProxyEnv: "127.0.0.1:8080", noProxyEnv: "localhost, 127.0.0.1", wantTypeOf: direct{}},
  56                 {allProxyEnv: "ftp://example.com:8000", noProxyEnv: "localhost, 127.0.0.1", wantTypeOf: direct{}},
  57                 {allProxyEnv: "socks5://example.com:8080", noProxyEnv: "localhost, 127.0.0.1", wantTypeOf: &PerHost{}},
  58                 {allProxyEnv: "irc://example.com:8000", wantTypeOf: dummyDialer{}},
  59                 {noProxyEnv: "localhost, 127.0.0.1", wantTypeOf: direct{}},
  60                 {wantTypeOf: direct{}},
  61         }
  62
  63         for _, tt := range proxyFromEnvTests {
  64                 os.Setenv("ALL_PROXY", tt.allProxyEnv)
  65                 os.Setenv("NO_PROXY", tt.noProxyEnv)
  66                 ResetCachedEnvironment()
  67
  68                 d := FromEnvironment()
  69                 if got, want := fmt.Sprintf("%T", d), fmt.Sprintf("%T", tt.wantTypeOf); got != want {
  70                         t.Errorf("%v: got type = %T, want %T", tt, d, tt.wantTypeOf)
  71                 }
  72         }
  73 }
  74
  75 func TestFromURL(t *testing.T) {
  76         endSystem, err := net.Listen("tcp", "127.0.0.1:0")
  77         if err != nil {
  78                 t.Fatalf("net.Listen failed: %v", err)
  79         }
  80         defer endSystem.Close()
  81         gateway, err := net.Listen("tcp", "127.0.0.1:0")
  82         if err != nil {
  83                 t.Fatalf("net.Listen failed: %v", err)
  84         }
  85         defer gateway.Close()
  86
  87         var wg sync.WaitGroup
  88         wg.Add(1)
  89         go socks5Gateway(t, gateway, endSystem, socks5Domain, &wg)
  90
  91         url, err := url.Parse("socks5://user:password@" + gateway.Addr().String())
  92         if err != nil {
  93                 t.Fatalf("url.Parse failed: %v", err)
  94         }
  95         proxy, err := FromURL(url, Direct)
  96         if err != nil {
  97                 t.Fatalf("FromURL failed: %v", err)
  98         }
  99         _, port, err := net.SplitHostPort(endSystem.Addr().String())
 100         if err != nil {
 101                 t.Fatalf("net.SplitHostPort failed: %v", err)
 102         }
 103         if c, err := proxy.Dial("tcp", "localhost:"+port); err != nil {
 104                 t.Fatalf("FromURL.Dial failed: %v", err)
 105         } else {
 106                 c.Close()
 107         }
 108
 109         wg.Wait()
 110 }
 111
 112 func TestSOCKS5(t *testing.T) {
 113         endSystem, err := net.Listen("tcp", "127.0.0.1:0")
 114         if err != nil {
 115                 t.Fatalf("net.Listen failed: %v", err)
 116         }
 117         defer endSystem.Close()
 118         gateway, err := net.Listen("tcp", "127.0.0.1:0")
 119         if err != nil {
 120                 t.Fatalf("net.Listen failed: %v", err)
 121         }
 122         defer gateway.Close()
 123
 124         var wg sync.WaitGroup
 125         wg.Add(1)
 126         go socks5Gateway(t, gateway, endSystem, socks5IP4, &wg)
 127
 128         proxy, err := SOCKS5("tcp", gateway.Addr().String(), nil, Direct)
 129         if err != nil {
 130                 t.Fatalf("SOCKS5 failed: %v", err)
 131         }
 132         if c, err := proxy.Dial("tcp", endSystem.Addr().String()); err != nil {
 133                 t.Fatalf("SOCKS5.Dial failed: %v", err)
 134         } else {
 135                 c.Close()
 136         }
 137
 138         wg.Wait()
 139 }
 140
 141 func socks5Gateway(t *testing.T, gateway, endSystem net.Listener, typ byte, wg *sync.WaitGroup) {
 142         defer wg.Done()
 143
 144         c, err := gateway.Accept()
 145         if err != nil {
 146                 t.Errorf("net.Listener.Accept failed: %v", err)
 147                 return
 148         }
 149         defer c.Close()
 150
 151         b := make([]byte, 32)
 152         var n int
 153         if typ == socks5Domain {
 154                 n = 4
 155         } else {
 156                 n = 3
 157         }
 158         if _, err := io.ReadFull(c, b[:n]); err != nil {
 159                 t.Errorf("io.ReadFull failed: %v", err)
 160                 return
 161         }
 162         if _, err := c.Write([]byte{socks5Version, socks5AuthNone}); err != nil {
 163                 t.Errorf("net.Conn.Write failed: %v", err)
 164                 return
 165         }
 166         if typ == socks5Domain {
 167                 n = 16
 168         } else {
 169                 n = 10
 170         }
 171         if _, err := io.ReadFull(c, b[:n]); err != nil {
 172                 t.Errorf("io.ReadFull failed: %v", err)
 173                 return
 174         }
 175         if b[0] != socks5Version || b[1] != socks5Connect || b[2] != 0x00 || b[3] != typ {
 176                 t.Errorf("got an unexpected packet: %#02x %#02x %#02x %#02x", b[0], b[1], b[2], b[3])
 177                 return
 178         }
 179         if typ == socks5Domain {
 180                 copy(b[:5], []byte{socks5Version, 0x00, 0x00, socks5Domain, 9})
 181                 b = append(b, []byte("localhost")...)
 182         } else {
 183                 copy(b[:4], []byte{socks5Version, 0x00, 0x00, socks5IP4})
 184         }
 185         host, port, err := net.SplitHostPort(endSystem.Addr().String())
 186         if err != nil {
 187                 t.Errorf("net.SplitHostPort failed: %v", err)
 188                 return
 189         }
 190         b = append(b, []byte(net.ParseIP(host).To4())...)
 191         p, err := strconv.Atoi(port)
 192         if err != nil {
 193                 t.Errorf("strconv.Atoi failed: %v", err)
 194                 return
 195         }
 196         b = append(b, []byte{byte(p >> 8), byte(p)}...)
 197         if _, err := c.Write(b); err != nil {
 198                 t.Errorf("net.Conn.Write failed: %v", err)
 199                 return
 200         }
 201 }
 202
 203 func ResetProxyEnv() {
 204         for _, env := range []*envOnce{allProxyEnv, noProxyEnv} {
 205                 for _, v := range env.names {
 206                         os.Setenv(v, "")
 207                 }
 208         }
 209         ResetCachedEnvironment()
 210 }
 211
 212 func ResetCachedEnvironment() {
 213         allProxyEnv.reset()
 214         noProxyEnv.reset()
 215 }