git.zd.zone/hrpc/hrpc@v0.0.12/filter/filter.go (about)

     1  package filter
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"google.golang.org/grpc"
     8  )
     9  
    10  // ServerFilter is an interceptor for a service as server role
    11  type ServerFilter struct {
    12  	Name string
    13  	Fn   FilterFunc
    14  }
    15  
    16  // HandlerFunc defines the handler of a function
    17  type HandlerFunc grpc.UnaryHandler
    18  
    19  // FilterFunc defines a new func type that represented the filter function
    20  type FilterFunc func(ctx context.Context, req any, handle HandlerFunc) (any, error)
    21  
    22  func convertToInterceptor(f FilterFunc) grpc.UnaryServerInterceptor {
    23  	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
    24  		return f(ctx, req, HandlerFunc(handler))
    25  	}
    26  }
    27  
    28  var (
    29  	mu            sync.RWMutex
    30  	serverFilters []ServerFilter
    31  )
    32  
    33  // Register can register a filter func to the framework
    34  func Register(name string, fn FilterFunc) {
    35  	mu.Lock()
    36  	defer mu.Unlock()
    37  	serverFilters = append(serverFilters, ServerFilter{
    38  		Name: name,
    39  		Fn:   fn,
    40  	})
    41  }
    42  
    43  // ToInterceptors will convert FilterFuncs to grpc.UnaryServerInterceptors.
    44  func ToInterceptors() []grpc.UnaryServerInterceptor {
    45  	mu.RLock()
    46  	defer mu.RUnlock()
    47  	var out []grpc.UnaryServerInterceptor
    48  	for _, f := range serverFilters {
    49  		out = append(out, convertToInterceptor(f.Fn))
    50  	}
    51  	return out
    52  }