Golang yamux.Server函数代码示例

func (c *clientConnection) run() {
	c.log.Debug("starting yamux on ws")

	session, err := yamux.Server(c.websocket, nil)
	if err != nil {
		c.log.Crit("could not initialise yamux session", log.Ctx{"error": err})
	}
	c.log.Debug("yamux session started")
	c.session = session

	c.log.Debug("listening for streams")
	// Accept a stream
	for {
		stream, id, err := c.acceptStream()
		if err != nil {
			if err != io.EOF {
				c.log.Error("error acception stream", log.Ctx{"error": err})
			}
			c.websocket.Close()
			c.session.Close()
			return
		}
		c.log.Debug("accepted stream", log.Ctx{"streamid": id})
		go c.handleStream(stream, id)
	}
}

func testHandshake(t *testing.T, list net.Listener, expect *HandshakeRequest) {
	client, err := list.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer client.Close()

	preamble := make([]byte, len(clientPreamble))
	n, err := client.Read(preamble)
	if err != nil || n != len(preamble) {
		t.Fatalf("err: %v", err)
	}

	server, _ := yamux.Server(client, yamux.DefaultConfig())
	conn, err := server.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer conn.Close()
	rpcCodec := msgpackrpc.NewCodec(true, true, conn)

	rpcSrv := rpc.NewServer()
	rpcSrv.RegisterName("Session", &TestHandshake{t, expect})

	err = rpcSrv.ServeRequest(rpcCodec)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
}

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
}

// ServeConn runs a single connection.
//
// ServeConn blocks, serving the connection until the client hangs up.
func (s *Server) ServeConn(conn io.ReadWriteCloser) {
	// First create the yamux server to wrap this connection
	mux, err := yamux.Server(conn, nil)
	if err != nil {
		conn.Close()
		log.Printf("[ERR] plugin: %s", err)
		return
	}

	// Accept the control connection
	control, err := mux.Accept()
	if err != nil {
		mux.Close()
		log.Printf("[ERR] plugin: %s", err)
		return
	}

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

	// Use the control connection to build the dispenser and serve the
	// connection.
	server := rpc.NewServer()
	server.RegisterName("Dispenser", &dispenseServer{
		ProviderFunc:    s.ProviderFunc,
		ProvisionerFunc: s.ProvisionerFunc,

		broker: broker,
	})
	server.ServeConn(control)
}

// handle multiplex-ed connection
func handleMux(conn io.ReadWriteCloser, target string) {
	// stream multiplex
	var mux *yamux.Session
	config := &yamux.Config{
		AcceptBacklog:          256,
		EnableKeepAlive:        true,
		KeepAliveInterval:      30 * time.Second,
		ConnectionWriteTimeout: 30 * time.Second,
		MaxStreamWindowSize:    16777216,
		LogOutput:              os.Stderr,
	}
	m, err := yamux.Server(conn, config)
	if err != nil {
		log.Println(err)
		return
	}
	mux = m
	defer mux.Close()

	for {
		p1, err := mux.Accept()
		if err != nil {
			log.Println(err)
			return
		}
		p2, err := net.DialTimeout("tcp", target, 5*time.Second)
		if err != nil {
			log.Println(err)
			return
		}
		go handleClient(p1, p2)
	}
}

func downstreamServe(conn net.Conn, downstream host) error {
	session, err := yamux.Server(conn, nil)
	if err != nil {
		return err
	}

	control, err := session.Accept()
	if err != nil {
		return err
	}

	laddr, err := ioutil.ReadAll(control)
	fmt.Println(string(laddr))
	if err != nil {
		return err
	}

	for {
		incoming, err := session.Accept()
		if err != nil {
			return err
		}

		outgoing, err := net.Dial("tcp", downstream.String())
		if err != nil {
			return err
		}

		go handle(proxy(incoming, outgoing))
	}

	return nil
}

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

	return newMuxBroker(s), nil
}

func (y *yamuxer) handleConn(g grim.GrimReaper, conn net.Conn) {
	defer g.Wait()

	conf := yamux.DefaultConfig()
	conf.LogOutput = y.logOutput
	session, _ := yamux.Server(conn, conf)

	streamCh := make(chan net.Conn)
	g.SpawnFunc(processStreams(g.New(), conn, streamCh, y.dispatcher))
	g.SpawnFunc(acceptStreams(y.logger, session, streamCh))
}

func (s *Service) acceptMux(c net.Conn) {
	defer s.wg.Done()

	session, err := yamux.Server(c, muxConfig)
	if err != nil {
		if eofish(err) {
			return
		}

		panic(err)
	}

	defer session.Close()

	acs := make(chan acceptStream, 1)

	debugf("new session for %s\n", c.RemoteAddr())

	data := &clientData{
		parent:     c,
		session:    session,
		inflight:   make(map[MessageId]*Delivery),
		ephemerals: make(map[string]*clientEphemeralInfo),
		done:       make(chan struct{}),
	}

	for {
		go func() {
			stream, err := session.AcceptStream()
			acs <- acceptStream{stream, err}
		}()

		select {
		case <-s.shutdown:
			session.Close()
			s.cleanupConn(c, data)
			return
		case ac := <-acs:
			if ac.err != nil {
				if eofish(ac.err) {
					debugf("eof detected starting a new stream\n")
					s.cleanupConn(c, data)
					return
				}

				panic(ac.err)
			}

			go s.handle(c, ac.stream, data)
		}
	}
}

// ServeConn runs a single connection.
//
// ServeConn blocks, serving the connection until the client hangs up.
func (s *RPCServer) ServeConn(conn io.ReadWriteCloser) {
	// First create the yamux server to wrap this connection
	mux, err := yamux.Server(conn, nil)
	if err != nil {
		conn.Close()
		log.Printf("[ERR] plugin: error creating yamux server: %s", err)
		return
	}

	// Accept the control connection
	control, err := mux.Accept()
	if err != nil {
		mux.Close()
		if err != io.EOF {
			log.Printf("[ERR] plugin: error accepting control connection: %s", err)
		}

		return
	}

	// Connect the stdstreams (in, out, err)
	stdstream := make([]net.Conn, 2)
	for i, _ := range stdstream {
		stdstream[i], err = mux.Accept()
		if err != nil {
			mux.Close()
			log.Printf("[ERR] plugin: accepting stream %d: %s", i, err)
			return
		}
	}

	// Copy std streams out to the proper place
	go copyStream("stdout", stdstream[0], s.Stdout)
	go copyStream("stderr", stdstream[1], s.Stderr)

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

	// Use the control connection to build the dispenser and serve the
	// connection.
	server := rpc.NewServer()
	server.RegisterName("Control", &controlServer{
		server: s,
	})
	server.RegisterName("Dispenser", &dispenseServer{
		broker:  broker,
		plugins: s.Plugins,
	})
	server.ServeConn(control)
}

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)
	}
}

// handleMultiplex handles a multiplexed connection.
func (r *rpcServer) handleMultiplex(conn net.Conn) {
	defer conn.Close()
	conf := yamux.DefaultConfig()
	conf.LogOutput = r.n.config.LogOutput
	server, _ := yamux.Server(conn, conf)
	for {
		sub, err := server.Accept()
		if err != nil {
			if err != io.EOF {
				r.n.logger.Printf("[ERR] onecache.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go r.handleConn(sub)
	}
}

// handleMultiplex is used to multiplex a single incoming connection
// using the Yamux multiplexer
func (s *Server) handleMultiplex(conn net.Conn) {
	defer conn.Close()
	conf := yamux.DefaultConfig()
	conf.LogOutput = s.config.LogOutput
	server, _ := yamux.Server(conn, conf)
	for {
		sub, err := server.Accept()
		if err != nil {
			if err != io.EOF {
				s.logger.Printf("[ERR] nomad.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go s.handleNomadConn(sub)
	}
}

func (li *Listener) getSession() (*yamux.Session, error) {
	li.mu.Lock()
	defer li.mu.Unlock()

	if li.session != nil {
		return li.session, nil
	}

	// connect to the socket master
	conn, err := net.Dial("tcp", li.socketMasterAddress)
	if err != nil {
		return nil, err
	}

	// bind to a port
	err = protocol.WriteHandshakeRequest(conn, protocol.HandshakeRequest{
		SocketDefinition: li.socketDefinition,
	})
	if err != nil {
		conn.Close()
		return nil, err
	}

	// see if that worked
	res, err := protocol.ReadHandshakeResponse(conn)
	if err != nil {
		conn.Close()
		return nil, err
	}
	if res.Status != "OK" {
		conn.Close()
		return nil, fmt.Errorf("%s", res.Status)
	}

	// start a new session
	session, err := yamux.Server(conn, yamux.DefaultConfig())
	if err != nil {
		conn.Close()
		return nil, err
	}

	return session, nil
}