Refactor GetToken

Author: everesio

cmd/plugin-googleid-provider/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"flag"
+	"github.com/grepplabs/kafka-proxy/pkg/apis"
 	"github.com/grepplabs/kafka-proxy/plugin/gateway-client/shared"
 	"github.com/hashicorp/go-plugin"
 	"golang.org/x/oauth2/google"
@@ -18,29 +19,29 @@ type TokenProvider struct {
 //TODO: caching, expiry
 //TODO: refresh in the half of time
 //TODO: send claims
-func (p TokenProvider) GetToken(claims []string) (int32, string, error) {
+func (p TokenProvider) GetToken(request apis.TokenRequest) (apis.TokenResponse, error) {
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(p.timeout)*time.Second)
 	defer cancel()
 
 	tokenSource, err := google.DefaultTokenSource(ctx, oauth2.UserinfoEmailScope)
 	if err != nil {
-		return tokenResponse(1, "")
+		return tokenResponse(false, 1, "")
 	}
 	token, err := tokenSource.Token()
 	if err != nil {
-		return tokenResponse(2, "")
+		return tokenResponse(false, 2, "")
 	}
 	if token.Extra("id_token") == nil {
-		return tokenResponse(3, "")
+		return tokenResponse(false, 3, "")
 	}
 	idToken := token.Extra("id_token").(string)
-	return tokenResponse(0, idToken)
+	return tokenResponse(true, 0, idToken)
 
 }
 
-func tokenResponse(status int32, token string) (int32, string, error) {
-	return status, token, nil
+func tokenResponse(success bool, status int32, token string) (apis.TokenResponse, error) {
+	return apis.TokenResponse{Success: success, Status: status, Token: token}, nil
 }
 
 func (f *TokenProvider) flagSet() *flag.FlagSet {

pkg/apis/gateway.go

@@ -1,5 +1,14 @@
 package apis
 
+type TokenRequest struct {
+	Params []string
+}
+type TokenResponse struct {
+	Success bool
+	Status  int32
+	Token   string
+}
+
 type TokenProvider interface {
-	GetToken(claims []string) (int32, string, error)
+	GetToken(request TokenRequest) (TokenResponse, error)
 }

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

@@ -2,12 +2,13 @@ syntax = "proto3";
 package proto;
 
 message TokenRequest {
-    repeated string claims = 1;
+    repeated string params = 1;
 }
 
 message TokenResponse {
-    int32 status = 1;
-    string token = 2;
+    bool success = 1;
+    int32 status = 2;
+    string token = 3;
     // TODO: expiry ?
 }
 

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

@@ -13,12 +13,9 @@ type GRPCClient struct {
 	client proto.TokenProviderClient
 }
 
-func (m *GRPCClient) GetToken(claims []string) (int32, string, error) {
-	resp, err := m.client.GetToken(context.Background(), &proto.TokenRequest{Claims: claims})
-	if err != nil {
-		return 0, "", err
-	}
-	return resp.Status, resp.Token, nil
+func (m *GRPCClient) GetToken(request apis.TokenRequest) (apis.TokenResponse, error) {
+	resp, err := m.client.GetToken(context.Background(), &proto.TokenRequest{Params: request.Params})
+	return apis.TokenResponse{Success: resp.Success, Status: resp.Status, Token: resp.Token}, err
 }
 
 // Here is the gRPC server that GRPCClient talks to.
@@ -30,6 +27,6 @@ type GRPCServer struct {
 func (m *GRPCServer) GetToken(
 	ctx context.Context,
 	req *proto.TokenRequest) (*proto.TokenResponse, error) {
-	s, t, err := m.Impl.GetToken(req.Claims)
-	return &proto.TokenResponse{Status: s, Token: t}, err
+	resp, err := m.Impl.GetToken(apis.TokenRequest{Params: req.Params})
+	return &proto.TokenResponse{Success: resp.Success, Status: resp.Status, Token: resp.Token}, err
 }

plugin/gateway-client/shared/grpc.go

@@ -7,23 +7,25 @@ import (
 
 type RPCClient struct{ client *rpc.Client }
 
-func (m *RPCClient) GetToken(claims []string) (int32, string, error) {
+func (m *RPCClient) GetToken(request apis.TokenRequest) (apis.TokenResponse, error) {
 	var resp map[string]interface{}
 	err := m.client.Call("Plugin.GetToken", map[string]interface{}{
-		"claims": claims,
+		"params": request.Params,
 	}, &resp)
-	return resp["status"].(int32), resp["token"].(string), err
+	return apis.TokenResponse{Success: resp["success"].(bool), Status: resp["status"].(int32), Token: resp["token"].(string)}, err
 }
 
 type RPCServer struct {
 	Impl apis.TokenProvider
 }
 
 func (m *RPCServer) GetToken(args map[string]interface{}, resp *map[string]interface{}) error {
-	s, t, err := m.Impl.GetToken(args["claims"].([]string))
+
+	r, err := m.Impl.GetToken(apis.TokenRequest{Params: args["params"].([]string)})
 	*resp = map[string]interface{}{
-		"status": s,
-		"token":  t,
+		"success": r.Success,
+		"status":  r.Status,
+		"token":   r.Token,
 	}
 	return err
 }

plugin/gateway-client/shared/rpc.go

@@ -23,16 +23,18 @@ type AuthClient struct {
 //TODO: reset deadlines after method - ok
 func (b *AuthClient) sendAndReceiveGatewayAuth(conn DeadlineReaderWriter) error {
 	//TODO: retrieve from plugin (with timeout)
-	status, data, err := b.tokenProvider.GetToken([]string{})
+	resp, err := b.tokenProvider.GetToken(apis.TokenRequest{})
 	if err != nil {
 		return err
 	}
-	if status != 0 {
-		return fmt.Errorf("get token failed with status: %d", status)
+	if !resp.Success {
+		return fmt.Errorf("get token unsuccessful with status: %d", resp.Status)
 	}
-	if data == "" {
-		return errors.New("received empty token")
+	if resp.Token == "" {
+		return errors.New("get token returned empty token")
 	}
+	data := resp.Token
+	logrus.Info(data)
 
 	length := len(b.method) + 1 + len(data)
 	// 8 - bytes magic, 4 bytes length