consul/internal/tools/protoc-gen-consul-rate-limit/postprocess/main.go

241 lines
6.0 KiB
Go
Raw Normal View History

// Copyright (c) HashiCorp, Inc.
[COMPLIANCE] License changes (#18443) * Adding explicit MPL license for sub-package This directory and its subdirectories (packages) contain files licensed with the MPLv2 `LICENSE` file in this directory and are intentionally licensed separately from the BSL `LICENSE` file at the root of this repository. * Adding explicit MPL license for sub-package This directory and its subdirectories (packages) contain files licensed with the MPLv2 `LICENSE` file in this directory and are intentionally licensed separately from the BSL `LICENSE` file at the root of this repository. * Updating the license from MPL to Business Source License Going forward, this project will be licensed under the Business Source License v1.1. Please see our blog post for more details at <Blog URL>, FAQ at www.hashicorp.com/licensing-faq, and details of the license at www.hashicorp.com/bsl. * add missing license headers * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 * Update copyright file headers to BUSL-1.1 --------- Co-authored-by: hashicorp-copywrite[bot] <110428419+hashicorp-copywrite[bot]@users.noreply.github.com>
2023-08-11 13:12:13 +00:00
// SPDX-License-Identifier: BUSL-1.1
2023-01-04 16:07:02 +00:00
package main
import (
"bytes"
"encoding/json"
"errors"
"flag"
"fmt"
"go/format"
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
"io/fs"
2023-01-04 16:07:02 +00:00
"os"
"path/filepath"
"sort"
"strings"
)
const (
usage = "Usage: %s -input=/proto-dir-1 -input=/proto-dir-2 -output=/mappings.go\n"
fileHeader = `// generated by protoc-gen-consul-rate-limit; DO NOT EDIT.
package middleware
import "github.com/hashicorp/consul/agent/consul/rate"
`
entTags = `//go:build consulent
// +build consulent
`
2023-01-04 16:07:02 +00:00
)
func main() {
var (
inputPaths sliceFlags
outputPath string
)
flag.Var(&inputPaths, "input", "")
flag.StringVar(&outputPath, "output", "", "")
flag.Parse()
if len(inputPaths) == 0 || outputPath == "" {
fmt.Fprintf(os.Stderr, usage, os.Args[0])
os.Exit(1)
}
if err := run(inputPaths, outputPath); err != nil {
fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
os.Exit(1)
}
}
func run(inputPaths []string, outputPath string) error {
if !strings.HasSuffix(outputPath, ".go") {
return errors.New("-output path must end in .go")
}
oss, ent, err := collectSpecs(inputPaths)
if err != nil {
return err
}
ossSource, err := generateOSS(oss)
if err != nil {
return err
}
if err := os.WriteFile(outputPath, ossSource, 0666); err != nil {
return fmt.Errorf("failed to write output file: %s - %w", outputPath, err)
}
// ent should only be non-zero in the enterprise repository.
if len(ent) > 0 {
entSource, err := generateENT(ent)
if err != nil {
return err
}
if err := os.WriteFile(enterpriseFileName(outputPath), entSource, 0666); err != nil {
return fmt.Errorf("failed to write output file: %s - %w", outputPath, err)
}
}
return nil
}
// enterpriseFileName adds the _ent filename suffix before the extension.
//
// Example:
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
//
2023-01-04 16:07:02 +00:00
// enterpriseFileName("bar/baz/foo.gen.go") => "bar/baz/foo_ent.gen.go"
func enterpriseFileName(filename string) string {
fileName := filepath.Base(filename)
extStart := strings.Index(fileName, ".")
return filepath.Join(
filepath.Dir(filename),
fileName[0:extStart]+"_ent"+fileName[extStart:],
)
}
type spec struct {
MethodName string
OperationType string
Enterprise bool
OperationCategory string
2023-01-04 16:07:02 +00:00
}
func (s spec) GoOperationType() string {
switch s.OperationType {
case "OPERATION_TYPE_WRITE":
return "rate.OperationTypeWrite"
case "OPERATION_TYPE_READ":
return "rate.OperationTypeRead"
case "OPERATION_TYPE_EXEMPT":
return "rate.OperationTypeExempt"
}
panic(fmt.Sprintf("unknown rate limit operation type: %s", s.OperationType))
}
func (s spec) GoOperationCategory() string {
switch s.OperationCategory {
case "OPERATION_CATEGORY_ACL":
return "rate.OperationCategoryACL"
case "OPERATION_CATEGORY_PEER_STREAM":
return "rate.OperationCategoryPeerStream"
case "OPERATION_CATEGORY_CONNECT_CA":
return "rate.OperationCategoryConnectCA"
case "OPERATION_CATEGORY_PARTITION":
return "rate.OperationCategoryPartition"
case "OPERATION_CATEGORY_PEERING":
return "rate.OperationCategoryPeering"
case "OPERATION_CATEGORY_SERVER_DISCOVERY":
return "rate.OperationCategoryServerDiscovery"
case "OPERATION_CATEGORY_DATAPLANE":
return "rate.OperationCategoryDataPlane"
case "OPERATION_CATEGORY_DNS":
return "rate.OperationCategoryDNS"
case "OPERATION_CATEGORY_SUBSCRIBE":
return "rate.OperationCategorySubscribe"
case "OPERATION_CATEGORY_OPERATOR":
return "rate.OperationCategoryOperator"
case "OPERATION_CATEGORY_RESOURCE":
return "rate.OperationCategoryResource"
}
panic(fmt.Sprintf("unknown rate limit operation category: %s found in method: %s", s.OperationCategory, s.MethodName))
}
2023-01-04 16:07:02 +00:00
func collectSpecs(inputPaths []string) ([]spec, []spec, error) {
var specs []spec
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
var specFiles []string
2023-01-04 16:07:02 +00:00
for _, protoPath := range inputPaths {
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
err := filepath.WalkDir(protoPath, func(path string, info fs.DirEntry, err error) error {
if err != nil {
return err
}
if info.Name() == ".ratelimit.tmp" {
specFiles = append(specFiles, path)
}
return nil
})
2023-01-04 16:07:02 +00:00
if err != nil {
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
return nil, nil, fmt.Errorf("failed to walk directory: %s - %w", protoPath, err)
2023-01-04 16:07:02 +00:00
}
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
}
2023-01-04 16:07:02 +00:00
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
for _, file := range specFiles {
b, err := os.ReadFile(file)
if err != nil {
return nil, nil, fmt.Errorf("failed to read ratelimit file: %w", err)
}
2023-01-04 16:07:02 +00:00
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
var fileSpecs []spec
if err := json.Unmarshal(b, &fileSpecs); err != nil {
return nil, nil, fmt.Errorf("failed to unmarshal ratelimit file %s - %w", file, err)
2023-01-04 16:07:02 +00:00
}
Protobuf Refactoring for Multi-Module Cleanliness (#16302) Protobuf Refactoring for Multi-Module Cleanliness This commit includes the following: Moves all packages that were within proto/ to proto/private Rewrites imports to account for the packages being moved Adds in buf.work.yaml to enable buf workspaces Names the proto-public buf module so that we can override the Go package imports within proto/buf.yaml Bumps the buf version dependency to 1.14.0 (I was trying out the version to see if it would get around an issue - it didn't but it also doesn't break things and it seemed best to keep up with the toolchain changes) Why: In the future we will need to consume other protobuf dependencies such as the Google HTTP annotations for openapi generation or grpc-gateway usage. There were some recent changes to have our own ratelimiting annotations. The two combined were not working when I was trying to use them together (attempting to rebase another branch) Buf workspaces should be the solution to the problem Buf workspaces means that each module will have generated Go code that embeds proto file names relative to the proto dir and not the top level repo root. This resulted in proto file name conflicts in the Go global protobuf type registry. The solution to that was to add in a private/ directory into the path within the proto/ directory. That then required rewriting all the imports. Is this safe? AFAICT yes The gRPC wire protocol doesn't seem to care about the proto file names (although the Go grpc code does tack on the proto file name as Metadata in the ServiceDesc) Other than imports, there were no changes to any generated code as a result of this.
2023-02-17 21:14:46 +00:00
specs = append(specs, fileSpecs...)
2023-01-04 16:07:02 +00:00
}
sort.Slice(specs, func(a, b int) bool {
return specs[a].MethodName < specs[b].MethodName
})
var oss, ent []spec
for _, spec := range specs {
if spec.Enterprise {
ent = append(ent, spec)
} else {
oss = append(oss, spec)
}
}
return oss, ent, nil
}
func generateOSS(specs []spec) ([]byte, error) {
var output bytes.Buffer
output.WriteString(fileHeader)
fmt.Fprintln(&output, `var rpcRateLimitSpecs = map[string]rate.OperationSpec{`)
2023-01-04 16:07:02 +00:00
for _, spec := range specs {
fmt.Fprintf(&output, `"%s": {Type: %s, Category: %s},`, spec.MethodName, spec.GoOperationType(), spec.GoOperationCategory())
2023-01-04 16:07:02 +00:00
output.WriteString("\n")
}
output.WriteString("}")
formatted, err := format.Source(output.Bytes())
if err != nil {
return nil, fmt.Errorf("failed to format source in oss: %w", err)
2023-01-04 16:07:02 +00:00
}
return formatted, nil
}
func generateENT(specs []spec) ([]byte, error) {
var output bytes.Buffer
output.WriteString(entTags)
output.WriteString(fileHeader)
output.WriteString("func init() {\n")
for _, spec := range specs {
fmt.Fprintf(&output, `rpcRateLimitSpecs["%s"] = rate.OperationSpec{Type: %s, Category: %s}`, spec.MethodName, spec.GoOperationType(), spec.GoOperationCategory())
2023-01-04 16:07:02 +00:00
output.WriteString("\n")
}
output.WriteString("}")
formatted, err := format.Source(output.Bytes())
if err != nil {
return nil, fmt.Errorf("failed to format source in ent: %w", err)
2023-01-04 16:07:02 +00:00
}
return formatted, nil
}
type sliceFlags []string
func (i *sliceFlags) Set(value string) error {
*i = append(*i, value)
return nil
}
func (i *sliceFlags) String() string { return "" }