generator.go

package main

import (
	"bytes"
	"fmt"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
	"strings"
)

type generator struct {
	Request  *plugin.CodeGeneratorRequest
	Response *plugin.CodeGeneratorResponse

	output *bytes.Buffer
	indent string
}

func newGenerator(req *plugin.CodeGeneratorRequest) *generator {
	return &generator{
		Request:  req,
		Response: nil,
		output:   bytes.NewBuffer(nil),
		indent:   "",
	}
}

func (g *generator) Generate() error {
	g.Response = &plugin.CodeGeneratorResponse{}

	for _, file := range g.getProtoFiles() {
		err := g.processFile(file)
		if err != nil {
			return err
		}
	}

	return nil
}

func (g *generator) processFile(file *descriptor.FileDescriptorProto) error {
	if file.Options.GetJavaGenericServices() {
		return fmt.Errorf("twirp_java_jaxrs cannot not work with java_generic_services option")
	}

	g.Response.File = append(g.Response.File, g.generateProvider(file))

	for _, service := range file.GetService() {
		out := g.generateServiceInterface(file, service)
		g.Response.File = append(g.Response.File, out)

		out = g.generateServiceClient(file, service)
		g.Response.File = append(g.Response.File, out)
	}

	return nil
}

func (g *generator) generateProvider(file *descriptor.FileDescriptorProto) *plugin.CodeGeneratorResponse_File {

	multi := file.Options.GetJavaMultipleFiles()
	serviceName := "ProtoBufMessageProvider"

	if multi {
		pkg := getJavaPackage(file)
		g.P(`// Code generated by protoc-gen-twirp_java_jaxrs, DO NOT EDIT.`)
		g.P(`// source: `, file.GetName())
		g.P()
		if pkg != "" {
			g.P(`package `, pkg, `;`)
			g.P()
		}
	}

	static := ""
	if !multi {
		static = "static "
	}
	g.P(`@javax.ws.rs.ext.Provider
@javax.ws.rs.Produces({"application/protobuf", "application/json"})
@javax.ws.rs.Consumes({"application/protobuf", "application/json"})
public `, static, `class `, serviceName, ` implements javax.ws.rs.ext.MessageBodyWriter<com.google.protobuf.Message>, javax.ws.rs.ext.MessageBodyReader<com.google.protobuf.Message> {

        @Override
        public boolean isWriteable(Class<?> type, java.lang.reflect.Type genericType, java.lang.annotation.Annotation[] annotations, javax.ws.rs.core.MediaType mediaType) {
            return com.google.protobuf.Message.class.isAssignableFrom(type) && ("json".equals(mediaType.getSubtype()) || "protobuf".equals(mediaType.getSubtype()));
        }

        @Override
        public long getSize(com.google.protobuf.Message t, Class<?> type, java.lang.reflect.Type genericType, java.lang.annotation.Annotation[] annotations, javax.ws.rs.core.MediaType mediaType) {
            if (t == null) {
                return -1;
            }
            java.io.ByteArrayOutputStream out = new java.io.ByteArrayOutputStream();
            try {
                writeTo(t, type, genericType, annotations, mediaType, null, out);
            } catch (java.io.IOException e) {
                return -1;
            }
            return out.size();
        }

        @Override
        public void writeTo(com.google.protobuf.Message t, Class<?> type, java.lang.reflect.Type genericType, java.lang.annotation.Annotation[] annotations, javax.ws.rs.core.MediaType mediaType, javax.ws.rs.core.MultivaluedMap<String, Object> httpHeaders, java.io.OutputStream entityStream) throws java.io.IOException, javax.ws.rs.WebApplicationException {
            switch (mediaType.getSubtype()) {
                case "protobuf":
                    t.writeTo(entityStream);
                    break;
                case "json":
                    entityStream.write(com.google.protobuf.util.JsonFormat.printer().print(t).getBytes("UTF-8"));
                    break;
                default:
                    throw new javax.ws.rs.WebApplicationException("MediaType not supported!");
            }

        }

        @Override
        public boolean isReadable(Class<?> type, java.lang.reflect.Type genericType, java.lang.annotation.Annotation[] annotations, javax.ws.rs.core.MediaType mediaType) {
            return com.google.protobuf.Message.class.isAssignableFrom(type) && ("json".equals(mediaType.getSubtype()) || "protobuf".equals(mediaType.getSubtype()));
        }

        @Override
        public com.google.protobuf.Message readFrom(Class<com.google.protobuf.Message> type, java.lang.reflect.Type genericType, java.lang.annotation.Annotation[] annotations, javax.ws.rs.core.MediaType mediaType, javax.ws.rs.core.MultivaluedMap<String, String> httpHeaders, java.io.InputStream entityStream) throws java.io.IOException, javax.ws.rs.WebApplicationException {
            try {
                switch (mediaType.getSubtype()) {
                    case "protobuf":
                        java.lang.reflect.Method m = type.getMethod("parseFrom", java.io.InputStream.class);
                        return (com.google.protobuf.Message) m.invoke(null, entityStream);
                    case "json":
                        com.google.protobuf.Message.Builder msg = (com.google.protobuf.Message.Builder)type.getMethod("newBuilder").invoke(null);
                        com.google.protobuf.util.JsonFormat.parser().merge(new java.io.InputStreamReader(entityStream), msg);
                        return msg.build();
                    default:
                        throw new javax.ws.rs.WebApplicationException("MediaType not supported!");
                }
            } catch (Exception e) {
                throw new javax.ws.rs.WebApplicationException(e);
            }
        }
}
`)

	out := &plugin.CodeGeneratorResponse_File{}
	out.Content = proto.String(g.output.String())
	if multi {
		out.Name = proto.String(getJavaServiceClientClassFileByString(file, serviceName))
	} else {
		out.Name = proto.String(getJavaOuterClassFile(file))
		out.InsertionPoint = proto.String("outer_class_scope")
	}
	g.Reset()
	return out
}

func (g *generator) generateServiceClient(file *descriptor.FileDescriptorProto, service *descriptor.ServiceDescriptorProto) *plugin.CodeGeneratorResponse_File {
	multi := file.Options.GetJavaMultipleFiles()

	if multi {
		pkg := getJavaPackage(file)
		g.P(`// Code generated by protoc-gen-twirp_java_jaxrs, DO NOT EDIT.`)
		g.P(`// source: `, file.GetName())
		g.P()
		if pkg != "" {
			g.P(`package `, pkg, `;`)
			g.P()
		}
	}

	// TODO add comment

	serviceClass := getJavaServiceClientClassName(file, service)
	servicePath := g.getServicePath(file, service)
	interfaceClass := getJavaType(file, getJavaServiceClassName(file, service))
	provider := getJavaType(file, "ProtoBufMessageProvider")

	static := ""
	if !multi {
		static = "static "
	}
	g.P(`public `, static, `class `, serviceClass, ` implements `, interfaceClass, ` {`)
	g.P(`  private final javax.ws.rs.client.WebTarget target;`)
	g.P()
	g.P(`  public `, serviceClass, `(javax.ws.rs.client.WebTarget target) {`)
	g.P(`    this.target = target.path("`, servicePath, `");`)
	g.P(`    this.target.register(new `, provider, `());`)
	g.P(`  }`)
	g.P()
	g.P(`  private <R> R call(String path, com.google.protobuf.Message req, Class<R> responseClass) {`)
	g.P(`    javax.ws.rs.core.Response response = target.path(path)`)
	g.P(`        .request("application/protobuf")`)
	g.P(`        .post(javax.ws.rs.client.Entity.entity(req, "application/protobuf"));`)
	g.P(`	 if (response.getStatusInfo().getFamily() == javax.ws.rs.core.Response.Status.Family.SUCCESSFUL) {`)
	g.P(`        R r = response.readEntity(responseClass);`)
	g.P(`        response.close();`)
	g.P(`        return r;`)
	g.P(`    } else {`)
	g.P(`        throw new javax.ws.rs.WebApplicationException(response);`)
	g.P(`    }`)
	g.P(`  }`)
	g.P()

	for _, method := range service.GetMethod() {
		inputType := getJavaType(file, method.GetInputType())
		outputType := getJavaType(file, method.GetOutputType())
		methodName := lowerCamelCase(method.GetName())
		methodPath := camelCase(method.GetName())

		g.P()
		// add comment
		g.P(`  @Override`)
		g.P(`  public `, outputType, ` `, methodName, `(`, inputType, ` request) {`)
		g.P(`    return call("/`, methodPath, `", request, `, outputType, `.class);`)
		g.P(`  }`)
	}

	g.P(`}`)
	g.P()

	out := &plugin.CodeGeneratorResponse_File{}
	out.Content = proto.String(g.output.String())
	if multi {
		out.Name = proto.String(getJavaServiceClientClassFile(file, service))
	} else {
		out.Name = proto.String(getJavaOuterClassFile(file))
		out.InsertionPoint = proto.String("outer_class_scope")
	}
	g.Reset()

	return out
}

func (g *generator) generateServiceInterface(file *descriptor.FileDescriptorProto, service *descriptor.ServiceDescriptorProto) *plugin.CodeGeneratorResponse_File {
	// TODO add comment

	serviceClass := getJavaServiceClassName(file, service)
	servicePath := g.getServicePath(file, service)
	multi := file.Options.GetJavaMultipleFiles()

	if multi {
		pkg := getJavaPackage(file)
		g.P(`// Code generated by protoc-gen-twirp_java_jaxrs, DO NOT EDIT.`)
		g.P(`// source: `, file.GetName())
		g.P()
		if pkg != "" {
			g.P(`package `, pkg, `;`)
			g.P()
		}
	}

	g.P(`@javax.ws.rs.Path( "/`, servicePath, `" )`)
	g.P(`public interface `, serviceClass, ` {`)

	for _, method := range service.GetMethod() {
		inputType := getJavaType(file, method.GetInputType())
		outputType := getJavaType(file, method.GetOutputType())
		methodName := lowerCamelCase(method.GetName())

		// add comment
		g.P(`  `, `@javax.ws.rs.POST`)
		g.P(`  `, `@javax.ws.rs.Path( "/`, strings.Title(methodName), `" )`)
		g.P(`  `, `@javax.ws.rs.Consumes({"application/protobuf", "application/json"})`)
		g.P(`  `, `@javax.ws.rs.Produces({"application/protobuf", "application/json"})`)
		g.P(`  `, outputType, ` `, methodName, `(`, inputType, ` request);`)
	}

	g.P(`}`)
	g.P()

	out := &plugin.CodeGeneratorResponse_File{}
	out.Content = proto.String(g.output.String())

	if multi {
		out.Name = proto.String(getJavaServiceClassFile(file, service))
	} else {
		out.Name = proto.String(getJavaOuterClassFile(file))
		out.InsertionPoint = proto.String("outer_class_scope")
	}
	g.Reset()

	return out
}

func (g *generator) Reset() {
	g.indent = ""
	g.output.Reset()
}

func (g *generator) In() {
	g.indent += "  "
}

func (g *generator) Out() {
	g.indent = g.indent[2:]
}

func (g *generator) P(str ...string) {
	for _, v := range str {
		g.output.WriteString(v)
	}
	g.output.WriteByte('\n')
}

func (g *generator) getProtoFiles() []*descriptor.FileDescriptorProto {
	files := make([]*descriptor.FileDescriptorProto, 0)
	for _, fname := range g.Request.GetFileToGenerate() {
		for _, proto := range g.Request.GetProtoFile() {
			if proto.GetName() == fname {
				files = append(files, proto)
			}
		}
	}
	return files
}

func (g *generator) getServicePath(file *descriptor.FileDescriptorProto, service *descriptor.ServiceDescriptorProto) string {
	name := camelCase(service.GetName())
	pkg := file.GetPackage()
	if pkg != "" {
		name = pkg + "." + name
	}
	return name
}