serve.go

package entrypoints

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	_ "net/http/pprof" // Required to serve application.
	"strings"

	"github.com/pkg/errors"
	"google.golang.org/grpc/credentials"

	"github.com/lyft/flyteadmin/pkg/common"

	"github.com/lyft/flytestdlib/logger"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	flyteService "github.com/lyft/flyteidl/gen/pb-go/flyteidl/service"

	"github.com/lyft/flyteadmin/pkg/config"
	"github.com/lyft/flyteadmin/pkg/rpc/adminservice"

	"github.com/spf13/cobra"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/lyft/flytestdlib/contextutils"
	"github.com/lyft/flytestdlib/promutils/labeled"
	"google.golang.org/grpc"
)

type ServingOptions struct {
	Secure bool
	// Optional Arguments, Should be provided to enable secure mode
	CertFile string
	KeyFile  string
}

var serviceOpts = ServingOptions{}

// serveCmd represents the serve command
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Launches the Flyte admin server",
	RunE: func(cmd *cobra.Command, args []string) error {
		ctx := context.Background()
		cfg := config.GetConfig()
		if serviceOpts.Secure {
			return serveGatewaySecure(ctx, cfg, serviceOpts)
		}
		return serveGatewayInsecure(ctx, cfg)
	},
}

func init() {
	// Command information
	RootCmd.AddCommand(serveCmd)
	serveCmd.Flags().BoolVarP(&serviceOpts.Secure, "secure", "s", false, "Use ssl")
	serveCmd.Flags().StringVarP(&serviceOpts.CertFile, "cert-file", "c", "", "Path of file that contains x509 certificate")
	serveCmd.Flags().StringVarP(&serviceOpts.KeyFile, "key-file", "k", "", "Path of file that contains x509 client key")

	// Set Keys
	labeled.SetMetricKeys(contextutils.AppNameKey, contextutils.ProjectKey, contextutils.DomainKey,
		contextutils.ExecIDKey, contextutils.WorkflowIDKey, contextutils.NodeIDKey, contextutils.TaskIDKey,
		contextutils.TaskTypeKey, common.RuntimeTypeKey, common.RuntimeVersionKey)
}

// Creates a new gRPC Server with all the configuration
func newGRPCServer(_ context.Context, cfg *config.Config, opts ...grpc.ServerOption) (*grpc.Server, error) {
	serverOpts := []grpc.ServerOption{
		grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
		grpc.UnaryInterceptor(grpc_prometheus.UnaryServerInterceptor),
	}
	serverOpts = append(serverOpts, opts...)
	grpcServer := grpc.NewServer(serverOpts...)
	grpc_prometheus.Register(grpcServer)
	flyteService.RegisterAdminServiceServer(grpcServer, adminservice.NewAdminServer(cfg.KubeConfig, cfg.Master))
	return grpcServer, nil
}

func newHTTPServer(ctx context.Context, cfg *config.Config, grpcConnectionOpts []grpc.DialOption, grpcAddress string) (*http.ServeMux, error) {
	// Register the server that will serve HTTP/REST Traffic
	mux := http.NewServeMux()

	// Register healthcheck
	mux.HandleFunc("/healthcheck", func(w http.ResponseWriter, r *http.Request) {
		// A very simple health check.
		w.WriteHeader(http.StatusOK)
	})

	// Register OpenAPI endpoint
	// This endpoint will serve the OpenAPI2 spec generated by the swagger protoc plugin, and bundled by go-bindata
	mux.HandleFunc("/api/v1/openapi", func(w http.ResponseWriter, r *http.Request) {
		swaggerBytes, err := flyteService.Asset("admin.swagger.json")
		if err != nil {
			logger.Warningf(ctx, "Err %v", err)
			w.WriteHeader(http.StatusFailedDependency)
		} else {
			w.WriteHeader(http.StatusOK)
			_, err := w.Write(swaggerBytes)
			if err != nil {
				logger.Errorf(ctx, "failed to write openAPI information, error: %s", err.Error())
			}
		}
	})

	// Register the actual Server that will service gRPC traffic
	gwmux := runtime.NewServeMux(runtime.WithMarshalerOption("application/octet-stream", &runtime.ProtoMarshaller{}))
	err := flyteService.RegisterAdminServiceHandlerFromEndpoint(ctx, gwmux, grpcAddress, grpcConnectionOpts)
	if err != nil {
		return nil, errors.Wrap(err, "error registering admin service")
	}

	mux.Handle("/", gwmux)

	return mux, nil
}

func serveGatewayInsecure(ctx context.Context, cfg *config.Config) error {
	logger.Infof(ctx, "Serving FlyteAdmin Insecure")
	grpcServer, err := newGRPCServer(ctx, cfg)
	if err != nil {
		return errors.Wrap(err, "failed to create GRPC server")
	}

	logger.Infof(ctx, "Serving GRPC Traffic on: %s", cfg.GetGrpcHostAddress())
	lis, err := net.Listen("tcp", cfg.GetGrpcHostAddress())
	if err != nil {
		return errors.Wrapf(err, "failed to listen on GRPC port: %s", cfg.GetGrpcHostAddress())
	}

	go func() {
		err := grpcServer.Serve(lis)
		logger.Fatalf(ctx, "Failed to create GRPC Server, Err: ", err)
	}()

	logger.Infof(ctx, "Starting HTTP/1 Gateway server on %s", cfg.GetHostAddress())
	httpServer, err := newHTTPServer(ctx, cfg, []grpc.DialOption{grpc.WithInsecure()}, cfg.GetGrpcHostAddress())
	if err != nil {
		return err
	}
	err = http.ListenAndServe(cfg.GetHostAddress(), httpServer)
	if err != nil {
		return errors.Wrapf(err, "failed to Start HTTP Server")
	}

	return nil
}

// grpcHandlerFunc returns an http.Handler that delegates to grpcServer on incoming gRPC
// connections or otherHandler otherwise.
// See https://github.com/philips/grpc-gateway-example/blob/master/cmd/serve.go for reference
func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// This is a partial recreation of gRPC's internal checks
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			logger.Infof(context.TODO(), "Received GRPC request, %s, %s, %s, %d, %d", r.RequestURI, r.Header, r.Method, r.ProtoMajor, r.ProtoMinor)
			grpcServer.ServeHTTP(w, r)
		} else {
			logger.Infof(context.TODO(), "Received regular request - %s, %s, %s, %d, %d", r.RequestURI, r.Header, r.Method, r.ProtoMajor, r.ProtoMinor)
			otherHandler.ServeHTTP(w, r)
		}
	})
}

func serveGatewaySecure(ctx context.Context, cfg *config.Config, opts ServingOptions) error {
	// This support single cert right now?
	var err error
	cert, err := tls.LoadX509KeyPair(opts.CertFile, opts.KeyFile)
	if err != nil {
		return err
	}
	certPool := x509.NewCertPool()
	data, err := ioutil.ReadFile(opts.CertFile)
	if err != nil {
		return errors.Wrapf(err, "failed to read server cert file: %s", opts.CertFile)
	}
	if ok := certPool.AppendCertsFromPEM([]byte(data)); !ok {
		return fmt.Errorf("failed to load certificate into the pool")
	}

	grpcServer, err := newGRPCServer(ctx, cfg, grpc.Creds(credentials.NewClientTLSFromCert(certPool, cfg.GetHostAddress())))
	if err != nil {
		return errors.Wrap(err, "failed to create GRPC server")
	}

	dialCreds := credentials.NewTLS(&tls.Config{
		ServerName: cfg.GetHostAddress(),
		RootCAs:    certPool,
	})
	httpServer, err := newHTTPServer(ctx, cfg, []grpc.DialOption{grpc.WithTransportCredentials(dialCreds)}, cfg.GetHostAddress())
	if err != nil {
		return err
	}

	conn, err := net.Listen("tcp", cfg.GetHostAddress())
	if err != nil {
		panic(err)
	}

	srv := &http.Server{
		Addr:    cfg.GetHostAddress(),
		Handler: grpcHandlerFunc(grpcServer, httpServer),
		TLSConfig: &tls.Config{
			Certificates: []tls.Certificate{cert},
			NextProtos:   []string{"h2"},
		},
	}

	err = srv.Serve(tls.NewListener(conn, srv.TLSConfig))

	if err != nil {
		return errors.Wrapf(err, "failed to Start HTTP/2 Server")
	}
	return nil
}