github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/auth/auth.go (about)

     1  package auth
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/collection"
     8  	"github.com/lingyao2333/mo-zero/core/stores/redis"
     9  	"google.golang.org/grpc/codes"
    10  	"google.golang.org/grpc/metadata"
    11  	"google.golang.org/grpc/status"
    12  )
    13  
    14  const defaultExpiration = 5 * time.Minute
    15  
    16  // An Authenticator is used to authenticate the rpc requests.
    17  type Authenticator struct {
    18  	store  *redis.Redis
    19  	key    string
    20  	cache  *collection.Cache
    21  	strict bool
    22  }
    23  
    24  // NewAuthenticator returns an Authenticator.
    25  func NewAuthenticator(store *redis.Redis, key string, strict bool) (*Authenticator, error) {
    26  	cache, err := collection.NewCache(defaultExpiration)
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  
    31  	return &Authenticator{
    32  		store:  store,
    33  		key:    key,
    34  		cache:  cache,
    35  		strict: strict,
    36  	}, nil
    37  }
    38  
    39  // Authenticate authenticates the given ctx.
    40  func (a *Authenticator) Authenticate(ctx context.Context) error {
    41  	md, ok := metadata.FromIncomingContext(ctx)
    42  	if !ok {
    43  		return status.Error(codes.Unauthenticated, missingMetadata)
    44  	}
    45  
    46  	apps, tokens := md[appKey], md[tokenKey]
    47  	if len(apps) == 0 || len(tokens) == 0 {
    48  		return status.Error(codes.Unauthenticated, missingMetadata)
    49  	}
    50  
    51  	app, token := apps[0], tokens[0]
    52  	if len(app) == 0 || len(token) == 0 {
    53  		return status.Error(codes.Unauthenticated, missingMetadata)
    54  	}
    55  
    56  	return a.validate(app, token)
    57  }
    58  
    59  func (a *Authenticator) validate(app, token string) error {
    60  	expect, err := a.cache.Take(app, func() (interface{}, error) {
    61  		return a.store.Hget(a.key, app)
    62  	})
    63  	if err != nil {
    64  		if a.strict {
    65  			return status.Error(codes.Internal, err.Error())
    66  		}
    67  
    68  		return nil
    69  	}
    70  
    71  	if token != expect {
    72  		return status.Error(codes.Unauthenticated, accessDenied)
    73  	}
    74  
    75  	return nil
    76  }