Golang yamux.Client函數代碼示例

func (c *Client) Session() (*yamux.Session, error) {
	if c.sess == nil {
		s, err := net.Dial("tcp", c.addr)
		if err != nil {
			return nil, err
		}

		if c.secure {
			sec, err := seconn.NewClient(s)
			if err != nil {
				return nil, err
			}

			c.conn = sec
		} else {
			c.conn = s
		}

		sess, err := yamux.Client(c.conn, muxConfig)
		if err != nil {
			return nil, err
		}

		c.sess = sess
	}

	return c.sess, nil
}

// NewClient creates a client from an already-open connection-like value.
// Dial is typically used instead.
func NewClient(conn io.ReadWriteCloser) (*Client, error) {
	// Create the yamux client so we can multiplex
	mux, err := yamux.Client(conn, nil)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Connect to the control stream.
	control, err := mux.Open()
	if err != nil {
		mux.Close()
		return nil, err
	}

	// Create the broker and start it up
	broker := newMuxBroker(mux)
	go broker.Run()

	// Build the client using our broker and control channel.
	return &Client{
		broker:  broker,
		control: rpc.NewClient(control),
	}, nil
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(addr string, timeout time.Duration) (*Conn, error) {
	// Try to dial the conn
	conn, err := tls.DialWithDialer(&net.Dialer{Timeout: timeout}, "tcp", addr, p.config)
	if err != nil {
		return nil, err
	}

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

	// Create a multiplexed session
	session, err := yamux.Client(conn, conf)
	if err != nil {
		conn.Close()
		return nil, err
	}

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

func startTLSServer(config *Config) (net.Conn, chan error) {
	errc := make(chan error, 1)

	tlsConfigServer, err := config.IncomingTLSConfig()
	if err != nil {
		errc <- err
		return nil, errc
	}

	client, server := net.Pipe()

	// Use yamux to buffer the reads, otherwise it's easy to deadlock
	muxConf := yamux.DefaultConfig()
	serverSession, _ := yamux.Server(server, muxConf)
	clientSession, _ := yamux.Client(client, muxConf)
	clientConn, _ := clientSession.Open()
	serverConn, _ := serverSession.Accept()

	go func() {
		tlsServer := tls.Server(serverConn, tlsConfigServer)
		if err := tlsServer.Handshake(); err != nil {
			errc <- err
		}
		close(errc)
		// Because net.Pipe() is unbuffered, if both sides
		// Close() simultaneously, we will deadlock as they
		// both send an alert and then block. So we make the
		// server read any data from the client until error or
		// EOF, which will allow the client to Close(), and
		// *then* we Close() the server.
		io.Copy(ioutil.Discard, tlsServer)
		tlsServer.Close()
	}()
	return clientConn, errc
}

func Main() {
	kingpin.Parse()

	// Create websocket session for tunnels
	origin := fmt.Sprintf("http://%s/", *host)
	url := fmt.Sprintf("ws://%s/", *host)
	conn, err := websocket.Dial(url, "", origin)
	if err != nil {
		log.Fatal(fmt.Sprintf("Could not connect to subway server (%s)!", url))
	}

	session, err := yamux.Client(conn, nil)
	if err != nil {
		panic(err)
	}

	// Start REST api
	router := mux.NewRouter()
	router.HandleFunc("/", handlers.CreateTunnel)

	n := negroni.New()
	n.Use(negronilogrus.NewMiddleware())
	render := render.New(render.Options{})
	subway := context.CreateSubway(session, render)
	n.Use(subway)
	n.UseHandler(router)
	n.Run(fmt.Sprintf(":%v", *port))

}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(region 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(region, conn)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

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

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

	// Create a multiplexed session
	session, err := yamux.Client(conn, conf)
	if err != nil {
		conn.Close()
		return nil, err
	}

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

func newMuxBrokerClient(rwc io.ReadWriteCloser) (*muxBroker, error) {
	s, err := yamux.Client(rwc, nil)
	if err != nil {
		return nil, err
	}

	return newMuxBroker(s), nil
}

func TestProvider_Disconnect(t *testing.T) {
	config := testProviderConfig()
	p, err := NewProvider(config)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer p.Shutdown()

	// Setup RPC client
	a, b := testConn(t)
	client, _ := yamux.Client(a, yamux.DefaultConfig())
	server, _ := yamux.Server(b, yamux.DefaultConfig())
	go p.handleSession(client, make(chan struct{}))

	stream, _ := server.Open()
	cc := msgpackrpc.NewCodec(false, false, stream)

	// Make the connect rpc
	args := &DisconnectRequest{
		NoRetry: true,
		Backoff: 300 * time.Second,
	}
	resp := &DisconnectResponse{}
	err = msgpackrpc.CallWithCodec(cc, "Client.Disconnect", args, resp)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	p.backoffLock.Lock()
	defer p.backoffLock.Unlock()

	if p.backoff != 300*time.Second {
		t.Fatalf("bad: %v", p.backoff)
	}
	if !p.noRetry {
		t.Fatalf("bad")
	}

	p.sessionLock.Lock()
	defer p.sessionLock.Unlock()

	if p.sessionID != "" {
		t.Fatalf("Bad: %v", p.sessionID)
	}
	if p.sessionAuth {
		t.Fatalf("Bad: %v", p.sessionAuth)
	}
}

func TestProvider_Connect(t *testing.T) {
	config := testProviderConfig()
	config.Service.Capabilities["foo"] = 1
	config.Handlers["foo"] = fooCapability(t)
	p, err := NewProvider(config)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer p.Shutdown()

	// Setup RPC client
	a, b := testConn(t)
	client, _ := yamux.Client(a, yamux.DefaultConfig())
	server, _ := yamux.Server(b, yamux.DefaultConfig())
	go p.handleSession(client, make(chan struct{}))

	stream, _ := server.Open()
	cc := msgpackrpc.NewCodec(false, false, stream)

	// Make the connect rpc
	args := &ConnectRequest{
		Capability: "foo",
		Meta: map[string]string{
			"zip": "zap",
		},
	}
	resp := &ConnectResponse{}
	err = msgpackrpc.CallWithCodec(cc, "Client.Connect", args, resp)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Should be successful!
	if !resp.Success {
		t.Fatalf("bad")
	}

	// At this point, we should be connected
	out := make([]byte, 9)
	n, err := stream.Read(out)
	if err != nil {
		t.Fatalf("err: %v %d", err, n)
	}

	if string(out) != "foobarbaz" {
		t.Fatalf("bad: %s", out)
	}
}

func openSession(d *Daemon) error {

	err := openWebsocket(d)
	if err != nil {
		d.log.Error("could not open websocket", log.Ctx{"error": err})
		return err
	}

	// Setup client side of yamux
	session, err := yamux.Client(d.wsconn, nil)
	if err != nil {
		d.log.Error("Yamux session failed!", log.Ctx{"error": err})
	}

	d.session = session

	return nil
}

// Accepts new tunnel session in a blocking fashion. When it returns it closes the listener
// and returns with the error why it returned.
func Accept(accepted chan Tunnel, publicHostname string, listener net.Listener) error {
	defer listener.Close()

	for id := 0; ; id++ {

		conn, err := listener.Accept()
		if err != nil {
			log.Printf("FATAL: %v", err)
			return err
		}

		log.Printf("connection accepted %v\n", conn.RemoteAddr())

		// perform handshake
		go func(conn net.Conn, id int) {
			name := petname.Generate(2, "-")

			hostname := fmt.Sprintf("%v.%v", name, publicHostname)
			publicAddress := fmt.Sprintf("http://%v", hostname)

			reader := bufio.NewReader(conn)
			request, err := http.ReadRequest(reader)
			if err != nil {
				log.Println(err.Error())
				return
			}

			localAddress := request.Header.Get("X-Publichost-Local")
			session, err := yamux.Client(conn, nil)
			if err != nil {
				log.Println(err.Error())
				return
			}

			log.Printf("tunnel created %v->%v\n", publicAddress, localAddress)
			accepted <- Tunnel{
				id,
				hostname,
				session,
				localAddress,
			}
		}(conn, id)
	}
}

func TestProvider_Flash(t *testing.T) {
	config := testProviderConfig()
	buf := bytes.NewBuffer(nil)
	config.LogOutput = buf
	p, err := NewProvider(config)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer p.Shutdown()

	// Setup RPC client
	a, b := testConn(t)
	client, _ := yamux.Client(a, yamux.DefaultConfig())
	server, _ := yamux.Server(b, yamux.DefaultConfig())
	go p.handleSession(client, make(chan struct{}))

	stream, _ := server.Open()
	cc := msgpackrpc.NewCodec(false, false, stream)

	// Make the connect rpc
	args := &FlashRequest{
		Severity: "INFO",
		Message:  "TESTING",
	}
	resp := &FlashResponse{}
	err = msgpackrpc.CallWithCodec(cc, "Client.Flash", args, resp)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Wait until we are disconnected
	start := time.Now()
	for time.Now().Sub(start) < time.Second {
		if bytes.Contains(buf.Bytes(), []byte("TESTING")) {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if !bytes.Contains(buf.Bytes(), []byte("TESTING")) {
		t.Fatalf("missing: %s", buf)
	}
}

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

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

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

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

	// Create a multiplexed session
	session, err := yamux.Client(con, conf)
	if err != nil {
		con.Close()
		return nil, err
	}

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

// initClient does the common initialization
func initClient(conn net.Conn, opts *Opts) (*Client, error) {
	// Send the preamble
	_, err := conn.Write([]byte(clientPreamble))
	if err != nil {
		return nil, fmt.Errorf("preamble write failed: %v", err)
	}

	// Wrap the connection in yamux for multiplexing
	ymConf := yamux.DefaultConfig()
	if opts.LogOutput != nil {
		ymConf.LogOutput = opts.LogOutput
	}
	client, _ := yamux.Client(conn, ymConf)

	// Create the client
	c := &Client{
		conn:   conn,
		client: client,
	}
	return c, nil
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(dc string, addr net.Addr, version int) (*Conn, error) {
	// Get a new, raw connection.
	conn, _, err := p.Dial(dc, addr)
	if err != nil {
		return nil, err
	}

	// Switch the multiplexing based on version
	var session muxSession
	if version < 2 {
		conn.Close()
		return nil, fmt.Errorf("cannot make client connection, unsupported protocol version %d", version)
	} 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
}