func doPerRPCCreds(tc testpb.TestServiceClient) {
jsonKey := getServiceAccountJSONKey()
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
FillUsername: proto.Bool(true),
FillOauthScope: proto.Bool(true),
}
token := getToken()
kv := map[string]string{"authorization": token.TokenType + " " + token.AccessToken}
ctx := metadata.NewContext(context.Background(), metadata.MD{"authorization": []string{kv["authorization"]}})
reply, err := tc.UnaryCall(ctx, req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
user := reply.GetUsername()
scope := reply.GetOauthScope()
if !strings.Contains(string(jsonKey), user) {
grpclog.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
}
if !strings.Contains(*oauthScope, scope) {
grpclog.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, *oauthScope)
}
grpclog.Println("PerRPCCreds done")
}
func doEmptyUnaryCall(tc testpb.TestServiceClient) {
reply, err := tc.EmptyCall(context.Background(), &testpb.Empty{})
if err != nil {
grpclog.Fatal("/TestService/EmptyCall RPC failed: ", err)
}
if !proto.Equal(&testpb.Empty{}, reply) {
grpclog.Fatalf("/TestService/EmptyCall receives %v, want %v", reply, testpb.Empty{})
}
grpclog.Println("EmptyUnaryCall done")
}
// DoUnaryCall performs an unary RPC with given stub and request and response sizes.
func DoUnaryCall(tc testpb.TestServiceClient, reqSize, respSize int) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, reqSize)
req := &testpb.SimpleRequest{
ResponseType: pl.Type,
ResponseSize: int32(respSize),
Payload: pl,
}
if _, err := tc.UnaryCall(context.Background(), req); err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
}
func doLargeUnaryCall(tc testpb.TestServiceClient) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
}
reply, err := tc.UnaryCall(context.Background(), req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
t := reply.GetPayload().GetType()
s := len(reply.GetPayload().GetBody())
if t != testpb.PayloadType_COMPRESSABLE || s != largeRespSize {
grpclog.Fatalf("Got the reply with type %d len %d; want %d, %d", t, s, testpb.PayloadType_COMPRESSABLE, largeRespSize)
}
grpclog.Println("LargeUnaryCall done")
}
func doJWTTokenCreds(tc testpb.TestServiceClient) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
FillUsername: proto.Bool(true),
}
reply, err := tc.UnaryCall(context.Background(), req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
jsonKey := getServiceAccountJSONKey()
user := reply.GetUsername()
if !strings.Contains(string(jsonKey), user) {
grpclog.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
}
grpclog.Println("JWTtokenCreds done")
}
func doComputeEngineCreds(tc testpb.TestServiceClient) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
FillUsername: proto.Bool(true),
FillOauthScope: proto.Bool(true),
}
reply, err := tc.UnaryCall(context.Background(), req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
user := reply.GetUsername()
scope := reply.GetOauthScope()
if user != *defaultServiceAccount {
grpclog.Fatalf("Got user name %q, want %q.", user, *defaultServiceAccount)
}
if !strings.Contains(*oauthScope, scope) {
grpclog.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, *oauthScope)
}
grpclog.Println("ComputeEngineCreds done")
}
func main() {
flag.Parse()
serverAddr := net.JoinHostPort(*serverHost, strconv.Itoa(*serverPort))
var opts []grpc.DialOption
if *useTLS {
var sn string
if *tlsServerName != "" {
sn = *tlsServerName
}
var creds credentials.TransportAuthenticator
if *caFile != "" {
var err error
creds, err = credentials.NewClientTLSFromFile(*caFile, sn)
if err != nil {
grpclog.Fatalf("Failed to create TLS credentials %v", err)
}
} else {
creds = credentials.NewClientTLSFromCert(nil, sn)
}
opts = append(opts, grpc.WithTransportCredentials(creds))
if *testCase == "compute_engine_creds" {
opts = append(opts, grpc.WithPerRPCCredentials(oauth.NewComputeEngine()))
} else if *testCase == "service_account_creds" {
jwtCreds, err := oauth.NewServiceAccountFromFile(*serviceAccountKeyFile, *oauthScope)
if err != nil {
grpclog.Fatalf("Failed to create JWT credentials: %v", err)
}
opts = append(opts, grpc.WithPerRPCCredentials(jwtCreds))
} else if *testCase == "jwt_token_creds" {
jwtCreds, err := oauth.NewJWTAccessFromFile(*serviceAccountKeyFile)
if err != nil {
grpclog.Fatalf("Failed to create JWT credentials: %v", err)
}
opts = append(opts, grpc.WithPerRPCCredentials(jwtCreds))
} else if *testCase == "oauth2_auth_token" {
opts = append(opts, grpc.WithPerRPCCredentials(oauth.NewOauthAccess(getToken())))
}
} else {
opts = append(opts, grpc.WithInsecure())
}
conn, err := grpc.Dial(serverAddr, opts...)
if err != nil {
grpclog.Fatalf("Fail to dial: %v", err)
}
defer conn.Close()
tc := testpb.NewTestServiceClient(conn)
switch *testCase {
case "empty_unary":
doEmptyUnaryCall(tc)
case "large_unary":
doLargeUnaryCall(tc)
case "client_streaming":
doClientStreaming(tc)
case "server_streaming":
doServerStreaming(tc)
case "ping_pong":
doPingPong(tc)
case "empty_stream":
doEmptyStream(tc)
case "timeout_on_sleeping_server":
doTimeoutOnSleepingServer(tc)
case "compute_engine_creds":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute compute_engine_creds test case.")
}
doComputeEngineCreds(tc)
case "service_account_creds":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute service_account_creds test case.")
}
doServiceAccountCreds(tc)
case "jwt_token_creds":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute jwt_token_creds test case.")
}
doJWTTokenCreds(tc)
case "per_rpc_creds":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute per_rpc_creds test case.")
}
doPerRPCCreds(tc)
case "oauth2_auth_token":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute oauth2_auth_token test case.")
}
doOauth2TokenCreds(tc)
case "cancel_after_begin":
doCancelAfterBegin(tc)
case "cancel_after_first_response":
doCancelAfterFirstResponse(tc)
default:
grpclog.Fatal("Unsupported test case: ", *testCase)
}
}