Add context parameter to GetToken

Author: everesio

cmd/plugin-googleid-provider/main.go

@@ -17,11 +17,11 @@ type TokenProvider struct {
 }
 
 //TODO: caching, expiry
-//TODO: refresh in the half of time
-//TODO: send claims
-func (p TokenProvider) GetToken(request apis.TokenRequest) (apis.TokenResponse, error) {
-
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(p.timeout)*time.Second)
+//TODO: go routine - on startup, refresh in the half of time (back off - > use the backoff lib)
+//TODO: read / write mutex
+//TODO: const status codes + prometheus metrics
+func (p TokenProvider) GetToken(parent context.Context, request apis.TokenRequest) (apis.TokenResponse, error) {
+	ctx, cancel := context.WithTimeout(parent, time.Duration(p.timeout)*time.Second)
 	defer cancel()
 
 	tokenSource, err := google.DefaultTokenSource(ctx, oauth2.UserinfoEmailScope)
@@ -52,7 +52,7 @@ func (f *TokenProvider) flagSet() *flag.FlagSet {
 func main() {
 	tokenProvider := &TokenProvider{}
 	fs := tokenProvider.flagSet()
-	fs.IntVar(&tokenProvider.timeout, "timeout", 5, "Request timeout")
+	fs.IntVar(&tokenProvider.timeout, "timeout", 5, "Request timeout in seconds")
 
 	fs.Parse(os.Args[1:])
 

pkg/apis/gateway.go

@@ -1,5 +1,9 @@
 package apis
 
+import (
+	"context"
+)
+
 type TokenRequest struct {
 	Params []string
 }
@@ -10,5 +14,5 @@ type TokenResponse struct {
 }
 
 type TokenProvider interface {
-	GetToken(request TokenRequest) (TokenResponse, error)
+	GetToken(ctx context.Context, request TokenRequest) (TokenResponse, error)
 }

plugin/gateway-client/proto/token-provider.proto

@@ -9,7 +9,6 @@ message TokenResponse {
     bool success = 1;
     int32 status = 2;
     string token = 3;
-    // TODO: expiry ?
 }
 
 service TokenProvider {

plugin/gateway-client/shared/grpc.go

@@ -13,8 +13,8 @@ type GRPCClient struct {
 	client proto.TokenProviderClient
 }
 
-func (m *GRPCClient) GetToken(request apis.TokenRequest) (apis.TokenResponse, error) {
-	resp, err := m.client.GetToken(context.Background(), &proto.TokenRequest{Params: request.Params})
+func (m *GRPCClient) GetToken(ctx context.Context, request apis.TokenRequest) (apis.TokenResponse, error) {
+	resp, err := m.client.GetToken(ctx, &proto.TokenRequest{Params: request.Params})
 	return apis.TokenResponse{Success: resp.Success, Status: resp.Status, Token: resp.Token}, err
 }
 
@@ -27,6 +27,6 @@ type GRPCServer struct {
 func (m *GRPCServer) GetToken(
 	ctx context.Context,
 	req *proto.TokenRequest) (*proto.TokenResponse, error) {
-	resp, err := m.Impl.GetToken(apis.TokenRequest{Params: req.Params})
+	resp, err := m.Impl.GetToken(ctx, apis.TokenRequest{Params: req.Params})
 	return &proto.TokenResponse{Success: resp.Success, Status: resp.Status, Token: resp.Token}, err
 }

plugin/gateway-client/shared/rpc.go

@@ -3,6 +3,7 @@ package shared
 import (
 	"github.com/grepplabs/kafka-proxy/pkg/apis"
 	"net/rpc"
+	"context"
 )
 
 type RPCClient struct{ client *rpc.Client }
@@ -21,7 +22,7 @@ type RPCServer struct {
 
 func (m *RPCServer) GetToken(args map[string]interface{}, resp *map[string]interface{}) error {
 
-	r, err := m.Impl.GetToken(apis.TokenRequest{Params: args["params"].([]string)})
+	r, err := m.Impl.GetToken(context.Background(), apis.TokenRequest{Params: args["params"].([]string)})
 	*resp = map[string]interface{}{
 		"success": r.Success,
 		"status":  r.Status,

proxy/auth.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"strings"
 	"time"
+	"context"
 )
 
 type AuthClient struct {
@@ -23,7 +24,12 @@ type AuthClient struct {
 //TODO: reset deadlines after method - ok
 func (b *AuthClient) sendAndReceiveGatewayAuth(conn DeadlineReaderWriter) error {
 	//TODO: retrieve from plugin (with timeout)
-	resp, err := b.tokenProvider.GetToken(apis.TokenRequest{})
+	/*
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(p.timeout)*time.Second)
+	defer cancel()
+   */
+
+	resp, err := b.tokenProvider.GetToken(context.Background(), apis.TokenRequest{})
 	if err != nil {
 		return err
 	}