vendor/github.com/hashicorp/go-plugin/grpc_server.go

   1 package plugin
   2
   3 import (
   4         "bytes"
   5         "crypto/tls"
   6         "encoding/json"
   7         "fmt"
   8         "io"
   9         "net"
  10
  11         "google.golang.org/grpc"
  12         "google.golang.org/grpc/credentials"
  13         "google.golang.org/grpc/health"
  14         "google.golang.org/grpc/health/grpc_health_v1"
  15 )
  16
  17 // GRPCServiceName is the name of the service that the health check should
  18 // return as passing.
  19 const GRPCServiceName = "plugin"
  20
  21 // DefaultGRPCServer can be used with the "GRPCServer" field for Server
  22 // as a default factory method to create a gRPC server with no extra options.
  23 func DefaultGRPCServer(opts []grpc.ServerOption) *grpc.Server {
  24         return grpc.NewServer(opts...)
  25 }
  26
  27 // GRPCServer is a ServerType implementation that serves plugins over
  28 // gRPC. This allows plugins to easily be written for other languages.
  29 //
  30 // The GRPCServer outputs a custom configuration as a base64-encoded
  31 // JSON structure represented by the GRPCServerConfig config structure.
  32 type GRPCServer struct {
  33         // Plugins are the list of plugins to serve.
  34         Plugins map[string]Plugin
  35
  36         // Server is the actual server that will accept connections. This
  37         // will be used for plugin registration as well.
  38         Server func([]grpc.ServerOption) *grpc.Server
  39
  40         // TLS should be the TLS configuration if available. If this is nil,
  41         // the connection will not have transport security.
  42         TLS *tls.Config
  43
  44         // DoneCh is the channel that is closed when this server has exited.
  45         DoneCh chan struct{}
  46
  47         // Stdout/StderrLis are the readers for stdout/stderr that will be copied
  48         // to the stdout/stderr connection that is output.
  49         Stdout io.Reader
  50         Stderr io.Reader
  51
  52         config GRPCServerConfig
  53         server *grpc.Server
  54         broker *GRPCBroker
  55 }
  56
  57 // ServerProtocol impl.
  58 func (s *GRPCServer) Init() error {
  59         // Create our server
  60         var opts []grpc.ServerOption
  61         if s.TLS != nil {
  62                 opts = append(opts, grpc.Creds(credentials.NewTLS(s.TLS)))
  63         }
  64         s.server = s.Server(opts)
  65
  66         // Register the health service
  67         healthCheck := health.NewServer()
  68         healthCheck.SetServingStatus(
  69                 GRPCServiceName, grpc_health_v1.HealthCheckResponse_SERVING)
  70         grpc_health_v1.RegisterHealthServer(s.server, healthCheck)
  71
  72         // Register the broker service
  73         brokerServer := newGRPCBrokerServer()
  74         RegisterGRPCBrokerServer(s.server, brokerServer)
  75         s.broker = newGRPCBroker(brokerServer, s.TLS)
  76         go s.broker.Run()
  77
  78         // Register all our plugins onto the gRPC server.
  79         for k, raw := range s.Plugins {
  80                 p, ok := raw.(GRPCPlugin)
  81                 if !ok {
  82                         return fmt.Errorf("%q is not a GRPC-compatible plugin", k)
  83                 }
  84
  85                 if err := p.GRPCServer(s.broker, s.server); err != nil {
  86                         return fmt.Errorf("error registring %q: %s", k, err)
  87                 }
  88         }
  89
  90         return nil
  91 }
  92
  93 // Stop calls Stop on the underlying grpc.Server
  94 func (s *GRPCServer) Stop() {
  95         s.server.Stop()
  96 }
  97
  98 // GracefulStop calls GracefulStop on the underlying grpc.Server
  99 func (s *GRPCServer) GracefulStop() {
 100         s.server.GracefulStop()
 101 }
 102
 103 // Config is the GRPCServerConfig encoded as JSON then base64.
 104 func (s *GRPCServer) Config() string {
 105         // Create a buffer that will contain our final contents
 106         var buf bytes.Buffer
 107
 108         // Wrap the base64 encoding with JSON encoding.
 109         if err := json.NewEncoder(&buf).Encode(s.config); err != nil {
 110                 // We panic since ths shouldn't happen under any scenario. We
 111                 // carefully control the structure being encoded here and it should
 112                 // always be successful.
 113                 panic(err)
 114         }
 115
 116         return buf.String()
 117 }
 118
 119 func (s *GRPCServer) Serve(lis net.Listener) {
 120         // Start serving in a goroutine
 121         go s.server.Serve(lis)
 122
 123         // Wait until graceful completion
 124         <-s.DoneCh
 125 }
 126
 127 // GRPCServerConfig is the extra configuration passed along for consumers
 128 // to facilitate using GRPC plugins.
 129 type GRPCServerConfig struct {
 130         StdoutAddr string `json:"stdout_addr"`
 131         StderrAddr string `json:"stderr_addr"`
 132 }