auth.go 1.5 KB

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