github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/conf/jwt/jwt_auth.go (about)

     1  package jwt
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  
     7  	"github.com/pkg/errors"
     8  
     9  	"github.com/machinefi/w3bstream/pkg/depends/x/contextx"
    10  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/must"
    11  )
    12  
    13  type Auth struct {
    14  	AuthInQuery  string `name:"authorization,omitempty" in:"query"  validate:"@string[1,]"`
    15  	AuthInHeader string `name:"Authorization,omitempty" in:"header" validate:"@string[1,]"`
    16  }
    17  
    18  func (r Auth) ContextKey() interface{} { return keyAuth{} }
    19  
    20  func (r Auth) Output(ctx context.Context) (pl interface{}, err error) {
    21  	jwt, ok := ConfFromContext(ctx)
    22  	if !ok {
    23  		return nil, ErrEmptyJwtContext
    24  	}
    25  
    26  	av := r.AuthInQuery
    27  	if av == "" {
    28  		av = r.AuthInHeader
    29  	}
    30  	tok := strings.TrimSpace(strings.Replace(av, "Bearer", " ", 1))
    31  
    32  	if WithAnonymousPublisherFn != nil {
    33  		tok, err = WithAnonymousPublisherFn(ctx, tok)
    34  	}
    35  
    36  	ok = false
    37  	if BuiltInTokenValidateFn != nil {
    38  		pl, err, ok = BuiltInTokenValidateFn(ctx, tok)
    39  	}
    40  	if !ok {
    41  		var claims *Claims
    42  		if claims, err = jwt.ParseToken(tok); err == nil {
    43  			pl = claims.Payload
    44  		}
    45  	}
    46  
    47  	if err != nil {
    48  		return nil, err
    49  	}
    50  
    51  	if WithPermissionFn != nil && !WithPermissionFn(pl) {
    52  		return nil, ErrNoPermission
    53  	}
    54  
    55  	return
    56  }
    57  
    58  type keyConf struct{}
    59  
    60  func WithConfContext(jwt *Jwt) func(context.Context) context.Context {
    61  	return func(ctx context.Context) context.Context {
    62  		return contextx.WithValue(ctx, keyConf{}, jwt)
    63  	}
    64  }
    65  
    66  func ConfFromContext(ctx context.Context) (*Jwt, bool) {
    67  	j, ok := ctx.Value(keyConf{}).(*Jwt)
    68  	return j, ok
    69  }
    70  
    71  func MustConfFromContext(ctx context.Context) *Jwt {
    72  	j, ok := ctx.Value(keyConf{}).(*Jwt)
    73  	must.BeTrue(ok)
    74  	return j
    75  }
    76  
    77  type keyAuth struct{}
    78  
    79  func AuthFromContext(ctx context.Context) interface{} {
    80  	return ctx.Value(keyAuth{})
    81  }
    82  
    83  var (
    84  	ErrEmptyJwtContext = errors.New("empty jwt context")
    85  	ErrNoPermission    = errors.New("no permission")
    86  )
    87  
    88  var BuiltInTokenValidateFn func(context.Context, string) (interface{}, error, bool)
    89  
    90  func SetBuiltInTokenFn(f func(context.Context, string) (interface{}, error, bool)) {
    91  	BuiltInTokenValidateFn = f
    92  }
    93  
    94  var WithPermissionFn func(interface{}) bool
    95  
    96  func SetWithPermissionFn(f func(interface{}) bool) {
    97  	WithPermissionFn = f
    98  }
    99  
   100  var WithAnonymousPublisherFn func(context.Context, string) (string, error)
   101  
   102  func SetWithAnonymousPublisherFn(f func(context.Context, string) (string, error)) {
   103  	WithAnonymousPublisherFn = f
   104  }