vitess.io/vitess@v0.16.2/go/vt/servenv/grpc_auth.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package servenv
    18  
    19  import (
    20  	"context"
    21  
    22  	"github.com/spf13/pflag"
    23  	"google.golang.org/grpc"
    24  	"google.golang.org/grpc/codes"
    25  	"google.golang.org/grpc/metadata"
    26  	"google.golang.org/grpc/status"
    27  
    28  	"vitess.io/vitess/go/vt/log"
    29  )
    30  
    31  var grpcAuthServerFlagHooks []func(*pflag.FlagSet)
    32  
    33  // RegisterGRPCServerAuthFlags registers flags required to enable server-side
    34  // authentication in vitess gRPC services.
    35  //
    36  // `go/cmd/*` entrypoints should call this function before
    37  // ParseFlags(WithArgs)? if they wish to expose Authenticator functionality.
    38  func RegisterGRPCServerAuthFlags() {
    39  	OnParse(func(fs *pflag.FlagSet) {
    40  		fs.StringVar(&gRPCAuth, "grpc_auth_mode", gRPCAuth, "Which auth plugin implementation to use (eg: static)")
    41  
    42  		for _, fn := range grpcAuthServerFlagHooks {
    43  			fn(fs)
    44  		}
    45  	})
    46  }
    47  
    48  // GRPCAuth returns the value of the `--grpc_auth_mode` flag.
    49  func GRPCAuth() string {
    50  	return gRPCAuth
    51  }
    52  
    53  // Authenticator provides an interface to implement auth in Vitess in
    54  // grpc server
    55  type Authenticator interface {
    56  	Authenticate(ctx context.Context, fullMethod string) (context.Context, error)
    57  }
    58  
    59  // authPlugins is a registry of AuthPlugin initializers.
    60  var authPlugins = make(map[string]func() (Authenticator, error))
    61  
    62  // RegisterAuthPlugin registers an implementation of AuthServer.
    63  func RegisterAuthPlugin(name string, authPlugin func() (Authenticator, error)) {
    64  	if _, ok := authPlugins[name]; ok {
    65  		log.Fatalf("AuthPlugin named %v already exists", name)
    66  	}
    67  	authPlugins[name] = authPlugin
    68  }
    69  
    70  // GetAuthenticator returns an AuthPlugin by name, or log.Fatalf.
    71  func GetAuthenticator(name string) func() (Authenticator, error) {
    72  	authPlugin, ok := authPlugins[name]
    73  	if !ok {
    74  		log.Fatalf("no AuthPlugin name %v registered", name)
    75  	}
    76  	return authPlugin
    77  }
    78  
    79  // FakeAuthStreamInterceptor fake interceptor to test plugin
    80  func FakeAuthStreamInterceptor(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    81  	if fakeDummyAuthenticate(stream.Context()) {
    82  		return handler(srv, stream)
    83  	}
    84  	return status.Errorf(codes.Unauthenticated, "username and password must be provided")
    85  }
    86  
    87  // FakeAuthUnaryInterceptor fake interceptor to test plugin
    88  func FakeAuthUnaryInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
    89  	if fakeDummyAuthenticate(ctx) {
    90  		return handler(ctx, req)
    91  	}
    92  	return nil, status.Errorf(codes.Unauthenticated, "username and password must be provided")
    93  }
    94  
    95  func fakeDummyAuthenticate(ctx context.Context) bool {
    96  	if md, ok := metadata.FromIncomingContext(ctx); ok {
    97  		if len(md["username"]) == 0 || len(md["password"]) == 0 {
    98  			return false
    99  		}
   100  		username := md["username"][0]
   101  		password := md["password"][0]
   102  		if username == "valid" && password == "valid" {
   103  			return true
   104  		}
   105  		return false
   106  	}
   107  	return false
   108  }