github.com/pachyderm/pachyderm@v1.13.4/src/client/pkg/grpcutil/server.go (about)

     1  package grpcutil
     2  
     3  import (
     4  	"context"
     5  	gotls "crypto/tls"
     6  	"fmt"
     7  	"math"
     8  	"net"
     9  	"time"
    10  
    11  	"golang.org/x/sync/errgroup"
    12  	"google.golang.org/grpc"
    13  	"google.golang.org/grpc/credentials"
    14  	// Import registers the grpc GZIP decoder
    15  	_ "google.golang.org/grpc/encoding/gzip"
    16  	"google.golang.org/grpc/keepalive"
    17  
    18  	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
    19  	"github.com/pachyderm/pachyderm/src/client/pkg/tls"
    20  	"github.com/pachyderm/pachyderm/src/client/pkg/tracing"
    21  	log "github.com/sirupsen/logrus"
    22  )
    23  
    24  // Server is a convenience wrapper to gRPC servers that simplifies their
    25  // setup and execution
    26  type Server struct {
    27  	Server *grpc.Server
    28  	eg     *errgroup.Group
    29  }
    30  
    31  // NewServer creates a new gRPC server, but does not start serving yet.
    32  //
    33  // If 'publicPortTLSAllowed' is set, grpcutil may enable TLS. This should be
    34  // set for public ports that serve GRPC services to 3rd party clients. If set,
    35  // the criterion for actually serving over TLS is: if a signed TLS cert and
    36  // corresponding private key in 'TLSVolumePath', this will serve GRPC traffic
    37  // over TLS. If either are missing this will serve GRPC traffic over
    38  // unencrypted HTTP,
    39  func NewServer(ctx context.Context, publicPortTLSAllowed bool) (*Server, error) {
    40  	opts := []grpc.ServerOption{
    41  		grpc.MaxConcurrentStreams(math.MaxUint32),
    42  		grpc.MaxRecvMsgSize(MaxMsgSize),
    43  		grpc.MaxSendMsgSize(MaxMsgSize),
    44  		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
    45  			MinTime:             5 * time.Second,
    46  			PermitWithoutStream: true,
    47  		}),
    48  		grpc.UnaryInterceptor(tracing.UnaryServerInterceptor()),
    49  		grpc.StreamInterceptor(tracing.StreamServerInterceptor()),
    50  	}
    51  
    52  	var cLoader *tls.CertLoader
    53  	if publicPortTLSAllowed {
    54  		// Validate environment
    55  		certPath, keyPath, err := tls.GetCertPaths()
    56  		if err != nil {
    57  			log.Warnf("TLS disabled: %v", err)
    58  		} else {
    59  			cLoader = tls.NewCertLoader(certPath, keyPath, tls.CertCheckFrequency)
    60  			// Read TLS cert and key
    61  			err := cLoader.LoadAndStart()
    62  			if err != nil {
    63  				return nil, errors.Wrapf(err, "couldn't build transport creds: %v", err)
    64  			}
    65  			transportCreds := credentials.NewTLS(&gotls.Config{GetCertificate: cLoader.GetCertificate})
    66  			opts = append(opts, grpc.Creds(transportCreds))
    67  		}
    68  	}
    69  
    70  	server := grpc.NewServer(opts...)
    71  	eg, ctx := errgroup.WithContext(ctx)
    72  
    73  	eg.Go(func() error {
    74  		<-ctx.Done()
    75  		server.GracefulStop() // This also closes the listeners
    76  		if cLoader != nil {
    77  			cLoader.Stop()
    78  		}
    79  		return nil
    80  	})
    81  
    82  	return &Server{
    83  		Server: server,
    84  		eg:     eg,
    85  	}, nil
    86  }
    87  
    88  // ListenTCP causes the gRPC server to listen on a given TCP host and port
    89  func (s *Server) ListenTCP(host string, port uint16) (net.Listener, error) {
    90  	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", host, port))
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  
    95  	s.eg.Go(func() error {
    96  		return s.Server.Serve(listener)
    97  	})
    98  
    99  	return listener, nil
   100  }
   101  
   102  // Wait causes the gRPC server to wait until it finishes, returning any errors
   103  // that happened
   104  func (s *Server) Wait() error {
   105  	return s.eg.Wait()
   106  }