Matteo Scandolo | a6a3aee | 2019-11-26 13:30:14 -0700 | [diff] [blame] | 1 | package desc |
| 2 | |
| 3 | import ( |
| 4 | "fmt" |
| 5 | "path/filepath" |
| 6 | "reflect" |
| 7 | "strings" |
| 8 | "sync" |
| 9 | |
| 10 | "github.com/golang/protobuf/proto" |
| 11 | dpb "github.com/golang/protobuf/protoc-gen-go/descriptor" |
| 12 | ) |
| 13 | |
| 14 | var ( |
| 15 | globalImportPathConf map[string]string |
| 16 | globalImportPathMu sync.RWMutex |
| 17 | ) |
| 18 | |
| 19 | // RegisterImportPath registers an alternate import path for a given registered |
| 20 | // proto file path. For more details on why alternate import paths may need to |
| 21 | // be configured, see ImportResolver. |
| 22 | // |
| 23 | // This method panics if provided invalid input. An empty importPath is invalid. |
| 24 | // An un-registered registerPath is also invalid. For example, if an attempt is |
| 25 | // made to register the import path "foo/bar.proto" as "bar.proto", but there is |
| 26 | // no "bar.proto" registered in the Go protobuf runtime, this method will panic. |
| 27 | // This method also panics if an attempt is made to register the same import |
| 28 | // path more than once. |
| 29 | // |
| 30 | // This function works globally, applying to all descriptors loaded by this |
| 31 | // package. If you instead want more granular support for handling alternate |
| 32 | // import paths -- such as for a single invocation of a function in this |
| 33 | // package or when the alternate path is only used from one file (so you don't |
| 34 | // want the alternate path used when loading every other file), use an |
| 35 | // ImportResolver instead. |
| 36 | func RegisterImportPath(registerPath, importPath string) { |
| 37 | if len(importPath) == 0 { |
| 38 | panic("import path cannot be empty") |
| 39 | } |
| 40 | desc := proto.FileDescriptor(registerPath) |
| 41 | if len(desc) == 0 { |
| 42 | panic(fmt.Sprintf("path %q is not a registered proto file", registerPath)) |
| 43 | } |
| 44 | globalImportPathMu.Lock() |
| 45 | defer globalImportPathMu.Unlock() |
| 46 | if reg := globalImportPathConf[importPath]; reg != "" { |
| 47 | panic(fmt.Sprintf("import path %q already registered for %s", importPath, reg)) |
| 48 | } |
| 49 | if globalImportPathConf == nil { |
| 50 | globalImportPathConf = map[string]string{} |
| 51 | } |
| 52 | globalImportPathConf[importPath] = registerPath |
| 53 | } |
| 54 | |
| 55 | // ResolveImport resolves the given import path. If it has been registered as an |
| 56 | // alternate via RegisterImportPath, the registered path is returned. Otherwise, |
| 57 | // the given import path is returned unchanged. |
| 58 | func ResolveImport(importPath string) string { |
| 59 | importPath = clean(importPath) |
| 60 | globalImportPathMu.RLock() |
| 61 | defer globalImportPathMu.RUnlock() |
| 62 | reg := globalImportPathConf[importPath] |
| 63 | if reg == "" { |
| 64 | return importPath |
| 65 | } |
| 66 | return reg |
| 67 | } |
| 68 | |
| 69 | // ImportResolver lets you work-around linking issues that are caused by |
| 70 | // mismatches between how a particular proto source file is registered in the Go |
| 71 | // protobuf runtime and how that same file is imported by other files. The file |
| 72 | // is registered using the same relative path given to protoc when the file is |
| 73 | // compiled (i.e. when Go code is generated). So if any file tries to import |
| 74 | // that source file, but using a different relative path, then a link error will |
| 75 | // occur when this package tries to load a descriptor for the importing file. |
| 76 | // |
| 77 | // For example, let's say we have two proto source files: "foo/bar.proto" and |
| 78 | // "fubar/baz.proto". The latter imports the former using a line like so: |
| 79 | // import "foo/bar.proto"; |
| 80 | // However, when protoc is invoked, the command-line args looks like so: |
| 81 | // protoc -Ifoo/ --go_out=foo/ bar.proto |
| 82 | // protoc -I./ -Ifubar/ --go_out=fubar/ baz.proto |
| 83 | // Because the path given to protoc is just "bar.proto" and "baz.proto", this is |
| 84 | // how they are registered in the Go protobuf runtime. So, when loading the |
| 85 | // descriptor for "fubar/baz.proto", we'll see an import path of "foo/bar.proto" |
| 86 | // but will find no file registered with that path: |
| 87 | // fd, err := desc.LoadFileDescriptor("baz.proto") |
| 88 | // // err will be non-nil, complaining that there is no such file |
| 89 | // // found named "foo/bar.proto" |
| 90 | // |
| 91 | // This can be remedied by registering alternate import paths using an |
| 92 | // ImportResolver. Continuing with the example above, the code below would fix |
| 93 | // any link issue: |
| 94 | // var r desc.ImportResolver |
| 95 | // r.RegisterImportPath("bar.proto", "foo/bar.proto") |
| 96 | // fd, err := r.LoadFileDescriptor("baz.proto") |
| 97 | // // err will be nil; descriptor successfully loaded! |
| 98 | // |
| 99 | // If there are files that are *always* imported using a different relative |
| 100 | // path then how they are registered, consider using the global |
| 101 | // RegisterImportPath function, so you don't have to use an ImportResolver for |
| 102 | // every file that imports it. |
| 103 | type ImportResolver struct { |
| 104 | children map[string]*ImportResolver |
| 105 | importPaths map[string]string |
| 106 | |
| 107 | // By default, an ImportResolver will fallback to consulting any paths |
| 108 | // registered via the top-level RegisterImportPath function. Setting this |
| 109 | // field to true will cause the ImportResolver to skip that fallback and |
| 110 | // only examine its own locally registered paths. |
| 111 | SkipFallbackRules bool |
| 112 | } |
| 113 | |
| 114 | // ResolveImport resolves the given import path in the context of the given |
| 115 | // source file. If a matching alternate has been registered with this resolver |
| 116 | // via a call to RegisterImportPath or RegisterImportPathFrom, then the |
| 117 | // registered path is returned. Otherwise, the given import path is returned |
| 118 | // unchanged. |
| 119 | func (r *ImportResolver) ResolveImport(source, importPath string) string { |
| 120 | if r != nil { |
| 121 | res := r.resolveImport(clean(source), clean(importPath)) |
| 122 | if res != "" { |
| 123 | return res |
| 124 | } |
| 125 | if r.SkipFallbackRules { |
| 126 | return importPath |
| 127 | } |
| 128 | } |
| 129 | return ResolveImport(importPath) |
| 130 | } |
| 131 | |
| 132 | func (r *ImportResolver) resolveImport(source, importPath string) string { |
| 133 | if source == "" { |
| 134 | return r.importPaths[importPath] |
| 135 | } |
| 136 | var car, cdr string |
| 137 | idx := strings.IndexRune(source, filepath.Separator) |
| 138 | if idx < 0 { |
| 139 | car, cdr = source, "" |
| 140 | } else { |
| 141 | car, cdr = source[:idx], source[idx+1:] |
| 142 | } |
| 143 | ch := r.children[car] |
| 144 | if ch != nil { |
| 145 | if reg := ch.resolveImport(cdr, importPath); reg != "" { |
| 146 | return reg |
| 147 | } |
| 148 | } |
| 149 | return r.importPaths[importPath] |
| 150 | } |
| 151 | |
| 152 | // RegisterImportPath registers an alternate import path for a given registered |
| 153 | // proto file path with this resolver. Any appearance of the given import path |
| 154 | // when linking files will instead try to link the given registered path. If the |
| 155 | // registered path cannot be located, then linking will fallback to the actual |
| 156 | // imported path. |
| 157 | // |
| 158 | // This method will panic if given an empty path or if the same import path is |
| 159 | // registered more than once. |
| 160 | // |
| 161 | // To constrain the contexts where the given import path is to be re-written, |
| 162 | // use RegisterImportPathFrom instead. |
| 163 | func (r *ImportResolver) RegisterImportPath(registerPath, importPath string) { |
| 164 | r.RegisterImportPathFrom(registerPath, importPath, "") |
| 165 | } |
| 166 | |
| 167 | // RegisterImportPathFrom registers an alternate import path for a given |
| 168 | // registered proto file path with this resolver, but only for imports in the |
| 169 | // specified source context. |
| 170 | // |
| 171 | // The source context can be the name of a folder or a proto source file. Any |
| 172 | // appearance of the given import path in that context will instead try to link |
| 173 | // the given registered path. To be in context, the file that is being linked |
| 174 | // (i.e. the one whose import statement is being resolved) must be the same |
| 175 | // relative path of the source context or be a sub-path (i.e. a descendant of |
| 176 | // the source folder). |
| 177 | // |
| 178 | // If the registered path cannot be located, then linking will fallback to the |
| 179 | // actual imported path. |
| 180 | // |
| 181 | // This method will panic if given an empty path. The source context, on the |
| 182 | // other hand, is allowed to be blank. A blank source matches all files. This |
| 183 | // method also panics if the same import path is registered in the same source |
| 184 | // context more than once. |
| 185 | func (r *ImportResolver) RegisterImportPathFrom(registerPath, importPath, source string) { |
| 186 | importPath = clean(importPath) |
| 187 | if len(importPath) == 0 { |
| 188 | panic("import path cannot be empty") |
| 189 | } |
| 190 | registerPath = clean(registerPath) |
| 191 | if len(registerPath) == 0 { |
| 192 | panic("registered path cannot be empty") |
| 193 | } |
| 194 | r.registerImportPathFrom(registerPath, importPath, clean(source)) |
| 195 | } |
| 196 | |
| 197 | func (r *ImportResolver) registerImportPathFrom(registerPath, importPath, source string) { |
| 198 | if source == "" { |
| 199 | if r.importPaths == nil { |
| 200 | r.importPaths = map[string]string{} |
| 201 | } else if reg := r.importPaths[importPath]; reg != "" { |
| 202 | panic(fmt.Sprintf("already registered import path %q as %q", importPath, registerPath)) |
| 203 | } |
| 204 | r.importPaths[importPath] = registerPath |
| 205 | return |
| 206 | } |
| 207 | var car, cdr string |
| 208 | idx := strings.IndexRune(source, filepath.Separator) |
| 209 | if idx < 0 { |
| 210 | car, cdr = source, "" |
| 211 | } else { |
| 212 | car, cdr = source[:idx], source[idx+1:] |
| 213 | } |
| 214 | ch := r.children[car] |
| 215 | if ch == nil { |
| 216 | if r.children == nil { |
| 217 | r.children = map[string]*ImportResolver{} |
| 218 | } |
| 219 | ch = &ImportResolver{} |
| 220 | r.children[car] = ch |
| 221 | } |
| 222 | ch.registerImportPathFrom(registerPath, importPath, cdr) |
| 223 | } |
| 224 | |
| 225 | // LoadFileDescriptor is the same as the package function of the same name, but |
| 226 | // any alternate paths configured in this resolver are used when linking the |
| 227 | // given descriptor proto. |
| 228 | func (r *ImportResolver) LoadFileDescriptor(filePath string) (*FileDescriptor, error) { |
| 229 | return loadFileDescriptor(filePath, r) |
| 230 | } |
| 231 | |
| 232 | // LoadMessageDescriptor is the same as the package function of the same name, |
| 233 | // but any alternate paths configured in this resolver are used when linking |
| 234 | // files for the returned descriptor. |
| 235 | func (r *ImportResolver) LoadMessageDescriptor(msgName string) (*MessageDescriptor, error) { |
| 236 | return loadMessageDescriptor(msgName, r) |
| 237 | } |
| 238 | |
| 239 | // LoadMessageDescriptorForMessage is the same as the package function of the |
| 240 | // same name, but any alternate paths configured in this resolver are used when |
| 241 | // linking files for the returned descriptor. |
| 242 | func (r *ImportResolver) LoadMessageDescriptorForMessage(msg proto.Message) (*MessageDescriptor, error) { |
| 243 | return loadMessageDescriptorForMessage(msg, r) |
| 244 | } |
| 245 | |
| 246 | // LoadMessageDescriptorForType is the same as the package function of the same |
| 247 | // name, but any alternate paths configured in this resolver are used when |
| 248 | // linking files for the returned descriptor. |
| 249 | func (r *ImportResolver) LoadMessageDescriptorForType(msgType reflect.Type) (*MessageDescriptor, error) { |
| 250 | return loadMessageDescriptorForType(msgType, r) |
| 251 | } |
| 252 | |
| 253 | // LoadEnumDescriptorForEnum is the same as the package function of the same |
| 254 | // name, but any alternate paths configured in this resolver are used when |
| 255 | // linking files for the returned descriptor. |
| 256 | func (r *ImportResolver) LoadEnumDescriptorForEnum(enum protoEnum) (*EnumDescriptor, error) { |
| 257 | return loadEnumDescriptorForEnum(enum, r) |
| 258 | } |
| 259 | |
| 260 | // LoadEnumDescriptorForType is the same as the package function of the same |
| 261 | // name, but any alternate paths configured in this resolver are used when |
| 262 | // linking files for the returned descriptor. |
| 263 | func (r *ImportResolver) LoadEnumDescriptorForType(enumType reflect.Type) (*EnumDescriptor, error) { |
| 264 | return loadEnumDescriptorForType(enumType, r) |
| 265 | } |
| 266 | |
| 267 | // LoadFieldDescriptorForExtension is the same as the package function of the |
| 268 | // same name, but any alternate paths configured in this resolver are used when |
| 269 | // linking files for the returned descriptor. |
| 270 | func (r *ImportResolver) LoadFieldDescriptorForExtension(ext *proto.ExtensionDesc) (*FieldDescriptor, error) { |
| 271 | return loadFieldDescriptorForExtension(ext, r) |
| 272 | } |
| 273 | |
| 274 | // CreateFileDescriptor is the same as the package function of the same name, |
| 275 | // but any alternate paths configured in this resolver are used when linking the |
| 276 | // given descriptor proto. |
| 277 | func (r *ImportResolver) CreateFileDescriptor(fdp *dpb.FileDescriptorProto, deps ...*FileDescriptor) (*FileDescriptor, error) { |
| 278 | return createFileDescriptor(fdp, deps, r) |
| 279 | } |
| 280 | |
| 281 | // CreateFileDescriptors is the same as the package function of the same name, |
| 282 | // but any alternate paths configured in this resolver are used when linking the |
| 283 | // given descriptor protos. |
| 284 | func (r *ImportResolver) CreateFileDescriptors(fds []*dpb.FileDescriptorProto) (map[string]*FileDescriptor, error) { |
| 285 | return createFileDescriptors(fds, r) |
| 286 | } |
| 287 | |
| 288 | // CreateFileDescriptorFromSet is the same as the package function of the same |
| 289 | // name, but any alternate paths configured in this resolver are used when |
| 290 | // linking the descriptor protos in the given set. |
| 291 | func (r *ImportResolver) CreateFileDescriptorFromSet(fds *dpb.FileDescriptorSet) (*FileDescriptor, error) { |
| 292 | return createFileDescriptorFromSet(fds, r) |
| 293 | } |
| 294 | |
| 295 | // CreateFileDescriptorsFromSet is the same as the package function of the same |
| 296 | // name, but any alternate paths configured in this resolver are used when |
| 297 | // linking the descriptor protos in the given set. |
| 298 | func (r *ImportResolver) CreateFileDescriptorsFromSet(fds *dpb.FileDescriptorSet) (map[string]*FileDescriptor, error) { |
| 299 | return createFileDescriptorsFromSet(fds, r) |
| 300 | } |
| 301 | |
| 302 | const dotPrefix = "." + string(filepath.Separator) |
| 303 | |
| 304 | func clean(path string) string { |
| 305 | if path == "" { |
| 306 | return "" |
| 307 | } |
| 308 | path = filepath.Clean(path) |
| 309 | if path == "." { |
| 310 | return "" |
| 311 | } |
| 312 | return strings.TrimPrefix(path, dotPrefix) |
| 313 | } |