go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/buildbucket/appengine/internal/clients/pubsub.go (about)

     1  // Copyright 2022 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package clients
    16  
    17  import (
    18  	"context"
    19  	"regexp"
    20  	"strings"
    21  
    22  	"cloud.google.com/go/pubsub"
    23  	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
    24  	"google.golang.org/api/option"
    25  	"google.golang.org/grpc"
    26  	"google.golang.org/grpc/credentials"
    27  
    28  	"go.chromium.org/luci/common/errors"
    29  	"go.chromium.org/luci/grpc/grpcmon"
    30  	"go.chromium.org/luci/server/auth"
    31  )
    32  
    33  var (
    34  	mockPubsubClientKey = "mock pubsub clients key for testing only"
    35  
    36  	// cloudProjectIDRE is the cloud project identifier regex derived from
    37  	// https://cloud.google.com/resource-manager/docs/creating-managing-projects#before_you_begin
    38  	cloudProjectIDRE = regexp.MustCompile(`^[a-z]([a-z0-9-]){4,28}[a-z0-9]$`)
    39  	// topicNameRE is the full topic name regex derived from https://cloud.google.com/pubsub/docs/admin#resource_names
    40  	topicNameRE = regexp.MustCompile(`^projects/(.*)/topics/(.*)$`)
    41  	// topicIDRE is the topic id regex derived from https://cloud.google.com/pubsub/docs/admin#resource_names
    42  	topicIDRE = regexp.MustCompile(`^[A-Za-z]([0-9A-Za-z\._\-~+%]){3,255}$`)
    43  )
    44  
    45  // NewPubsubClient creates a pubsub client with the authority of a given
    46  // luciProject or the current service if luciProject is empty.
    47  func NewPubsubClient(ctx context.Context, cloudProject, luciProject string) (*pubsub.Client, error) {
    48  	if mockClients, ok := ctx.Value(&mockPubsubClientKey).(map[string]*pubsub.Client); ok {
    49  		if mockClient, exist := mockClients[cloudProject]; exist {
    50  			return mockClient, nil
    51  		}
    52  		return nil, errors.Reason("couldn't find mock pubsub client for %s", cloudProject).Err()
    53  	}
    54  
    55  	var creds credentials.PerRPCCredentials
    56  	var err error
    57  	if luciProject == "" {
    58  		creds, err = auth.GetPerRPCCredentials(ctx, auth.AsSelf, auth.WithScopes(auth.CloudOAuthScopes...))
    59  	} else {
    60  		creds, err = auth.GetPerRPCCredentials(ctx, auth.AsProject, auth.WithProject(luciProject), auth.WithScopes(auth.CloudOAuthScopes...))
    61  	}
    62  	if err != nil {
    63  		return nil, err
    64  	}
    65  	client, err := pubsub.NewClient(
    66  		ctx, cloudProject,
    67  		option.WithGRPCDialOption(grpc.WithStatsHandler(&grpcmon.ClientRPCStatsMonitor{})),
    68  		option.WithGRPCDialOption(grpc.WithUnaryInterceptor(otelgrpc.UnaryClientInterceptor())),
    69  		option.WithGRPCDialOption(grpc.WithStreamInterceptor(otelgrpc.StreamClientInterceptor())),
    70  		option.WithGRPCDialOption(grpc.WithPerRPCCredentials(creds)),
    71  	)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  	return client, nil
    76  }
    77  
    78  // ValidatePubSubTopicName validates the format of topic, extract the cloud project and topic id, and return them.
    79  func ValidatePubSubTopicName(topic string) (string, string, error) {
    80  	matches := topicNameRE.FindAllStringSubmatch(topic, -1)
    81  	if matches == nil || len(matches[0]) != 3 {
    82  		return "", "", errors.Reason("topic %q does not match %q", topic, topicNameRE).Err()
    83  	}
    84  
    85  	cloudProj := matches[0][1]
    86  	topicID := matches[0][2]
    87  	// Only internal App Engine projects start "google.com:" with go/gae4g-setup#choosing-the-right-app-engine-version,
    88  	// all other project ids conform to cloudProjectIDRE.
    89  	if !strings.HasPrefix(cloudProj, "google.com:") && !cloudProjectIDRE.MatchString(cloudProj) {
    90  		return "", "", errors.Reason("cloud project id %q does not match %q", cloudProj, cloudProjectIDRE).Err()
    91  	}
    92  	if strings.HasPrefix(topicID, "goog") {
    93  		return "", "", errors.Reason("topic id %q shouldn't begin with the string goog", topicID).Err()
    94  	}
    95  	if !topicIDRE.MatchString(topicID) {
    96  		return "", "", errors.Reason("topic id %q does not match %q", topicID, topicIDRE).Err()
    97  	}
    98  	return cloudProj, topicID, nil
    99  }