1 // Copyright (c) 2009-2021, Google LLC
2 // All rights reserved.
3 //
4 // Redistribution and use in source and binary forms, with or without
5 // modification, are permitted provided that the following conditions are met:
6 // * Redistributions of source code must retain the above copyright
7 // notice, this list of conditions and the following disclaimer.
8 // * Redistributions in binary form must reproduce the above copyright
9 // notice, this list of conditions and the following disclaimer in the
10 // documentation and/or other materials provided with the distribution.
11 // * Neither the name of Google LLC nor the
12 // names of its contributors may be used to endorse or promote products
13 // derived from this software without specific prior written permission.
14 //
15 // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
16 // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
18 // ARE DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY DIRECT,
19 // INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
20 // (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
21 // LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
22 // ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
23 // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
24 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
25
26 #ifndef UPB_UPBC_PLUGIN_H_
27 #define UPB_UPBC_PLUGIN_H_
28
29 #include <stdio.h>
30
31 #include <string>
32 #include <vector>
33 #ifdef _WIN32
34 #include <fcntl.h>
35 #include <io.h>
36 #endif
37
38 // begin:google_only
39 // #ifndef UPB_BOOTSTRAP_STAGE0
40 // #include "net/proto2/compiler/proto/plugin.upb.h"
41 // #include "net/proto2/proto/descriptor.upb.h"
42 // #else
43 // #include "google/protobuf/compiler/plugin.upb.h"
44 // #include "google/protobuf/descriptor.upb.h"
45 // #endif
46 // end:google_only
47
48 // begin:github_only
49 #include "google/protobuf/compiler/plugin.upb.h"
50 #include "google/protobuf/descriptor.upb.h"
51 // end:github_only
52
53 #include "absl/container/flat_hash_set.h"
54 #include "absl/log/absl_log.h"
55 #include "absl/strings/str_split.h"
56 #include "absl/strings/string_view.h"
57 #include "upb/reflection/def.hpp"
58
59 // Must be last.
60 #include "upb/port/def.inc"
61
62 namespace upbc {
63
ParseGeneratorParameter(const absl::string_view text)64 inline std::vector<std::pair<std::string, std::string>> ParseGeneratorParameter(
65 const absl::string_view text) {
66 std::vector<std::pair<std::string, std::string>> ret;
67 for (absl::string_view sp : absl::StrSplit(text, ',', absl::SkipEmpty())) {
68 std::string::size_type equals_pos = sp.find_first_of('=');
69 std::pair<std::string, std::string> value;
70 if (equals_pos == std::string::npos) {
71 value.first = std::string(sp);
72 } else {
73 value.first = std::string(sp.substr(0, equals_pos));
74 value.second = std::string(sp.substr(equals_pos + 1));
75 }
76 ret.push_back(std::move(value));
77 }
78 return ret;
79 }
80
81 class Plugin {
82 public:
Plugin()83 Plugin() { ReadRequest(); }
~Plugin()84 ~Plugin() { WriteResponse(); }
85
parameter()86 absl::string_view parameter() const {
87 return ToStringView(
88 UPB_DESC(compiler_CodeGeneratorRequest_parameter)(request_));
89 }
90
91 template <class T>
GenerateFilesRaw(T && func)92 void GenerateFilesRaw(T&& func) {
93 absl::flat_hash_set<absl::string_view> files_to_generate;
94 size_t size;
95 const upb_StringView* file_to_generate = UPB_DESC(
96 compiler_CodeGeneratorRequest_file_to_generate)(request_, &size);
97 for (size_t i = 0; i < size; i++) {
98 files_to_generate.insert(
99 {file_to_generate[i].data, file_to_generate[i].size});
100 }
101
102 const UPB_DESC(FileDescriptorProto)* const* files =
103 UPB_DESC(compiler_CodeGeneratorRequest_proto_file)(request_, &size);
104 for (size_t i = 0; i < size; i++) {
105 upb::Status status;
106 absl::string_view name =
107 ToStringView(UPB_DESC(FileDescriptorProto_name)(files[i]));
108 func(files[i], files_to_generate.contains(name));
109 }
110 }
111
112 template <class T>
GenerateFiles(T && func)113 void GenerateFiles(T&& func) {
114 GenerateFilesRaw(
115 [this, &func](const UPB_DESC(FileDescriptorProto) * file_proto,
116 bool generate) {
117 upb::Status status;
118 upb::FileDefPtr file = pool_.AddFile(file_proto, &status);
119 if (!file) {
120 absl::string_view name =
121 ToStringView(UPB_DESC(FileDescriptorProto_name)(file_proto));
122 ABSL_LOG(FATAL) << "Couldn't add file " << name
123 << " to DefPool: " << status.error_message();
124 }
125 if (generate) func(file);
126 });
127 }
128
SetError(absl::string_view error)129 void SetError(absl::string_view error) {
130 char* data =
131 static_cast<char*>(upb_Arena_Malloc(arena_.ptr(), error.size()));
132 memcpy(data, error.data(), error.size());
133 UPB_DESC(compiler_CodeGeneratorResponse_set_error)
134 (response_, upb_StringView_FromDataAndSize(data, error.size()));
135 }
136
AddOutputFile(absl::string_view filename,absl::string_view content)137 void AddOutputFile(absl::string_view filename, absl::string_view content) {
138 UPB_DESC(compiler_CodeGeneratorResponse_File)* file = UPB_DESC(
139 compiler_CodeGeneratorResponse_add_file)(response_, arena_.ptr());
140 UPB_DESC(compiler_CodeGeneratorResponse_File_set_name)
141 (file, StringDup(filename));
142 UPB_DESC(compiler_CodeGeneratorResponse_File_set_content)
143 (file, StringDup(content));
144 }
145
146 private:
147 upb::Arena arena_;
148 upb::DefPool pool_;
149 UPB_DESC(compiler_CodeGeneratorRequest) * request_;
150 UPB_DESC(compiler_CodeGeneratorResponse) * response_;
151
ToStringView(upb_StringView sv)152 static absl::string_view ToStringView(upb_StringView sv) {
153 return absl::string_view(sv.data, sv.size);
154 }
155
StringDup(absl::string_view s)156 upb_StringView StringDup(absl::string_view s) {
157 char* data =
158 reinterpret_cast<char*>(upb_Arena_Malloc(arena_.ptr(), s.size()));
159 memcpy(data, s.data(), s.size());
160 return upb_StringView_FromDataAndSize(data, s.size());
161 }
162
ReadAllStdinBinary()163 std::string ReadAllStdinBinary() {
164 std::string data;
165 #ifdef _WIN32
166 _setmode(_fileno(stdin), _O_BINARY);
167 _setmode(_fileno(stdout), _O_BINARY);
168 #endif
169 char buf[4096];
170 while (size_t len = fread(buf, 1, sizeof(buf), stdin)) {
171 data.append(buf, len);
172 }
173 return data;
174 }
175
ReadRequest()176 void ReadRequest() {
177 std::string data = ReadAllStdinBinary();
178 request_ = UPB_DESC(compiler_CodeGeneratorRequest_parse)(
179 data.data(), data.size(), arena_.ptr());
180 if (!request_) {
181 ABSL_LOG(FATAL) << "Failed to parse CodeGeneratorRequest";
182 }
183 response_ = UPB_DESC(compiler_CodeGeneratorResponse_new)(arena_.ptr());
184 UPB_DESC(compiler_CodeGeneratorResponse_set_supported_features)
185 (response_,
186 UPB_DESC(compiler_CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL));
187 }
188
WriteResponse()189 void WriteResponse() {
190 size_t size;
191 char* serialized = UPB_DESC(compiler_CodeGeneratorResponse_serialize)(
192 response_, arena_.ptr(), &size);
193 if (!serialized) {
194 ABSL_LOG(FATAL) << "Failed to serialize CodeGeneratorResponse";
195 }
196
197 if (fwrite(serialized, 1, size, stdout) != size) {
198 ABSL_LOG(FATAL) << "Failed to write response to stdout";
199 }
200 }
201 };
202
203 } // namespace upbc
204
205 #include "upb/port/undef.inc"
206
207 #endif // UPB_UPBC_PLUGIN_H_
208