/
interceptor.go
34 lines (29 loc) · 1022 Bytes
/
interceptor.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
package grpc
import (
"context"
grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
"google.golang.org/grpc"
)
type Interceptor = func(ctx context.Context, method string) (context.Context, error)
func UnaryServerInterceptor(intercept Interceptor) grpc.ServerOption {
return grpc.UnaryInterceptor(
func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
newCtx, err := intercept(ctx, info.FullMethod)
if err != nil {
return nil, err
}
return handler(newCtx, req)
})
}
func StreamServerInterceptor(intercept Interceptor) grpc.ServerOption {
return grpc.StreamInterceptor(
func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
newCtx, err := intercept(stream.Context(), info.FullMethod)
if err != nil {
return err
}
wrapped := grpc_middleware.WrapServerStream(stream)
wrapped.WrappedContext = newCtx
return handler(srv, wrapped)
})
}