grpcutil/grpcutil.go

   1 // Encapsulates common RPC server setup.
   2 package grpcutil
   3
   4 import (
   5         "crypto/tls"
   6         "crypto/x509"
   7         "flag"
   8         "fmt"
   9         "io/ioutil"
  10         "net"
  11         "net/http"
  12         "strings"
  13
  14         "github.com/Debian/dcs/internal/addrfd"
  15         "github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing"
  16         "golang.org/x/net/http2"
  17         "golang.org/x/net/trace"
  18         "google.golang.org/grpc"
  19         "google.golang.org/grpc/credentials"
  20         "google.golang.org/grpc/reflection"
  21 )
  22
  23 var (
  24         requireClientAuth = flag.Bool("tls_require_client_auth",
  25                 true,
  26                 "Require TLS Client Authentication")
  27 )
  28
  29 func init() {
  30         // Disable grpc tracing until
  31         // https://github.com/grpc/grpc-go/issues/695 is fixed.
  32         grpc.EnableTracing = false
  33 }
  34
  35 // grpcHandlerFunc returns an http.Handler that delegates to grpcServer on incoming gRPC
  36 // connections or otherHandler otherwise. Copied from cockroachdb.
  37 func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
  38         return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  39                 // This is a partial recreation of gRPC's internal checks:
  40                 // https://github.com/grpc/grpc-go/blob/7834b974e55fbf85a5b01afb5821391c71084efd/transport/handler_server.go#L61
  41                 if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
  42                         grpcServer.ServeHTTP(w, r)
  43                 } else {
  44                         otherHandler.ServeHTTP(w, r)
  45                 }
  46         })
  47 }
  48
  49 func DialTLS(addr, certFile, keyFile string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
  50         cert, err := tls.LoadX509KeyPair(certFile, keyFile)
  51         if err != nil {
  52                 return nil, err
  53         }
  54         roots := x509.NewCertPool()
  55         contents, err := ioutil.ReadFile(certFile)
  56         if err != nil {
  57                 return nil, err
  58         }
  59         if !roots.AppendCertsFromPEM(contents) {
  60                 return nil, fmt.Errorf("Could not parse %q as PEM file (contents: %q)", certFile, contents)
  61         }
  62         auth := credentials.NewTLS(&tls.Config{
  63                 RootCAs:      roots,
  64                 Certificates: []tls.Certificate{cert}})
  65
  66         return grpc.Dial(addr,
  67                 append([]grpc.DialOption{
  68                         grpc.WithTransportCredentials(auth),
  69                         grpc.WithStreamInterceptor(grpc_opentracing.StreamClientInterceptor()),
  70                         grpc.WithUnaryInterceptor(grpc_opentracing.UnaryClientInterceptor()),
  71                 }, opts...)...)
  72 }
  73
  74 func ListenAndServeTLS(addr, certFile, keyFile string, register func(s *grpc.Server)) error {
  75         ln, err := net.Listen("tcp", addr)
  76         if err != nil {
  77                 return err
  78         }
  79
  80         auth, err := credentials.NewServerTLSFromFile(certFile, keyFile)
  81         if err != nil {
  82                 return err
  83         }
  84
  85         s := grpc.NewServer(
  86                 grpc.Creds(auth),
  87                 grpc.StreamInterceptor(grpc_opentracing.StreamServerInterceptor()),
  88                 grpc.UnaryInterceptor(grpc_opentracing.UnaryServerInterceptor()))
  89
  90         register(s)
  91         reflection.Register(s)
  92
  93         srv := http.Server{
  94                 Addr:    addr,
  95                 Handler: grpcHandlerFunc(s, http.DefaultServeMux),
  96         }
  97         if err := http2.ConfigureServer(&srv, nil); err != nil {
  98                 return err
  99         }
 100         roots := x509.NewCertPool()
 101         contents, err := ioutil.ReadFile(certFile)
 102         if err != nil {
 103                 return err
 104         }
 105         if !roots.AppendCertsFromPEM(contents) {
 106                 return fmt.Errorf("Could not parse %q as PEM file (contents: %q)", certFile, contents)
 107         }
 108
 109         if *requireClientAuth {
 110                 srv.TLSConfig.ClientCAs = roots
 111                 srv.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
 112                 trace.AuthRequest = func(req *http.Request) (bool, bool) {
 113                         return true, true
 114                 }
 115         }
 116         srv.TLSConfig.Certificates = make([]tls.Certificate, 1)
 117         srv.TLSConfig.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
 118         addrfd.MustWrite(ln.Addr().String())
 119         return srv.Serve(tls.NewListener(ln, srv.TLSConfig))
 120 }