client.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/clientpool

     1  // Copyright 2020 Chaos Mesh Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package clientpool
    15  
    16  import (
    17  	"errors"
    18  	"net/http"
    19  	"strings"
    20  	"sync"
    21  
    22  	lru "github.com/hashicorp/golang-lru"
    23  	"k8s.io/apimachinery/pkg/runtime"
    24  	authorizationv1 "k8s.io/client-go/kubernetes/typed/authorization/v1"
    25  	"k8s.io/client-go/rest"
    26  	pkgclient "sigs.k8s.io/controller-runtime/pkg/client"
    27  
    28  	"github.com/chaos-mesh/chaos-mesh/pkg/mock"
    29  )
    30  
    31  // K8sClients is an object of Clients
    32  var K8sClients Clients
    33  
    34  type Clients interface {
    35  	Client(token string) (pkgclient.Client, error)
    36  	AuthClient(token string) (authorizationv1.AuthorizationV1Interface, error)
    37  	Num() int
    38  	Contains(token string) bool
    39  }
    40  
    41  type LocalClient struct {
    42  	client     pkgclient.Client
    43  	authClient authorizationv1.AuthorizationV1Interface
    44  }
    45  
    46  func NewLocalClient(localConfig *rest.Config, scheme *runtime.Scheme) (Clients, error) {
    47  	client, err := pkgclient.New(localConfig, pkgclient.Options{
    48  		Scheme: scheme,
    49  	})
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	authCli, err := authorizationv1.NewForConfig(localConfig)
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  
    59  	return &LocalClient{
    60  		client:     client,
    61  		authClient: authCli,
    62  	}, nil
    63  }
    64  
    65  // Client returns the local k8s client
    66  func (c *LocalClient) Client(token string) (pkgclient.Client, error) {
    67  	return c.client, nil
    68  }
    69  
    70  func (c *LocalClient) AuthClient(token string) (authorizationv1.AuthorizationV1Interface, error) {
    71  	return c.authClient, nil
    72  }
    73  
    74  // Num returns the num of clients
    75  func (c *LocalClient) Num() int {
    76  	return 1
    77  }
    78  
    79  // Contains return false for LocalClient
    80  func (c *LocalClient) Contains(token string) bool {
    81  	return false
    82  }
    83  
    84  // Clients is the client pool of k8s client
    85  type ClientsPool struct {
    86  	sync.RWMutex
    87  
    88  	scheme      *runtime.Scheme
    89  	localConfig *rest.Config
    90  	clients     *lru.Cache
    91  	authClients *lru.Cache
    92  }
    93  
    94  // New creates a new Clients
    95  func NewClientPool(localConfig *rest.Config, scheme *runtime.Scheme, maxClientNum int) (Clients, error) {
    96  	clients, err := lru.New(maxClientNum)
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  
   101  	authClients, err := lru.New(maxClientNum)
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  
   106  	return &ClientsPool{
   107  		localConfig: localConfig,
   108  		scheme:      scheme,
   109  		clients:     clients,
   110  		authClients: authClients,
   111  	}, nil
   112  }
   113  
   114  // Client returns a k8s client according to the token
   115  func (c *ClientsPool) Client(token string) (pkgclient.Client, error) {
   116  	c.Lock()
   117  	defer c.Unlock()
   118  
   119  	if len(token) == 0 {
   120  		return nil, errors.New("token is empty")
   121  	}
   122  
   123  	value, ok := c.clients.Get(token)
   124  	if ok {
   125  		return value.(pkgclient.Client), nil
   126  	}
   127  
   128  	config := rest.CopyConfig(c.localConfig)
   129  	config.BearerToken = token
   130  	config.BearerTokenFile = ""
   131  
   132  	newFunc := pkgclient.New
   133  
   134  	if mockNew := mock.On("MockCreateK8sClient"); mockNew != nil {
   135  		newFunc = mockNew.(func(config *rest.Config, options pkgclient.Options) (pkgclient.Client, error))
   136  	}
   137  
   138  	client, err := newFunc(config, pkgclient.Options{
   139  		Scheme: c.scheme,
   140  	})
   141  	if err != nil {
   142  		return nil, err
   143  	}
   144  
   145  	_ = c.clients.Add(token, client)
   146  
   147  	return client, nil
   148  }
   149  
   150  func (c *ClientsPool) AuthClient(token string) (authorizationv1.AuthorizationV1Interface, error) {
   151  	c.Lock()
   152  	defer c.Unlock()
   153  
   154  	if len(token) == 0 {
   155  		return nil, errors.New("token is empty")
   156  	}
   157  
   158  	value, ok := c.authClients.Get(token)
   159  	if ok {
   160  		return value.(authorizationv1.AuthorizationV1Interface), nil
   161  	}
   162  
   163  	config := rest.CopyConfig(c.localConfig)
   164  	config.BearerToken = token
   165  	config.BearerTokenFile = ""
   166  
   167  	authCli, err := authorizationv1.NewForConfig(config)
   168  	if err != nil {
   169  		return nil, err
   170  	}
   171  
   172  	_ = c.authClients.Add(token, authCli)
   173  
   174  	return authCli, nil
   175  }
   176  
   177  // Num returns the num of clients
   178  func (c *ClientsPool) Num() int {
   179  	return c.clients.Len()
   180  }
   181  
   182  // Contains return true if have client for the token
   183  func (c *ClientsPool) Contains(token string) bool {
   184  	c.RLock()
   185  	defer c.RUnlock()
   186  
   187  	_, ok := c.clients.Get(token)
   188  	return ok
   189  }
   190  
   191  // ExtractTokenFromHeader extracts token from http header
   192  func ExtractTokenFromHeader(header http.Header) string {
   193  	auth := header.Get("Authorization")
   194  	if strings.HasPrefix(auth, "Bearer ") {
   195  		return strings.TrimPrefix(auth, "Bearer ")
   196  	}
   197  
   198  	return ""
   199  }
   200  
   201  // ExtractTokenAndGetClient extracts token from http header, and get the k8s client of this token
   202  func ExtractTokenAndGetClient(header http.Header) (pkgclient.Client, error) {
   203  	token := ExtractTokenFromHeader(header)
   204  	return K8sClients.Client(token)
   205  }
   206  
   207  // ExtractTokenAndGetAuthClient extracts token from http header, and get the authority client of this token
   208  func ExtractTokenAndGetAuthClient(header http.Header) (authorizationv1.AuthorizationV1Interface, error) {
   209  	token := ExtractTokenFromHeader(header)
   210  	return K8sClients.AuthClient(token)
   211  }
   212