Golang bsonrpc.ServeCustomRPC函數代碼示例

// ServerSecurePort obtains a listener that accepts secure connections.
// If the provided port is zero, the listening is disabled.
func ServeSecurePort(securePort int, certFile, keyFile, caCertFile string) {
	if securePort == 0 {
		log.Info("Not listening on secure port")
		return
	}

	config := tls.Config{}

	// load the server cert / key
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Fatalf("SecureServe.LoadX509KeyPair(%v, %v) failed: %v", certFile, keyFile, err)
	}
	config.Certificates = []tls.Certificate{cert}

	// load the ca if necessary
	// FIXME(alainjobart) this doesn't quite work yet, have
	// to investigate
	if caCertFile != "" {
		config.ClientCAs = x509.NewCertPool()

		pemCerts, err := ioutil.ReadFile(caCertFile)
		if err != nil {
			log.Fatalf("SecureServe: cannot read ca file %v: %v", caCertFile, err)
		}
		if !config.ClientCAs.AppendCertsFromPEM(pemCerts) {
			log.Fatalf("SecureServe: AppendCertsFromPEM failed: %v", err)
		}

		config.ClientAuth = tls.RequireAndVerifyClientCert
	}
	l, err := tls.Listen("tcp", fmt.Sprintf(":%d", securePort), &config)
	if err != nil {
		log.Fatalf("Error listening on secure port %v: %v", securePort, err)
	}
	log.Infof("Listening on secure port %v", securePort)
	throttled := NewThrottledListener(l, *secureThrottle, *secureMaxBuffer)
	cl := proc.Published(throttled, "SecureConnections", "SecureAccepts")

	// rpc.HandleHTTP registers the default GOB handler at /_goRPC_
	// and the debug RPC service at /debug/rpc (it displays a list
	// of registered services and their methods).
	if ServiceMap["gob-vts"] {
		log.Infof("Registering GOB handler and /debug/rpc URL for vts port")
		secureRpcServer.HandleHTTP(rpcwrap.GetRpcPath("gob", false), rpcplus.DefaultDebugPath)
	}
	if ServiceMap["gob-auth-vts"] {
		log.Infof("Registering GOB handler and /debug/rpcs URL for SASL vts port")
		authenticatedSecureRpcServer.HandleHTTP(rpcwrap.GetRpcPath("gob", true), rpcplus.DefaultDebugPath+"s")
	}

	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, secureRpcServer, false)
	bsonrpc.ServeCustomRPC(handler, authenticatedSecureRpcServer, true)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(cl)
}

// the test here creates a fake server implementation, a fake client
// implementation, and runs the test suite against the setup.
func TestGoRPCBinlogStreamer(t *testing.T) {
	// Listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	fakeUpdateStream := binlogplayertest.NewFakeBinlogStreamer(t)
	server.Register(gorpcbinlogstreamer.New(fakeUpdateStream))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server, false)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a Go Rpc client to talk to the fake tablet
	client := &GoRpcBinlogPlayerClient{}

	// and send it to the test suite
	binlogplayertest.Run(t, client, listener.Addr().String(), fakeUpdateStream)
}

// StartActionLoop will start the action loop for a fake tablet,
// using ft.FakeMysqlDaemon as the backing mysqld.
func (ft *FakeTablet) StartActionLoop(t *testing.T, wr *wrangler.Wrangler) {
	if ft.Agent != nil {
		t.Fatalf("Agent for %v is already running", ft.Tablet.Alias)
	}

	// Listen on a random port
	var err error
	ft.Listener, err = net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}
	port := ft.Listener.Addr().(*net.TCPAddr).Port

	// create a test agent on that port, and re-read the record
	// (it has new ports and IP)
	ft.Agent = tabletmanager.NewTestActionAgent(context.Background(), wr.TopoServer(), ft.Tablet.Alias, port, ft.FakeMysqlDaemon)
	ft.Tablet = ft.Agent.Tablet().Tablet

	// create the RPC server
	ft.RPCServer = rpcplus.NewServer()
	gorpctmserver.RegisterForTest(ft.RPCServer, ft.Agent)

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, ft.RPCServer, false)
	ft.HTTPServer = http.Server{
		Handler: handler,
	}
	go ft.HTTPServer.Serve(ft.Listener)
}

// This test makes sure the go rpc service works
func TestGoRPCVTGateConn(t *testing.T) {
	// fake service
	service := vtgateconntest.CreateFakeServer(t)

	// listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcvtgateservice.New(service))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server, false)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a Go RPC client connecting to the server
	ctx := context.Background()
	client, err := dial(ctx, listener.Addr().String(), 30*time.Second)
	if err != nil {
		t.Fatalf("dial failed: %v", err)
	}

	// run the test suite
	vtgateconntest.TestSuite(t, client, service)

	// and clean up
	client.Close()
}

func TestMain(m *testing.M) {
	// fake service
	service := CreateFakeServer()

	// listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		panic(fmt.Sprintf("Cannot listen: %v", err))
	}

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcvtgateservice.New(service))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server, false)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	testAddress = listener.Addr().String()
	os.Exit(m.Run())
}

// TestGoRPCGoClient tests the go client using goRPC
func TestGoRPCGoClient(t *testing.T) {
	service := createService()

	// listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}
	defer listener.Close()

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcvtgateservice.New(service))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server, false)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// and run the test suite
	testGoClient(t, "gorpc", listener.Addr().String())
}

// the test here creates a fake server implementation, a fake client
// implementation, and runs the test suite against the setup.
func TestGoRPCBinlogStreamer(t *testing.T) {
	// Listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}
	host := listener.Addr().(*net.TCPAddr).IP.String()
	port := listener.Addr().(*net.TCPAddr).Port

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	fakeUpdateStream := binlogplayertest.NewFakeBinlogStreamer(t)
	server.Register(gorpcbinlogstreamer.New(fakeUpdateStream))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a Go Rpc client to talk to the fake tablet
	c := &client{}

	// and send it to the test suite
	binlogplayertest.Run(t, c, &pb.EndPoint{
		Host: host,
		PortMap: map[string]int32{
			"vt": int32(port),
		},
	}, fakeUpdateStream)
}

// the test here creates a fake server implementation, a fake client
// implementation, and runs the test suite against the setup.
func TestVtctlServer(t *testing.T) {
	ts := vtctlclienttest.CreateTopoServer(t)

	// Listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcvtctlserver.NewVtctlServer(ts))

	// Create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a VtctlClient Go Rpc client to talk to the fake server
	client, err := goRPCVtctlClientFactory(listener.Addr().String(), 30*time.Second)
	if err != nil {
		t.Fatalf("Cannot create client: %v", err)
	}
	defer client.Close()

	vtctlclienttest.TestSuite(t, ts, client)
}

// This test makes sure the go rpc service works
func testGoRPCTabletConn(t *testing.T) {
	// fake service
	service := tabletconntest.CreateFakeServer(t)

	// listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}
	defer listener.Close()
	port := listener.Addr().(*net.TCPAddr).Port

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcqueryservice.New(service))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// run the test suite
	tabletconntest.TestSuite(t, protocolName, &pb.EndPoint{
		Host: "localhost",
		PortMap: map[string]int32{
			"vt": int32(port),
		},
	}, service)
}

// NewVtctlPipe creates a new VtctlPipe based on the given topo server.
func NewVtctlPipe(t *testing.T, ts topo.Server) *VtctlPipe {
	// Listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcvtctlserver.NewVtctlServer(ts))

	// Create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server, false)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a VtctlClient Go Rpc client to talk to the fake server
	client, err := vtctlclient.New(listener.Addr().String(), 30*time.Second)
	if err != nil {
		t.Fatalf("Cannot create client: %v", err)
	}

	return &VtctlPipe{
		listener: listener,
		client:   client,
		t:        t,
	}
}

// ServeSocketFile listen to the named socket and serves RPCs on it.
func ServeSocketFile(name string) {
	if name == "" {
		log.Infof("Not listening on socket file")
		return
	}

	// try to delete if file exists
	if _, err := os.Stat(name); err == nil {
		err = os.Remove(name)
		if err != nil {
			log.Fatalf("Cannot remove socket file %v: %v", name, err)
		}
	}

	l, err := net.Listen("unix", name)
	if err != nil {
		log.Fatalf("Error listening on socket file %v: %v", name, err)
	}
	log.Infof("Listening on socket file %v", name)

	// HandleHTTP registers the default GOB handler at /_goRPC_
	// and the debug RPC service at /debug/rpc (it displays a list
	// of registered services and their methods).
	if ServiceMap["gob-unix"] {
		log.Infof("Registering GOB handler and /debug/rpc URL for unix socket")
		socketFileRpcServer.HandleHTTP(rpcwrap.GetRpcPath("gob", false), rpcplus.DefaultDebugPath)
	}
	if ServiceMap["gob-auth-unix"] {
		log.Infof("Registering GOB handler and /debug/rpcs URL for SASL unix socket")
		authenticatedSocketFileRpcServer.HandleHTTP(rpcwrap.GetRpcPath("gob", true), rpcplus.DefaultDebugPath+"s")
	}

	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, socketFileRpcServer, false)
	bsonrpc.ServeCustomRPC(handler, authenticatedSocketFileRpcServer, true)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(l)
}

// This test makes sure the go rpc service works
func testGoRPCTabletConn(t *testing.T, rpcOnlyInReply bool) {
	// fake service
	service := tabletconntest.CreateFakeServer(t)

	// listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}
	defer listener.Close()
	port := listener.Addr().(*net.TCPAddr).Port

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcqueryservice.New(service))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server, false)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)
	// Handle errors appropriately
	*tabletserver.RPCErrorOnlyInReply = rpcOnlyInReply

	// Create a Go RPC client connecting to the server
	ctx := context.Background()
	client, err := DialTablet(ctx, topo.EndPoint{
		Host: "localhost",
		NamedPortMap: map[string]int{
			"vt": port,
		},
	}, tabletconntest.TestKeyspace, tabletconntest.TestShard, 30*time.Second)
	if err != nil {
		t.Fatalf("dial failed: %v", err)
	}

	// run the test suite
	tabletconntest.TestSuite(t, client, service)

	// and clean up
	client.Close()
}

// TestGoRPCVTGateConn makes sure the gorpc (BsonRPC) service works
func TestGoRPCVTGateConn(t *testing.T) {
	// fake service
	service := vtgateconntest.CreateFakeServer(t)

	// listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcvtgateservice.New(service))
	// TODO(aaijazi): remove this flag once all VtGate Gorpc clients properly support the new behavior
	*vtgate.RPCErrorOnlyInReply = true

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a Go RPC client connecting to the server
	ctx := context.Background()
	client, err := dial(ctx, listener.Addr().String(), 30*time.Second)
	if err != nil {
		t.Fatalf("dial failed: %v", err)
	}
	vtgateconntest.RegisterTestDialProtocol(client)

	// run the test suite
	vtgateconntest.TestSuite(t, client, service)
	vtgateconntest.TestErrorSuite(t, service)

	// and clean up
	client.Close()
}

// TestGoRPCVTGateConn makes sure the gorpc (BsonRPC) service works
func TestGoRPCVTGateConn(t *testing.T) {
	// fake service
	service := vtgateconntest.CreateFakeServer(t)

	// listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	server.Register(gorpcvtgateservice.New(service))

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a Go RPC client connecting to the server
	ctx := context.Background()
	client, err := dial(ctx, listener.Addr().String(), 30*time.Second)
	if err != nil {
		t.Fatalf("dial failed: %v", err)
	}
	vtgateconntest.RegisterTestDialProtocol(client)

	// run the test suite
	vtgateconntest.TestSuite(t, client, service)
	// TODO(sougou/alainjobart): find out why this was commmented out for grpc.
	// This is now failing for gorpc also.
	//vtgateconntest.TestErrorSuite(t, service)

	// and clean up
	client.Close()
}

// the test here creates a fake server implementation, a fake client
// implementation, and runs the test suite against the setup.
func TestGoRPCTMServer(t *testing.T) {
	// Listen on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}
	defer listener.Close()
	port := int32(listener.Addr().(*net.TCPAddr).Port)

	// Create a Go Rpc server and listen on the port
	server := rpcplus.NewServer()
	fakeAgent := agentrpctest.NewFakeRPCAgent(t)
	server.Register(&TabletManager{fakeAgent})

	// create the HTTP server, serve the server from it
	handler := http.NewServeMux()
	bsonrpc.ServeCustomRPC(handler, server)
	httpServer := http.Server{
		Handler: handler,
	}
	go httpServer.Serve(listener)

	// Create a Go Rpc client to talk to the fake tablet
	client := &gorpctmclient.GoRPCTabletManagerClient{}
	ti := topo.NewTabletInfo(&pb.Tablet{
		Alias: &pb.TabletAlias{
			Cell: "test",
			Uid:  123,
		},
		Hostname: "localhost",
		PortMap: map[string]int32{
			"vt": port,
		},
	}, 0)

	// and run the test suite
	agentrpctest.Run(t, client, ti, fakeAgent)
}