github.com/kubeshop/testkube@v1.17.23/pkg/logs/client/client.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"crypto/x509"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  	"time"
    11  
    12  	"go.uber.org/zap"
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/credentials"
    15  	"google.golang.org/grpc/credentials/insecure"
    16  
    17  	"github.com/kubeshop/testkube/pkg/log"
    18  	"github.com/kubeshop/testkube/pkg/logs/events"
    19  	"github.com/kubeshop/testkube/pkg/logs/pb"
    20  )
    21  
    22  const (
    23  	buffer          = 100
    24  	requestDeadline = time.Minute * 5
    25  )
    26  
    27  // NewGrpcClient imlpements getter interface for log stream for given ID
    28  func NewGrpcClient(address string, creds credentials.TransportCredentials) StreamGetter {
    29  	return &GrpcClient{
    30  		log:     log.DefaultLogger.With("service", "logs-grpc-client"),
    31  		creds:   creds,
    32  		address: address,
    33  	}
    34  }
    35  
    36  type GrpcClient struct {
    37  	log     *zap.SugaredLogger
    38  	creds   credentials.TransportCredentials
    39  	address string
    40  }
    41  
    42  // Get returns channel with log stream chunks for given execution id connects through GRPC to log service
    43  func (c GrpcClient) Get(ctx context.Context, id string) (chan events.LogResponse, error) {
    44  	ch := make(chan events.LogResponse, buffer)
    45  
    46  	log := c.log.With("id", id)
    47  
    48  	log.Debugw("getting logs", "address", c.address)
    49  
    50  	go func() {
    51  		// Contact the server and print out its response.
    52  		ctx, cancel := context.WithTimeout(context.Background(), requestDeadline)
    53  		defer cancel()
    54  		defer close(ch)
    55  
    56  		// TODO add TLS to GRPC client
    57  		creds := insecure.NewCredentials()
    58  		if c.creds != nil {
    59  			creds = c.creds
    60  		}
    61  
    62  		conn, err := grpc.Dial(c.address, grpc.WithTransportCredentials(creds))
    63  		if err != nil {
    64  			ch <- events.LogResponse{Error: err}
    65  			return
    66  		}
    67  		defer conn.Close()
    68  		log.Debugw("connected to grpc server")
    69  
    70  		client := pb.NewLogsServiceClient(conn)
    71  
    72  		r, err := client.Logs(ctx, &pb.LogRequest{ExecutionId: id})
    73  		if err != nil {
    74  			ch <- events.LogResponse{Error: err}
    75  			log.Errorw("error getting logs", "error", err)
    76  			return
    77  		}
    78  
    79  		log.Debugw("client start streaming")
    80  		defer func() {
    81  			log.Debugw("client stopped streaming")
    82  		}()
    83  
    84  		for {
    85  			l, err := r.Recv()
    86  			if err == io.EOF {
    87  				log.Infow("client stream finished", "error", err)
    88  				return
    89  			} else if err != nil {
    90  				ch <- events.LogResponse{Error: err}
    91  				log.Errorw("error receiving log response", "error", err)
    92  				return
    93  			}
    94  
    95  			logChunk := pb.MapFromPB(l)
    96  
    97  			// catch finish event
    98  			if events.IsFinished(&logChunk) {
    99  				log.Infow("received finish", "log", l)
   100  				return
   101  			}
   102  
   103  			log.Debugw("grpc client log", "log", l)
   104  			// send to the channel
   105  			ch <- events.LogResponse{Log: logChunk}
   106  		}
   107  	}()
   108  
   109  	return ch, nil
   110  }
   111  
   112  // GrpcConnectionConfig contains GRPC connection parameters
   113  type GrpcConnectionConfig struct {
   114  	Secure     bool
   115  	SkipVerify bool
   116  	CertFile   string
   117  	KeyFile    string
   118  	CAFile     string
   119  }
   120  
   121  // GetGrpcTransportCredentials returns transport credentials for GRPC connection config
   122  func GetGrpcTransportCredentials(cfg GrpcConnectionConfig) (credentials.TransportCredentials, error) {
   123  	var creds credentials.TransportCredentials
   124  
   125  	if cfg.Secure {
   126  		var tlsConfig tls.Config
   127  
   128  		if cfg.SkipVerify {
   129  			tlsConfig.InsecureSkipVerify = true
   130  		} else {
   131  			if cfg.CertFile != "" && cfg.KeyFile != "" {
   132  				cert, err := tls.LoadX509KeyPair(cfg.CertFile, cfg.KeyFile)
   133  				if err != nil {
   134  					return nil, err
   135  				}
   136  
   137  				tlsConfig.Certificates = []tls.Certificate{cert}
   138  			}
   139  
   140  			if cfg.CAFile != "" {
   141  				caCertificate, err := os.ReadFile(cfg.CAFile)
   142  				if err != nil {
   143  					return nil, err
   144  				}
   145  
   146  				certPool := x509.NewCertPool()
   147  				if !certPool.AppendCertsFromPEM(caCertificate) {
   148  					return nil, fmt.Errorf("failed to add server CA's certificate")
   149  				}
   150  
   151  				tlsConfig.RootCAs = certPool
   152  			}
   153  		}
   154  
   155  		creds = credentials.NewTLS(&tlsConfig)
   156  	}
   157  
   158  	return creds, nil
   159  }