main.go

package main

import (
	"flag"
	"fmt"
	"os"
	"slices"
	"strings"

	"google.golang.org/genproto/googleapis/api/visibility"
	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/pluginpb"
)

const (
	importPath  = "go.flipt.io/flipt/sdk/go"
	emptyImport = "google.golang.org/protobuf/types/known/emptypb"

	ignoreDecl = "flipt:sdk:ignore"
)

func main() {
	var flags flag.FlagSet
	grpcAPIConfig := flags.String("grpc_api_configuration", "", "path to GRPC API configuration")

	protogen.Options{ParamFunc: flags.Set}.Run(func(gen *protogen.Plugin) error {
		if *grpcAPIConfig == "" {
			fmt.Fprintln(os.Stderr, "path to grpc API configuration required")
			os.Exit(1)
		}

		// We have some use of the optional feature in our proto3 definitions.
		// This broadcasts that our plugin supports it and hides the generated
		// warning.
		gen.SupportedFeatures |= uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
		for _, f := range gen.Files {
			f.Services = slices.DeleteFunc(f.Services, shouldIgnoreService)
			f.Generate = f.Generate && len(f.Services) > 0

			if !f.Generate {
				continue
			}

			generateSubSDK(gen, f)
		}

		generateSDK(gen)

		generateGRPC(gen)

		generateHTTP(gen, *grpcAPIConfig)

		return nil
	})
}

func generateSDK(gen *protogen.Plugin) {
	g := gen.NewGeneratedFile("sdk.gen.go", importPath)
	g.P("// Code generated by protoc-gen-go-flipt-sdk. DO NOT EDIT.")
	g.P()
	g.P("package sdk")
	g.P()
	g.P("var _ *", importPackage(g, "time")("Time"))
	g.P("var _ *", importPackage(g, "os")("File"))
	g.P("var _ *", importPackage(g, "sync")("Mutex"))
	g.P("var _ ", importPackage(g, "go.flipt.io/flipt/rpc/flipt/auth")("Method"))
	g.P()
	g.P("const (")
	g.P(`defaultServiceAccountTokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token"`)
	g.P(`defaultKubernetesExpiryLeeway  = 10 * time.Second`)
	g.P(")")
	g.P()
	g.P("type Transport interface {")
	var types [][2]string
	for _, file := range gen.Files {
		if !file.Generate {
			continue
		}

		var (
			typ     = strings.Title(string(file.GoPackageName))
			method  = typ + "Client"
			returns = method
		)

		if len(file.Services) == 1 {
			returns = relativeImport(g, file, file.Services[0].GoName+"Client")
		}

		types = append(types, [...]string{typ, method})
		g.P(method, "() ", returns)
	}
	g.P("}")
	g.P()

	g.P(sdkBase)
	g.P()

	for _, t := range types {
		g.P("func (s SDK) ", t[0], "() *", t[0], "{")
		g.P("return &", t[0], "{")
		g.P("transport: s.transport.", t[1], "(),")
		g.P("authenticationProvider: s.authenticationProvider,")
		g.P("}")
		g.P("}\n")
	}

	authenticateFunction(g)
}

// generateSubSDK generates a .pb.sdk.go file containing a single SDK structure
// which represents an entire package from within the entire Flipt SDK API.
func generateSubSDK(gen *protogen.Plugin, file *protogen.File) (typ, client string) {
	filename := string(file.GoPackageName) + ".sdk.gen.go"
	g := gen.NewGeneratedFile(filename, importPath)
	g.P("// Code generated by protoc-gen-go-flipt-sdk. DO NOT EDIT.")
	g.P()
	g.P("package sdk")
	g.P()

	context := importPackage(g, "context")

	oneServicePackage := len(file.Services) == 1

	// define client structure
	typ = strings.Title(string(file.GoPackageName))
	relativeImport(g, file, typ+"Client")

	// We generate an interface which conjoins all the client interfaces
	// generated by the gRPC protoc generator.
	// Our gRPC and HTTP wrapper generators will take care of
	// bundling these clients appropriately for the SDK to consume.
	if !oneServicePackage {
		client = typ + "Client"

		g.P("type ", typ, "Client interface {")
		for _, srv := range file.Services {
			g.P(srv.GoName+"Client", "()", relativeImport(g, file, srv.GoName+"Client"))
		}
		g.P("}\n")

		g.P("type ", typ, " struct {")
		g.P("transport ", typ, "Client")
		g.P("authenticationProvider ", "ClientAuthenticationProvider")
		g.P("}\n")
	}

	for _, srv := range file.Services {
		serviceName := srv.GoName
		if oneServicePackage {
			serviceName = typ
		}

		g.P("type ", serviceName, " struct {")
		g.P("transport ", relativeImport(g, file, srv.GoName+"Client"))
		g.P("authenticationProvider ", "ClientAuthenticationProvider")
		g.P("}\n")

		if !oneServicePackage {
			g.P("func (s ", typ, ") ", srv.GoName, "() *", srv.GoName, "{")
			g.P("return &", srv.GoName, "{")
			g.P("transport: s.transport.", srv.GoName+"Client", "(),")
			g.P("authenticationProvider: ", "s.authenticationProvider,")
			g.P("}")
			g.P("}")
		}

		for _, method := range srv.Methods {
			var (
				signature       = []any{"func (x *", serviceName, ") ", method.GoName, "(ctx ", context("Context")}
				returnStatement = []any{"x.transport.", method.GoName, "(ctx, "}
			)

			if method.Input.GoIdent.GoImportPath != emptyImport {
				signature = append(signature, ", v *", method.Input.GoIdent)
				returnStatement = append(returnStatement, "v)")
			} else {
				returnStatement = append(returnStatement, "&", method.Input.GoIdent, "{})")
			}

			if method.Output.GoIdent.GoImportPath != emptyImport {
				g.P(append(signature, ") (*", method.Output.GoIdent, ", error) {")...)
				g.P("ctx, err := authenticate(ctx, x.authenticationProvider)")
				g.P("if err != nil { return nil, err }")
				g.P(append([]any{"return "}, returnStatement...)...)
			} else {
				g.P(append(signature, ") error {")...)
				g.P("ctx, err := authenticate(ctx, x.authenticationProvider)")
				g.P("if err != nil { return err }")
				g.P(append([]any{"_, err = "}, returnStatement...)...)
				g.P("return err")
			}

			g.P("}\n")
		}
	}
	return
}

func authenticateFunction(g *protogen.GeneratedFile) {
	context := importPackage(g, "context")
	g.P("func authenticate(ctx ", context("Context"), ", p ClientAuthenticationProvider) (", context("Context"), ", error) {")
	metadata := importPackage(g, "google.golang.org/grpc/metadata")
	g.P("if p != nil {")
	g.P("authentication, err := p.Authentication(ctx)")
	g.P("if err != nil { return ctx, err }")
	g.P()
	g.P("ctx = ", metadata("AppendToOutgoingContext"), `(ctx, "authorization", authentication)`)
	g.P("}\n")
	g.P("return ctx, nil")
	g.P("}")
	g.P()
}

func unexport(v string) string {
	return strings.ToLower(v[:1]) + v[1:]
}

func importPackage(g *protogen.GeneratedFile, pkg string) func(string) string {
	return func(name string) string {
		return g.QualifiedGoIdent(protogen.GoIdent{
			GoImportPath: protogen.GoImportPath(pkg),
			GoName:       name,
		})
	}
}

func relativeImport(g *protogen.GeneratedFile, file *protogen.File, name string) string {
	return g.QualifiedGoIdent(protogen.GoIdent{
		GoImportPath: file.GoImportPath,
		GoName:       name,
	})
}

const sdkBase = `// ClientTokenProvider is a type which when requested provides a
// client token which can be used to authenticate RPC/API calls
// invoked through the SDK.
// Deprecated: Use ClientAuthenticationProvider instead.
type ClientTokenProvider interface {
	ClientToken() (string, error)
}

// WithClientTokenProviders returns an Option which configures
// any supplied SDK with the provided ClientTokenProvider.
// Deprecated: Use WithAuthenticationProvider instead.
func WithClientTokenProvider(p ClientTokenProvider) Option {
	return func(s *SDK) {
		s.authenticationProvider = authenticationProviderFunc(func(context.Context) (string, error) {
			clientToken, err := p.ClientToken()
			if err != nil {
				return "", err
			}

			return "Bearer " + string(clientToken), nil
		})
	}
}

type authenticationProviderFunc func(context.Context) (string, error)

func (f authenticationProviderFunc) Authentication(ctx context.Context) (string, error) {
	return f(ctx)
}

// StaticClientTokenProvider is a string which is supplied as a static client token
// on each RPC which requires authentication.
// Deprecated: Use StaticTokenAuthenticationProvider instead.
type StaticClientTokenProvider string

// ClientToken returns the underlying string that is the StaticClientTokenProvider.
// Deprecated: Use StaticTokenAuthenticationProvider instead.
func (p StaticClientTokenProvider) ClientToken() (string, error) {
	return string(p), nil
}

// ClientAuthenticationProvider is a type which when requested provides a
// client authentication which can be used to authenticate RPC/API calls
// invoked through the SDK.
type ClientAuthenticationProvider interface {
	Authentication(context.Context) (string, error)
}

// SDK is the definition of Flipt's Go SDK.
// It depends on a pluggable transport implementation and exposes
// a consistent API surface area across both transport implementations.
// It also provides consistent client-side instrumentation and authentication
// lifecycle support.
type SDK struct {
	transport Transport
    authenticationProvider ClientAuthenticationProvider
}

// Option is a functional option which configures the Flipt SDK.
type Option func(*SDK)

// WithAuthenticationProviders returns an Option which configures
// any supplied SDK with the provided ClientAuthenticationProvider.
func WithAuthenticationProvider(p ClientAuthenticationProvider) Option {
	return func(s *SDK) {
        s.authenticationProvider = p
    }
}

// StaticTokenAuthenticationProvider is a string which is supplied as a static client authentication
// on each RPC which requires authentication.
type StaticTokenAuthenticationProvider string

// Authentication returns the underlying string that is the StaticTokenAuthenticationProvider.
func (p StaticTokenAuthenticationProvider) Authentication(context.Context) (string, error) {
    return "Bearer " + string(p), nil
}

// JWTAuthenticationProvider is a string which is supplied as a JWT client authentication
// on each RPC which requires authentication.
type JWTAuthenticationProvider string

// Authentication returns the underlying string that is the JWTAuthenticationProvider.
func (p JWTAuthenticationProvider) Authentication(context.Context) (string, error) {
    return "JWT " + string(p), nil
}

// KubernetesAuthenticationProvider is an implementation of ClientAuthenticationProvider
// which automatically uses the service account token from the environment and exchanges
// it with Flipt for a client token.
// This provider keeps the client token up to date and refreshes it for a new client
// token before expiry. It re-reads the service account token as Kubernetes can and will refresh
// this token, as it also has its own expiry.
type KubernetesAuthenticationProvider struct {
	transport               Transport
	serviceAccountTokenPath string
	leeway                  time.Duration

	mu   sync.RWMutex
	resp *auth.VerifyServiceAccountResponse
}

// KubernetesAuthenticationProviderOption is a functional option for configuring KubernetesAuthenticationProvider.
type KubernetesAuthenticationProviderOption func(*KubernetesAuthenticationProvider)

// WithKubernetesServiceAccountTokenPath sets the path on the host to locate the kubernetes service account.
// The KubernetesAuthenticationProvider uses the default location set by Kubernetes.
// This option lets you override that if your path happens to differ.
func WithKubernetesServiceAccountTokenPath(p string) KubernetesAuthenticationProviderOption {
	return func(kctp *KubernetesAuthenticationProvider) {
		kctp.serviceAccountTokenPath = p
	}
}

// WithKubernetesExpiryLeeway configures the duration leeway for deciding when to refresh
// the client token. The default is 10 seconds, which ensures that tokens are automatically refreshed
// when their is less that 10 seconds of lifetime left on the previously fetched client token.
func WithKubernetesExpiryLeeway(d time.Duration) KubernetesAuthenticationProviderOption {
	return func(kctp *KubernetesAuthenticationProvider) {
		kctp.leeway = d
	}
}

// NewKubernetesAuthenticationProvider constructs and configures a new KubernetesAuthenticationProvider
// using the provided transport.
func NewKubernetesAuthenticationProvider(transport Transport, opts ...KubernetesAuthenticationProviderOption) *KubernetesAuthenticationProvider {
    k := &KubernetesAuthenticationProvider{
		transport:               transport,
		serviceAccountTokenPath: defaultServiceAccountTokenPath,
		leeway:                  defaultKubernetesExpiryLeeway,
	}

    for _, opt := range opts {
        opt(k)
    }

    return k
}

// Authentication returns the authentication header string to be used for a request
// by the client SDK. It is generated via exchanging the local service account token
// with Flipt for a client token. The token is then formatted appropriately for use
// in the Authentication header as a bearer token.
func (k *KubernetesAuthenticationProvider) Authentication(ctx context.Context) (string, error) {
	k.mu.RLock()
	resp := k.resp
	k.mu.RUnlock()
	if resp != nil && time.Now().UTC().Add(k.leeway).Before(resp.Authentication.ExpiresAt.AsTime()) {
		return StaticTokenAuthenticationProvider(k.resp.ClientToken).Authentication(ctx)
	}

	k.mu.Lock()
	defer k.mu.Unlock()
	saToken, err := os.ReadFile(k.serviceAccountTokenPath)
	if err != nil {
		return "", err
	}

	resp, err = k.transport.
		AuthClient().
		AuthenticationMethodKubernetesServiceClient().
		VerifyServiceAccount(ctx, &auth.VerifyServiceAccountRequest{
			ServiceAccountToken: string(saToken),
		})
	if err != nil {
		return "", err
	}

	k.resp = resp

	return StaticTokenAuthenticationProvider(k.resp.ClientToken).Authentication(ctx)
}

// New constructs and configures a Flipt SDK instance from
// the provided Transport implementation and options.
func New(t Transport, opts ...Option) SDK {
    sdk := SDK{transport: t}

    for _, opt := range opts { opt(&sdk) }

    return sdk
}`

func shouldIgnoreService(srv *protogen.Service) bool {
	if v := proto.GetExtension(srv.Desc.Options(), visibility.E_ApiVisibility).(*visibility.VisibilityRule); v != nil {
		return v.Restriction == ignoreDecl
	}
	return false
}