blob: 38a4961d39f51b05da719f091fcc961042395180 [file] [log] [blame]
// Copyright 2016 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main
import (
"bytes"
"compress/gzip"
"context"
"fmt"
"io/ioutil"
"os"
"os/exec"
"strings"
"text/template"
"github.com/golang/protobuf/proto"
"github.com/golang/protobuf/protoc-gen-go/descriptor"
)
const (
discoveryPackagePath = "go.chromium.org/luci/grpc/discovery"
)
// discoveryTmpl is template for generated Go discovery file.
// The result of execution will also be passed through gofmt.
var discoveryTmpl = template.Must(template.New("").Parse(strings.TrimSpace(`
// Code generated by cproto. DO NOT EDIT.
package {{.GoPkg}};
{{if .ImportDiscovery}}
import discovery "go.chromium.org/luci/grpc/discovery"
{{end}}
import "github.com/golang/protobuf/protoc-gen-go/descriptor"
func init() {
{{if .ImportDiscovery}}discovery.{{end}}RegisterDescriptorSetCompressed(
[]string{
{{range .ServiceNames}}"{{.}}",{{end}}
},
{{.CompressedBytes}},
)
}
// FileDescriptorSet returns a descriptor set for this proto package, which
// includes all defined services, and all transitive dependencies.
//
// Will not return nil.
//
// Do NOT modify the returned descriptor.
func FileDescriptorSet() *descriptor.FileDescriptorSet {
// We just need ONE of the service names to look up the FileDescriptorSet.
ret, err := {{if .ImportDiscovery}}discovery.{{end}}GetDescriptorSet("{{index .ServiceNames 0 }}")
if err != nil {
panic(err)
}
return ret
}
`)))
// genDiscoveryFile generates a Go discovery file that calls
// discovery.RegisterDescriptorSetCompressed(serviceNames, compressedDescBytes)
// in an init function.
func genDiscoveryFile(c context.Context, target, descFile, protoPkg, goPkg string) error {
descBytes, err := ioutil.ReadFile(descFile)
if err != nil {
return err
}
var desc descriptor.FileDescriptorSet
if err := proto.Unmarshal(descBytes, &desc); err != nil {
return fmt.Errorf("cannot parse generated descriptor file: %s", err)
}
var serviceNames []string
for _, f := range desc.File {
for _, s := range f.Service {
serviceNames = append(serviceNames, fmt.Sprintf("%s.%s", protoPkg, s.GetName()))
}
}
if len(serviceNames) == 0 {
// no services, no discovery.
return nil
}
compressedDescBytes, err := compress(descBytes)
if err != nil {
return err
}
inDiscoveryPackage, err := isInPackage(target, discoveryPackagePath)
if err != nil {
return err
}
var buf bytes.Buffer
err = discoveryTmpl.Execute(&buf, map[string]interface{}{
"GoPkg": goPkg,
"ImportDiscovery": !inDiscoveryPackage,
"ServiceNames": serviceNames,
"CompressedBytes": asByteArray(compressedDescBytes),
})
if err != nil {
return err
}
src := buf.Bytes()
formatted, err := gofmt(src)
if err != nil {
os.Stdout.Write(src)
return err
}
return ioutil.WriteFile(target, formatted, 0666)
}
// asByteArray converts blob to a valid []byte Go literal.
func asByteArray(blob []byte) string {
out := &bytes.Buffer{}
fmt.Fprintf(out, "[]byte{")
for i := 0; i < len(blob); i++ {
fmt.Fprintf(out, "%d, ", blob[i])
if i%14 == 1 {
fmt.Fprintln(out)
}
}
fmt.Fprintf(out, "}")
return out.String()
}
// gofmt applies "gofmt -s" to the content of the buffer.
func gofmt(blob []byte) ([]byte, error) {
out := bytes.Buffer{}
cmd := exec.Command("gofmt", "-s")
cmd.Stdin = bytes.NewReader(blob)
cmd.Stdout = &out
cmd.Stderr = os.Stderr
if err := cmd.Run(); err != nil {
return nil, err
}
return out.Bytes(), nil
}
// compress compresses data with gzip.
func compress(data []byte) ([]byte, error) {
var buf bytes.Buffer
w := gzip.NewWriter(&buf)
if _, err := w.Write(data); err != nil {
return nil, err
}
if err := w.Close(); err != nil {
return nil, err
}
return buf.Bytes(), nil
}