Golang muxado.Client函數代碼示例

func HTTPDialer(network, proxyUrl, addr string, tlsConfig *tls.Config) func() (muxado.Session, error) {
	return func() (sess muxado.Session, err error) {
		// parse the proxy address
		var parsedUrl *url.URL
		if parsedUrl, err = url.Parse(proxyUrl); err != nil {
			return
		}

		var proxyAuth string
		if parsedUrl.User != nil {
			proxyAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(parsedUrl.User.String()))
		}

		var dial func(net, addr string) (net.Conn, error)
		switch parsedUrl.Scheme {
		case "http":
			dial = net.Dial
		case "https":
			dial = func(net, addr string) (net.Conn, error) { return tls.Dial(net, addr, new(tls.Config)) }
		default:
			err = fmt.Errorf("Proxy URL scheme must be http or https, got: %s", parsedUrl.Scheme)
			return
		}

		// dial the proxy
		var conn net.Conn
		if conn, err = dial(network, parsedUrl.Host); err != nil {
			return
		}

		// send an HTTP proxy CONNECT message
		req, err := http.NewRequest("CONNECT", "https://"+addr, nil)
		if err != nil {
			return
		}

		if proxyAuth != "" {
			req.Header.Set("Proxy-Authorization", proxyAuth)
		}
		req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; ngrok)")
		req.Write(conn)

		// read the proxy's response
		resp, err := http.ReadResponse(bufio.NewReader(conn), req)
		if err != nil {
			return
		}
		resp.Body.Close()

		if resp.StatusCode != 200 {
			err = fmt.Errorf("Non-200 response from proxy server: %s", resp.Status)
			return
		}

		// upgrade to TLS
		conn = tls.Client(conn, tlsConfig)

		return muxado.Client(conn), nil
	}
}

func TestDisconnect(t *testing.T) {
	g := newTestingGateway("TestDisconnect", t)
	defer g.Close()

	if err := g.Disconnect("bar"); err == nil {
		t.Fatal("disconnect removed unconnected peer")
	}

	// dummy listener to accept connection
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatal("couldn't start listener:", err)
	}
	go func() {
		conn, err := l.Accept()
		if err != nil {
			t.Fatal("accept failed:", err)
		}
		conn.Close()
	}()
	// skip standard connection protocol
	conn, err := net.Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	id := g.mu.Lock()
	g.addPeer(&peer{addr: "foo", sess: muxado.Client(conn)})
	g.mu.Unlock(id)
	if err := g.Disconnect("foo"); err != nil {
		t.Fatal("disconnect failed:", err)
	}
}

func TestAddPeer(t *testing.T) {
	g := newTestingGateway("TestAddPeer", t)
	defer g.Close()
	id := g.mu.Lock()
	defer g.mu.Unlock(id)
	g.addPeer(&peer{addr: "foo", sess: muxado.Client(nil)})
	if len(g.peers) != 1 {
		t.Fatal("gateway did not add peer")
	}
}

// Connect establishes a persistent connection to a peer, and adds it to the
// Gateway's peer list.
func (g *Gateway) Connect(addr modules.NetAddress) error {
	if addr == g.Address() {
		return errors.New("can't connect to our own address")
	}

	id := g.mu.RLock()
	_, exists := g.peers[addr]
	g.mu.RUnlock(id)
	if exists {
		return errors.New("peer already added")
	}

	conn, err := net.DialTimeout("tcp", string(addr), dialTimeout)
	if err != nil {
		return err
	}
	// send our version
	if err := encoding.WriteObject(conn, "0.3.3"); err != nil {
		return err
	}
	// read version ack
	var remoteVersion string
	if err := encoding.ReadObject(conn, &remoteVersion, maxAddrLength); err != nil {
		return err
	} else if remoteVersion == "reject" {
		return errors.New("peer rejected connection")
	}
	// decide whether to accept this version
	if build.VersionCmp(remoteVersion, "0.3.3") < 0 {
		conn.Close()
		return errors.New("unacceptable version: " + remoteVersion)
	}

	g.log.Println("INFO: connected to new peer", addr)

	id = g.mu.Lock()
	g.addPeer(&peer{addr: addr, sess: muxado.Client(conn), inbound: false})
	g.mu.Unlock(id)

	// call initRPCs
	id = g.mu.RLock()
	var wg sync.WaitGroup
	wg.Add(len(g.initRPCs))
	for name, fn := range g.initRPCs {
		go func(name string, fn modules.RPCFunc) {
			// errors here are non-fatal
			g.RPC(addr, name, fn)
			wg.Done()
		}(name, fn)
	}
	g.mu.RUnlock(id)
	wg.Wait()

	return nil
}

func (t transport) NewConn(nc net.Conn, isServer bool) (smux.Conn, error) {
	var s muxado.Session
	if isServer {
		s = muxado.Server(nc)
	} else {
		s = muxado.Client(nc)
	}
	cl := make(chan struct{})
	go func() {
		s.Wait()
		close(cl)
	}()
	return &conn{ms: s, closed: cl}, nil
}

func SOCKS5Dialer(network, proxyAddr, user, password, addr string) (muxado.Session, error) {
	proxyAuth := &proxy.Auth{User: user, Password: password}

	proxyDialer, err := proxy.SOCKS5(network, proxyAddr, proxyAuth, proxy.Direct)
	if err != nil {
		return nil, err
	}

	conn, err := proxyDialer.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	// upgrade to TLS
	conn = tls.Client(conn, new(tls.Config))

	return muxado.Client(conn), nil
}

func TestRandomInboundPeer(t *testing.T) {
	g := newTestingGateway("TestRandomInboundPeer", t)
	defer g.Close()
	id := g.mu.Lock()
	defer g.mu.Unlock(id)
	_, err := g.randomInboundPeer()
	if err != errNoPeers {
		t.Fatal("expected errNoPeers, got", err)
	}

	g.addPeer(&peer{addr: "foo", sess: muxado.Client(new(dummyConn)), inbound: true})
	if len(g.peers) != 1 {
		t.Fatal("gateway did not add peer")
	}
	addr, err := g.randomInboundPeer()
	if err != nil || addr != "foo" {
		t.Fatal("gateway did not select random peer")
	}
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(dc string, addr net.Addr, version int) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsWrap != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		tlsConn, err := p.tlsWrap(dc, conn)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

	// Switch the multiplexing based on version
	var session muxSession
	if version < 2 {
		// Write the Consul multiplex byte to set the mode
		if _, err := conn.Write([]byte{byte(rpcMultiplex)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Create a multiplexed session
		session = &muxadoWrapper{muxado.Client(conn)}

	} else {
		// Write the Consul multiplex byte to set the mode
		if _, err := conn.Write([]byte{byte(rpcMultiplexV2)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Setup the logger
		conf := yamux.DefaultConfig()
		conf.LogOutput = p.logOutput

		// Create a multiplexed session
		session, _ = yamux.Client(conn, conf)
	}

	// Wrap the connection
	c := &Conn{
		refCount: 1,
		addr:     addr,
		session:  session,
		clients:  list.New(),
		lastUsed: time.Now(),
		version:  version,
		pool:     p,
	}
	return c, nil
}

// Client starts a new go-tunnel session on conn
func Client(conn net.Conn) *client.Session {
	return client.NewSession(muxado.Client(conn))
}

func TestListen(t *testing.T) {
	g := newTestingGateway("TestListen", t)
	defer g.Close()

	// compliant connect with old version
	conn, err := net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	addr := modules.NetAddress(conn.LocalAddr().String())
	// send version
	if err := encoding.WriteObject(conn, "0.1"); err != nil {
		t.Fatal("couldn't write version")
	}
	// read ack
	var ack string
	if err := encoding.ReadObject(conn, &ack, maxAddrLength); err != nil {
		t.Fatal(err)
	} else if ack != "reject" {
		t.Fatal("gateway should have rejected old version")
	}

	// a simple 'conn.Close' would not obey the muxado disconnect protocol
	muxado.Client(conn).Close()

	// compliant connect
	conn, err = net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	addr = modules.NetAddress(conn.LocalAddr().String())
	// send version
	if err := encoding.WriteObject(conn, build.Version); err != nil {
		t.Fatal("couldn't write version")
	}
	// read ack
	if err := encoding.ReadObject(conn, &ack, maxAddrLength); err != nil {
		t.Fatal(err)
	} else if ack == "reject" {
		t.Fatal("gateway should have given ack")
	}

	// g should add the peer
	var ok bool
	for !ok {
		id := g.mu.RLock()
		_, ok = g.peers[addr]
		g.mu.RUnlock(id)
	}

	muxado.Client(conn).Close()

	// g should remove the peer
	for ok {
		id := g.mu.RLock()
		_, ok = g.peers[addr]
		g.mu.RUnlock(id)
	}

	// uncompliant connect
	conn, err = net.Dial("tcp", string(g.Address()))
	if err != nil {
		t.Fatal("dial failed:", err)
	}
	if _, err := conn.Write([]byte("missing length prefix")); err != nil {
		t.Fatal("couldn't write malformed header")
	}
	// g should have closed the connection
	if n, err := conn.Write([]byte("closed")); err != nil && n > 0 {
		t.Error("write succeeded after closed connection")
	}
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(addr net.Addr) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsConfig != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		conn = tls.Client(conn, p.tlsConfig)
	}

	// Write the Consul multiplex byte to set the mode
	if _, err := conn.Write([]byte{byte(rpcMultiplex)}); err != nil {
		conn.Close()
		return nil, err
	}

	// Create a multiplexed session
	session := muxado.Client(conn)

	// Wrap the connection
	c := &Conn{
		refCount: 1,
		addr:     addr,
		session:  session,
		lastUsed: time.Now(),
	}

	// Monitor the session
	go func() {
		session.Wait()
		p.Lock()
		defer p.Unlock()
		if conn, ok := p.pool[addr.String()]; ok && conn.session == session {
			delete(p.pool, addr.String())
		}
	}()

	// Track this connection, handle potential race condition
	p.Lock()
	defer p.Unlock()
	if existing := p.pool[addr.String()]; existing != nil {
		session.Close()
		return existing, nil
	} else {
		p.pool[addr.String()] = c
		return c, nil
	}
}