GithubBackup/httprunner
1
0
Fork 0
mirror of https://github.com/httprunner/httprunner.git synced 2026-09-05 07:26:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add interceptors for grpc

Browse Source
This commit is contained in:
徐聪
2022-07-25 21:23:08 +08:00
parent b325232a6c
commit b121282525
16 changed files with 676 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hrp/internal/boomer/client_grpc.go
+100 -10
View File
@@ -8,9 +8,14 @@ import (
"sync/atomic"
"time"
"golang.org/x/oauth2"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/oauth"
"github.com/httprunner/httprunner/v4/hrp/internal/data"
"github.com/httprunner/httprunner/v4/hrp/internal/grpc/messager"
"github.com/rs/zerolog/log"
"google.golang.org/grpc"
)
type grpcClient struct {
@@ -31,14 +36,86 @@ type grpcClient struct {
}
type grpcClientConfig struct {
ctx context.Context
cancel context.CancelFunc // use cancel() to stop client
conn *grpc.ClientConn
biStream messager.Message_BidirectionalStreamingMessageClient
// ctx is used for the lifetime of the stream that may need to be canceled
// on client shutdown.
ctx context.Context
ctxCancel context.CancelFunc
conn *grpc.ClientConn
biStream messager.Message_BidirectionalStreamingMessageClient
mutex sync.RWMutex
}
const token = "httprunner-secret-token"
func logger(format string, a ...interface{}) {
log.Logger.Log().Msg(fmt.Sprintf(format, a...))
}
// unaryInterceptor is an example unary interceptor.
func unaryInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
var credsConfigured bool
for _, o := range opts {
_, ok := o.(grpc.PerRPCCredsCallOption)
if ok {
credsConfigured = true
break
}
}
if !credsConfigured {
opts = append(opts, grpc.PerRPCCredentials(oauth.NewOauthAccess(&oauth2.Token{
AccessToken: token,
})))
}
start := time.Now()
err := invoker(ctx, method, req, reply, cc, opts...)
end := time.Now()
logger("RPC: %s, start time: %s, end time: %s, err: %v", method, start.Format("Basic"), end.Format(time.RFC3339), err)
return err
}
// wrappedStream wraps around the embedded grpc.ClientStream, and intercepts the RecvMsg and
// SendMsg method call.
type wrappedStream struct {
grpc.ClientStream
}
func (w *wrappedStream) RecvMsg(m interface{}) error {
logger("Receive a message (Type: %T) at %v", m, time.Now().Format(time.RFC3339))
return w.ClientStream.RecvMsg(m)
}
func (w *wrappedStream) SendMsg(m interface{}) error {
logger("Send a message (Type: %T) at %v", m, time.Now().Format(time.RFC3339))
return w.ClientStream.SendMsg(m)
}
func newWrappedStream(s grpc.ClientStream) grpc.ClientStream {
return &wrappedStream{s}
}
// streamInterceptor is an example stream interceptor.
func streamInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
var credsConfigured bool
for _, o := range opts {
_, ok := o.(*grpc.PerRPCCredsCallOption)
if ok {
credsConfigured = true
break
}
}
if !credsConfigured {
opts = append(opts, grpc.PerRPCCredentials(oauth.NewOauthAccess(&oauth2.Token{
AccessToken: token,
})))
}
s, err := streamer(ctx, desc, cc, method, opts...)
if err != nil {
return nil, err
}
return newWrappedStream(s), nil
}
func (c *grpcClientConfig) getBiStreamClient() messager.Message_BidirectionalStreamingMessageClient {
c.mutex.RLock()
defer c.mutex.RUnlock()
@@ -64,9 +141,9 @@ func newClient(masterHost string, masterPort int, identity string) (client *grpc
disconnectedFromMaster: make(chan bool),
shutdownChan: make(chan bool),
config: &grpcClientConfig{
ctx: ctx,
cancel: cancel,
mutex: sync.RWMutex{},
ctx: ctx,
ctxCancel: cancel,
mutex: sync.RWMutex{},
},
}
return client
@@ -74,7 +151,20 @@ func newClient(masterHost string, masterPort int, identity string) (client *grpc
func (c *grpcClient) connect() (err error) {
addr := fmt.Sprintf("%v:%v", c.masterHost, c.masterPort)
c.config.conn, err = grpc.Dial(addr, grpc.WithInsecure(), grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(1024*1024*1024)))
// Create tls based credential.
creds, err := credentials.NewClientTLSFromFile(data.Path("x509/ca_cert.pem"), "x.test.example.com")
if err != nil {
log.Fatal().Msg(fmt.Sprintf("failed to load credentials: %v", err))
}
opts := []grpc.DialOption{
// oauth.NewOauthAccess requires the configuration of transport
// credentials.
grpc.WithTransportCredentials(creds),
grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(1024 * 1024 * 1024)),
grpc.WithUnaryInterceptor(unaryInterceptor),
grpc.WithStreamInterceptor(streamInterceptor),
}
c.config.conn, err = grpc.Dial(addr, opts...)
if err != nil {
log.Error().Err(err).Msg("failed to connect")
return err
@@ -112,7 +202,7 @@ func (c *grpcClient) reConnect() (err error) {
func (c *grpcClient) close() {
close(c.shutdownChan)
c.config.cancel()
c.config.ctxCancel()
if c.config.conn != nil {
c.config.conn.Close()
}
hrp/internal/boomer/server_grpc.go
+87 -1
View File
@@ -5,18 +5,93 @@ import (
"fmt"
"io"
"net"
"strings"
"sync"
"sync/atomic"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/reflection"
"google.golang.org/grpc/status"
"github.com/httprunner/httprunner/v4/hrp/internal/data"
"github.com/httprunner/httprunner/v4/hrp/internal/grpc/messager"
"github.com/rs/zerolog/log"
)
var (
errMissingMetadata = status.Errorf(codes.InvalidArgument, "missing metadata")
errInvalidToken = status.Errorf(codes.Unauthenticated, "invalid token")
)
// valid validates the authorization.
func valid(authorization []string) bool {
if len(authorization) < 1 {
return false
}
token := strings.TrimPrefix(authorization[0], "Bearer ")
// Perform the token validation here. For the sake of this example, the code
// here forgoes any of the usual OAuth2 token validation and instead checks
// for a token matching an arbitrary string.
return token == "httprunner-secret-token"
}
func serverUnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
// authentication (token verification)
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return nil, errMissingMetadata
}
if !valid(md["authorization"]) {
return nil, errInvalidToken
}
m, err := handler(ctx, req)
if err != nil {
logger("RPC failed with error %v", err)
}
return m, err
}
// serverWrappedStream wraps around the embedded grpc.ServerStream, and intercepts the RecvMsg and
// SendMsg method call.
type serverWrappedStream struct {
grpc.ServerStream
}
func (w *serverWrappedStream) RecvMsg(m interface{}) error {
logger("Receive a message (Type: %T) at %s", m, time.Now().Format(time.RFC3339))
return w.ServerStream.RecvMsg(m)
}
func (w *serverWrappedStream) SendMsg(m interface{}) error {
logger("Send a message (Type: %T) at %v", m, time.Now().Format(time.RFC3339))
return w.ServerStream.SendMsg(m)
}
func newServerWrappedStream(s grpc.ServerStream) grpc.ServerStream {
return &serverWrappedStream{s}
}
func serverStreamInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
// authentication (token verification)
md, ok := metadata.FromIncomingContext(ss.Context())
if !ok {
return errMissingMetadata
}
if !valid(md["authorization"]) {
return errInvalidToken
}
err := handler(srv, newServerWrappedStream(ss))
if err != nil {
logger("RPC failed with error %v", err)
}
return err
}
func (s *grpcServer) BidirectionalStreamingMessage(srv messager.Message_BidirectionalStreamingMessageServer) error {
s.wg.Add(1)
defer s.wg.Done()
@@ -158,13 +233,24 @@ func newServer(masterHost string, masterPort int) (server *grpcServer) {
func (s *grpcServer) start() (err error) {
addr := fmt.Sprintf("%v:%v", s.masterHost, s.masterPort)
// Create tls based credential.
creds, err := credentials.NewServerTLSFromFile(data.Path("x509/server_cert.pem"), data.Path("x509/server_key.pem"))
if err != nil {
log.Fatal().Msg(fmt.Sprintf("failed to load key pair: %s", err))
}
opts := []grpc.ServerOption{
grpc.UnaryInterceptor(serverUnaryInterceptor),
grpc.StreamInterceptor(serverStreamInterceptor),
// Enable TLS for all incoming connections.
grpc.Creds(creds),
}
lis, err := net.Listen("tcp", addr)
if err != nil {
log.Error().Err(err).Msg("failed to listen")
return
}
// create gRPC server
serv := grpc.NewServer()
serv := grpc.NewServer(opts...)
// register message server
messager.RegisterMessageServer(serv, s)
reflection.Register(serv)