63729 a23134379909 bfca3 ac3 c378 b48 (#452)

* Fix a bug in how we check for more than one extension in a probe config.

Use proto.ExtensionDescs(msg) to get extension descriptors instead of RegisteredExtensions(msg). Former returns extensions from the proto message, while the latter returns all extensions from the "type" of the proto message.

See #450 for more details.

PiperOrigin-RevId: 329415541

probes/probes.go

@@ -19,7 +19,6 @@ package probes
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"regexp"
 	"sync"
@@ -87,29 +86,23 @@ type ProbeInfo struct {
 }
 
 func getExtensionProbe(p *configpb.ProbeDef) (Probe, interface{}, error) {
-	extensions := proto.RegisteredExtensions(p)
-	if len(extensions) > 1 {
-		return nil, nil, fmt.Errorf("only one probe extension is allowed per probe, got %d extensions", len(extensions))
-	}
-	var field int
-	var desc *proto.ExtensionDesc
-	// There should be only one extension.
-	for f, d := range extensions {
-		field = int(f)
-		desc = d
+	extensions, err := proto.ExtensionDescs(p)
+	if err != nil {
+		return nil, nil, fmt.Errorf("error getting extensions from the probe config (%s): %v", p.String(), err)
 	}
-	if desc == nil {
-		return nil, nil, errors.New("no probe extension in probe config")
+	if len(extensions) != 1 {
+		return nil, nil, fmt.Errorf("there should be exactly one extension in the probe config (%s), got %d extensions", p.String(), len(extensions))
 	}
+	desc := extensions[0]
 	value, err := proto.GetExtension(p, desc)
 	if err != nil {
 		return nil, nil, err
 	}
 	extensionMapMu.Lock()
 	defer extensionMapMu.Unlock()
-	newProbeFunc, ok := extensionMap[field]
+	newProbeFunc, ok := extensionMap[int(desc.Field)]
 	if !ok {
-		return nil, nil, fmt.Errorf("no probes registered for the extension: %d", field)
+		return nil, nil, fmt.Errorf("no probes registered for the extension: %d", desc.Field)
 	}
 	return newProbeFunc(), value, nil
 }

probes/testdata/testdata.proto

@@ -7,9 +7,14 @@ import "github.com/google/cloudprober/probes/proto/config.proto";
 option go_package = "github.com/google/cloudprober/probes/testdata";
 
 message FancyProbe {
-  required string name = 1;
+  optional string name = 1;
+}
+
+message AnotherFancyProbe {
+  optional string name = 1;
 }
 
 extend probes.ProbeDef {
   optional FancyProbe fancy_probe = 200;
+  optional AnotherFancyProbe another_fancy_probe = 201;
 }

targets/targets.go

@@ -361,30 +361,24 @@ func New(targetsDef *targetspb.TargetsDef, ldLister lameduck.Lister, globalOpts
 }
 
 func getExtensionTargets(pb *targetspb.TargetsDef, l *logger.Logger) (Targets, error) {
-	extensions := proto.RegisteredExtensions(pb)
-	if len(extensions) > 1 {
-		return nil, fmt.Errorf("only one extension is allowed per targets definition, got %d extensions", len(extensions))
-	}
-	var field int
-	var desc *proto.ExtensionDesc
-	// There should be only one extension in one protobuf message.
-	for f, d := range extensions {
-		field = int(f)
-		desc = d
+	extensions, err := proto.ExtensionDescs(pb)
+	if err != nil {
+		return nil, fmt.Errorf("error getting extensions from the target config (%s): %v", pb.String(), err)
 	}
-	if desc == nil {
-		return nil, errors.New("unrecognized target type")
+	if len(extensions) != 1 {
+		return nil, fmt.Errorf("there should be exactly one extension in the targets config (%s), got %d extensions", pb.String(), len(extensions))
 	}
+	desc := extensions[0]
 	value, err := proto.GetExtension(pb, desc)
 	if err != nil {
 		return nil, err
 	}
-	l.Infof("Extension field: %d, value: %v", field, value)
+	l.Infof("Extension field: %d, value: %v", desc.Field, value)
 	extensionMapMu.Lock()
 	defer extensionMapMu.Unlock()
-	newTargetsFunc, ok := extensionMap[field]
+	newTargetsFunc, ok := extensionMap[int(desc.Field)]
 	if !ok {
-		return nil, fmt.Errorf("no targets type registered for the extension: %d", field)
+		return nil, fmt.Errorf("no targets type registered for the extension: %d", desc.Field)
 	}
 	return newTargetsFunc(value, l)
 }