<commit_before>// Copyright 2022 Lumberjack authors (see AUTHORS file)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package grpcrunner

import (
	"context"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"strings"
	"testing"
	"time"

	"cloud.google.com/go/bigquery"
	jvspb "github.com/abcxyz/jvs/apis/v0"
	"github.com/abcxyz/jvs/pkg/jvscrypto"
	"github.com/abcxyz/lumberjack/integration/testrunner/talkerpb"
	"github.com/abcxyz/lumberjack/integration/testrunner/utils"
	"github.com/golang-jwt/jwt"
	"github.com/google/uuid"
	"golang.org/x/oauth2"
	rpccode "google.golang.org/genproto/googleapis/rpc/code"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/oauth"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

type GRPC struct {
	ProjectID    string
	DatasetQuery string

	// IDToken and endpoint URL are required unless a custom TalkerClient is provided.
	IDToken      string
	EndpointURL  string
	TalkerClient talkerpb.TalkerClient

	Config               *utils.Config
	BigQueryClient       *bigquery.Client
	RequireJustification bool
}

// Matching public key here: https://github.com/abcxyz/lumberjack/pull/261/files#diff-f06321655121106c1e25ed2dd8cf773af6d7d5fb9b129abb2e1c04ba4d6dea5eR48
// and here: https://github.com/abcxyz/lumberjack/pull/264/files#diff-6c73b144011b225fa9f7602b77cc3bb4817b74b9852269e4f5f7625bd039ab4dR113
const privateKey = "-----<KEY>"

// TestGRPCEndpoint runs tests against a GRPC endpoint. The given GRPC must
// define a projectID and datasetQuery. If a TalkerClient or BigQueryClient are
// not provided, they are instantiated via the defaults.
func TestGRPCEndpoint(ctx context.Context, tb testing.TB, g *GRPC) {
	tb.Helper()

	if g.ProjectID == "" {
		tb.Fatal("ProjectID must be set")
	}

	if g.DatasetQuery == "" {
		tb.Fatal("DatasetQuery must be set")
	}

	if g.TalkerClient == nil {
		if g.IDToken == "" || g.EndpointURL == "" {
			tb.Fatal("IDToken and EndpointURL are required to create a TalkerClient")
		}

		conn := createConnection(tb, g.EndpointURL, g.IDToken)
		tb.Cleanup(func() {
			conn.Close()
		})
		g.TalkerClient = talkerpb.NewTalkerClient(conn)
	}

	signedToken, err := justificationToken()
	if err != nil {
		tb.Fatalf("couldn't generate justification token: %v", err)
	}

	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		md = metadata.New(map[string]string{})
	}
	md.Set("justification-token", signedToken)
	ctx = metadata.NewOutgoingContext(ctx, md)

	if g.BigQueryClient == nil {
		bqClient, err := utils.MakeClient(ctx, g.ProjectID)
		if err != nil {
			tb.Errorf("BigQuery request failed: %v", err)
		}
		tb.Cleanup(func() {
			if err := bqClient.Close(); err != nil {
				tb.Errorf("Failed to close the BQ client: %v", err)
			}
		})
		g.BigQueryClient = bqClient
	}

	g.runHelloCheck(ctx, tb)
	g.runFailCheck(ctx, tb)
	g.runFibonacciCheck(ctx, tb)
	g.runAdditionCheck(ctx, tb)
	g.runFailOnFourCheck(ctx, tb)
}

// create a justification token to pass in the call to services.
func justificationToken() (string, error) {
	now := time.Now().UTC()
	claims := jvspb.JVSClaims{
		StandardClaims: &jwt.StandardClaims{
			Audience:  "talker-app",
			ExpiresAt: now.Add(time.Hour).Unix(),
			Id:        uuid.New().String(),
			IssuedAt:  now.Unix(),
			Issuer:    "lumberjack-test-runner",
			NotBefore: now.Unix(),
			Subject:   "lumberjack-integ",
		},
		Justifications: []*jvspb.Justification{},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodES256, claims)
	token.Header["kid"] = "integ-key"

	block, _ := pem.Decode([]byte(privateKey))
	key, _ := x509.ParseECPrivateKey(block.Bytes)
	return jvscrypto.SignToken(token, key)
}

// End-to-end test for the fibonacci API, which is a test for server-side streaming.
func (g *GRPC) runFibonacciCheck(ctx context.Context, tb testing.TB) {
	tb.Helper()

	u := uuid.New()
	places := 5
	stream, err := g.TalkerClient.Fibonacci(ctx, &talkerpb.FibonacciRequest{Places: uint32(places), Target: u.String()})
	if err != nil {
		tb.Errorf("fibonacci call failed: %v", err)
	}
	for {
		place, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			// stream is finished
			break
		}
		if err != nil {
			tb.Errorf("Err while reading fibonacci stream: %v", err)
		}
		tb.Logf("Received value %v", place.Value)
	}
	query := g.makeQueryForGRPCStream(u)
	utils.QueryIfAuditLogsExistWithRetries(ctx, tb, query, g.Config, "server_stream_fibonacci", int64(places))
}

// End-to-end test for the addition API, which is a test for client-side streaming.
func (g *GRPC) runAdditionCheck(ctx context.Context, tb testing.TB) {
	tb.Helper()

	u := uuid.New()
	stream, err := g.TalkerClient.Addition(ctx)
	if err != nil {
		tb.Errorf("addition call failed: %v", err)<commit_msg>Fix a NPE bug in grpc test (#295)<commit_after>// Copyright 2022 Lumberjack authors (see AUTHORS file)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package grpcrunner

import (
	"context"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"strings"
	"testing"
	"time"

	"cloud.google.com/go/bigquery"
	jvspb "github.com/abcxyz/jvs/apis/v0"
	"github.com/abcxyz/jvs/pkg/jvscrypto"
	"github.com/abcxyz/lumberjack/integration/testrunner/talkerpb"
	"github.com/abcxyz/lumberjack/integration/testrunner/utils"
	"github.com/golang-jwt/jwt"
	"github.com/google/uuid"
	"golang.org/x/oauth2"
	rpccode "google.golang.org/genproto/googleapis/rpc/code"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/oauth"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

type GRPC struct {
	ProjectID    string
	DatasetQuery string

	// IDToken and endpoint URL are required unless a custom TalkerClient is provided.
	IDToken      string
	EndpointURL  string
	TalkerClient talkerpb.TalkerClient

	Config               *utils.Config
	BigQueryClient       *bigquery.Client
	RequireJustification bool
}

// Matching public key here: https://github.com/abcxyz/lumberjack/pull/261/files#diff-f06321655121106c1e25ed2dd8cf773af6d7d5fb9b129abb2e1c04ba4d6dea5eR48
// and here: https://github.com/abcxyz/lumberjack/pull/264/files#diff-6c73b144011b225fa9f7602b77cc3bb4817b74b9852269e4f5f7625bd039ab4dR113
const privateKey = "-----<KEY>"

// TestGRPCEndpoint runs tests against a GRPC endpoint. The given GRPC must
// define a projectID and datasetQuery. If a TalkerClient or BigQueryClient are
// not provided, they are instantiated via the defaults.
func TestGRPCEndpoint(ctx context.Context, tb testing.TB, g *GRPC) {
	tb.Helper()

	if g.ProjectID == "" {
		tb.Fatal("ProjectID must be set")
	}

	if g.DatasetQuery == "" {
		tb.Fatal("DatasetQuery must be set")
	}

	if g.TalkerClient == nil {
		if g.IDToken == "" || g.EndpointURL == "" {
			tb.Fatal("IDToken and EndpointURL are required to create a TalkerClient")
		}

		conn := createConnection(tb, g.EndpointURL, g.IDToken)
		tb.Cleanup(func() {
			conn.Close()
		})
		g.TalkerClient = talkerpb.NewTalkerClient(conn)
	}

	signedToken, err := justificationToken()
	if err != nil {
		tb.Fatalf("couldn't generate justification token: %v", err)
	}

	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		md = metadata.New(map[string]string{})
	}
	md.Set("justification-token", signedToken)
	ctx = metadata.NewOutgoingContext(ctx, md)

	if g.BigQueryClient == nil {
		bqClient, err := utils.MakeClient(ctx, g.ProjectID)
		if err != nil {
			tb.Errorf("BigQuery request failed: %v", err)
		}
		tb.Cleanup(func() {
			if err := bqClient.Close(); err != nil {
				tb.Errorf("Failed to close the BQ client: %v", err)
			}
		})
		g.BigQueryClient = bqClient
	}

	g.runHelloCheck(ctx, tb)
	g.runFailCheck(ctx, tb)
	g.runFibonacciCheck(ctx, tb)
	g.runAdditionCheck(ctx, tb)
	g.runFailOnFourCheck(ctx, tb)
}

// create a justification token to pass in the call to services.
func justificationToken() (string, error) {
	now := time.Now().UTC()
	claims := jvspb.JVSClaims{
		StandardClaims: &jwt.StandardClaims{
			Audience:  "talker-app",
			ExpiresAt: now.Add(time.Hour).Unix(),
			Id:        uuid.New().String(),
			IssuedAt:  now.Unix(),
			Issuer:    "lumberjack-test-runner",
			NotBefore: now.Unix(),
			Subject:   "lumberjack-integ",
		},
		Justifications: []*jvspb.Justification{},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodES256, claims)
	token.Header["kid"] = "integ-key"

	block, _ := pem.Decode([]byte(privateKey))
	key, _ := x509.ParseECPrivateKey(block.Bytes)
	return jvscrypto.SignToken(token, key)
}

// End-to-end test for the fibonacci API, which is a test for server-side streaming.
func (g *GRPC) runFibonacciCheck(ctx context.Context, tb testing.TB) {
	tb.Helper()

	u := uuid.New()
	places := 5
	stream, err := g.TalkerClient.Fibonacci(ctx, &talkerpb.FibonacciRequest{Places: uint32(places), Target: u.String()})
	if err != nil {
		tb.Errorf("fibonacci call failed: %v", err)
	}
	for {
		place, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			// stream is finished
			break
		}
		if err != nil {
			tb.Errorf("failed to read fibonacci stream: %v", err)
			break
		}
		tb.Logf("Received value %v", place.Value)
	}
	query := g.makeQueryForGRPCStream(u)
	utils.QueryIfAuditLogsExistWithRetries(ctx, tb, query, g.Config, "server_stream_fibonacci", int64(places))
}

// End-to-end test for the addition API, which is a test for client-side streaming.
func (g *GRPC) runAdditionCheck(ctx context.Context, tb testing.TB) {
	tb.Helper()

	u := uuid.New()
	stream, err := g.TalkerClient.Addition(ctx)
	if err != nil {
		tb.Errorf("addition call failed: %v", err)
{'fraction_non_alphanumeric': 0.13296057619408644, 'fraction_numerical': 0.02189158453373768, 'mean_word_length': 3.9872400756143667, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4614223', 'n_tokens_mistral': 4142, 'n_tokens_neox': 3790, 'n_words': 907}