auth/credentials/idtoken/validate_test.go - gocloud - Git at Google

 // Copyright 2023 Google LLC
 //
 // 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 idtoken

 import (
 	"bytes"
 	"context"
 	"crypto"
 	"crypto/ecdsa"
 	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/sha256"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io"
 	"math/big"
 	"net/http"
 	"testing"
 	"time"

 	"cloud.google.com/go/auth/internal/jwt"
 )

 const (
 	keyID              = "1234"
 	testAudience       = "test-audience"
 	expiry       int64 = 233431200
 )

 var (
 	beforeExp = func() time.Time { return time.Unix(expiry-1, 0) }
 	afterExp  = func() time.Time { return time.Unix(expiry+1, 0) }
 )

 func TestValidateRS256(t *testing.T) {
 	idToken, pk := createRS256JWT(t)
 	tests := []struct {
 		name    string
 		keyID   string
 		n       *big.Int
 		e       int
 		nowFunc func() time.Time
 		wantErr bool
 	}{
 		{
 			name:    "works",
 			keyID:   keyID,
 			n:       pk.N,
 			e:       pk.E,
 			nowFunc: beforeExp,
 			wantErr: false,
 		},
 		{
 			name:    "no matching key",
 			keyID:   "5678",
 			n:       pk.N,
 			e:       pk.E,
 			nowFunc: beforeExp,
 			wantErr: true,
 		},
 		{
 			name:    "sig does not match",
 			keyID:   keyID,
 			n:       new(big.Int).SetBytes([]byte("42")),
 			e:       42,
 			nowFunc: beforeExp,
 			wantErr: true,
 		},
 		{
 			name:    "token expired",
 			keyID:   keyID,
 			n:       pk.N,
 			e:       pk.E,
 			nowFunc: afterExp,
 			wantErr: true,
 		},
 	}

 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			client := &http.Client{
 				Transport: RoundTripFn(func(req *http.Request) *http.Response {
 					cr := certResponse{
 						Keys: []jwk{
 							{
 								Kid: tt.keyID,
 								N:   base64.RawURLEncoding.EncodeToString(tt.n.Bytes()),
 								E:   base64.RawURLEncoding.EncodeToString(new(big.Int).SetInt64(int64(tt.e)).Bytes()),
 							},
 						},
 					}
 					b, err := json.Marshal(&cr)
 					if err != nil {
 						t.Fatalf("unable to marshal response: %v", err)
 					}
 					return &http.Response{
 						StatusCode: 200,
 						Body:       io.NopCloser(bytes.NewReader(b)),
 						Header:     make(http.Header),
 					}
 				}),
 			}
 			oldNow := now
 			defer func() { now = oldNow }()
 			now = tt.nowFunc

 			v, err := NewValidator(&ValidatorOptions{
 				Client: client,
 			})
 			if err != nil {
 				t.Fatalf("NewValidator(...) = %q, want nil", err)
 			}
 			payload, err := v.Validate(context.Background(), idToken, testAudience)
 			if tt.wantErr && err != nil {
 				// Got the error we wanted.
 				return
 			}
 			if !tt.wantErr && err != nil {
 				t.Fatalf("Validate(ctx, %s, %s): got err %q, want nil", idToken, testAudience, err)
 			}
 			if tt.wantErr && err == nil {
 				t.Fatalf("Validate(ctx, %s, %s): got nil err, want err", idToken, testAudience)
 			}
 			if payload == nil {
 				t.Fatalf("Got nil payload, err: %v", err)
 			}
 			if payload.Audience != testAudience {
 				t.Fatalf("Validate(ctx, %s, %s): got %v, want %v", idToken, testAudience, payload.Audience, testAudience)
 			}
 			if len(payload.Claims) == 0 {
 				t.Fatalf("Validate(ctx, %s, %s): missing Claims map. payload.Claims = %+v", idToken, testAudience, payload.Claims)
 			}
 			if got, ok := payload.Claims["aud"]; !ok {
 				t.Fatalf("Validate(ctx, %s, %s): missing aud claim. payload.Claims = %+v", idToken, testAudience, payload.Claims)
 			} else {
 				got, ok := got.(string)
 				if !ok {
 					t.Fatalf("Validate(ctx, %s, %s): aud wasn't a string. payload.Claims = %+v", idToken, testAudience, payload.Claims)
 				}
 				if got != testAudience {
 					t.Fatalf("Validate(ctx, %s, %s): Payload[aud] want %v got %v", idToken, testAudience, testAudience, got)
 				}
 			}
 		})
 	}
 }

 func TestValidateES256(t *testing.T) {
 	idToken, pk := createES256JWT(t)
 	tests := []struct {
 		name    string
 		keyID   string
 		x       *big.Int
 		y       *big.Int
 		nowFunc func() time.Time
 		wantErr bool
 	}{
 		{
 			name:    "works",
 			keyID:   keyID,
 			x:       pk.X,
 			y:       pk.Y,
 			nowFunc: beforeExp,
 			wantErr: false,
 		},
 		{
 			name:    "no matching key",
 			keyID:   "5678",
 			x:       pk.X,
 			y:       pk.Y,
 			nowFunc: beforeExp,
 			wantErr: true,
 		},
 		{
 			name:    "sig does not match",
 			keyID:   keyID,
 			x:       new(big.Int),
 			y:       new(big.Int),
 			nowFunc: beforeExp,
 			wantErr: true,
 		},
 		{
 			name:    "token expired",
 			keyID:   keyID,
 			x:       pk.X,
 			y:       pk.Y,
 			nowFunc: afterExp,
 			wantErr: true,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			client := &http.Client{
 				Transport: RoundTripFn(func(req *http.Request) *http.Response {
 					cr := certResponse{
 						Keys: []jwk{
 							{
 								Kid: tt.keyID,
 								X:   base64.RawURLEncoding.EncodeToString(tt.x.Bytes()),
 								Y:   base64.RawURLEncoding.EncodeToString(tt.y.Bytes()),
 							},
 						},
 					}
 					b, err := json.Marshal(&cr)
 					if err != nil {
 						t.Fatalf("unable to marshal response: %v", err)
 					}
 					return &http.Response{
 						StatusCode: 200,
 						Body:       io.NopCloser(bytes.NewReader(b)),
 						Header:     make(http.Header),
 					}
 				}),
 			}
 			oldNow := now
 			defer func() { now = oldNow }()
 			now = tt.nowFunc

 			v, err := NewValidator(&ValidatorOptions{
 				Client: client,
 			})
 			if err != nil {
 				t.Fatalf("NewValidator(...) = %q, want nil", err)
 			}
 			payload, err := v.Validate(context.Background(), idToken, testAudience)
 			if !tt.wantErr && err != nil {
 				t.Fatalf("Validate(ctx, %s, %s) = %q, want nil", idToken, testAudience, err)
 			}
 			if !tt.wantErr && payload.Audience != testAudience {
 				t.Fatalf("got %v, want %v", payload.Audience, testAudience)
 			}
 		})
 	}
 }

 func TestParsePayload(t *testing.T) {
 	idToken, _ := createRS256JWT(t)
 	tests := []struct {
 		name                string
 		token               string
 		wantPayloadAudience string
 		wantErr             bool
 	}{{
 		name:                "valid token",
 		token:               idToken,
 		wantPayloadAudience: testAudience,
 	}, {
 		name:    "unparseable token",
 		token:   "aaa.bbb.ccc",
 		wantErr: true,
 	}}

 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			payload, err := ParsePayload(tt.token)
 			gotErr := err != nil
 			if gotErr != tt.wantErr {
 				t.Errorf("ParsePayload(%q) got error %v, wantErr = %v", tt.token, err, tt.wantErr)
 			}
 			if tt.wantPayloadAudience != "" {
 				if payload == nil || payload.Audience != tt.wantPayloadAudience {
 					t.Errorf("ParsePayload(%q) got payload %+v, want payload with audience = %q", tt.token, payload, tt.wantPayloadAudience)
 				}
 			}
 		})
 	}
 }

 func createES256JWT(t *testing.T) (string, ecdsa.PublicKey) {
 	t.Helper()
 	header, claims := commonToken(t, "ES256")
 	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 	if err != nil {
 		t.Fatalf("unable to generate key: %v", err)
 	}
 	signedContent := header + "." + claims
 	hashed := sha256.Sum256([]byte(signedContent))
 	hash := hashed[:]
 	r, s, err := ecdsa.Sign(rand.Reader, privateKey, hash)
 	if err != nil {
 		t.Fatalf("unable to sign content: %v", err)
 	}
 	rb := r.Bytes()
 	lPadded := make([]byte, es256KeySize)
 	copy(lPadded[es256KeySize-len(rb):], rb)
 	var sig []byte
 	sig = append(sig, lPadded...)
 	sig = append(sig, s.Bytes()...)
 	signature := base64.RawURLEncoding.EncodeToString(sig)
 	return fmt.Sprintf("%s.%s.%s", header, claims, signature), privateKey.PublicKey
 }

 func createRS256JWT(t *testing.T) (string, rsa.PublicKey) {
 	t.Helper()
 	header, claims := commonToken(t, jwt.HeaderAlgRSA256)
 	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
 	if err != nil {
 		t.Fatalf("unable to generate key: %v", err)
 	}
 	signedContent := header + "." + claims
 	hashed := sha256.Sum256([]byte(signedContent))
 	hash := hashed[:]
 	sig, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hash)
 	if err != nil {
 		t.Fatalf("unable to sign content: %v", err)
 	}
 	signature := base64.RawURLEncoding.EncodeToString(sig)
 	return fmt.Sprintf("%s.%s.%s", header, claims, signature), privateKey.PublicKey
 }

 // returns header and claims
 func commonToken(t *testing.T, alg string) (string, string) {
 	t.Helper()
 	header := jwt.Header{
 		KeyID:     keyID,
 		Algorithm: alg,
 		Type:      jwt.HeaderType,
 	}
 	payload := Payload{
 		Issuer:   "example.com",
 		Audience: testAudience,
 		Expires:  expiry,
 	}

 	hb, err := json.Marshal(&header)
 	if err != nil {
 		t.Fatalf("unable to marshall header: %v", err)
 	}
 	pb, err := json.Marshal(&payload)
 	if err != nil {
 		t.Fatalf("unable to marshall payload: %v", err)
 	}
 	eb := base64.RawURLEncoding.EncodeToString(hb)
 	ep := base64.RawURLEncoding.EncodeToString(pb)
 	return eb, ep
 }

 type RoundTripFn func(req *http.Request) *http.Response

 func (f RoundTripFn) RoundTrip(req *http.Request) (*http.Response, error) { return f(req), nil }
	// Copyright 2023 Google LLC
	//
	// 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 idtoken

	import (
	"bytes"
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"testing"
	"time"

	"cloud.google.com/go/auth/internal/jwt"
	)

	const (
	keyID = "1234"
	testAudience = "test-audience"
	expiry int64 = 233431200
	)

	var (
	beforeExp = func() time.Time { return time.Unix(expiry-1, 0) }
	afterExp = func() time.Time { return time.Unix(expiry+1, 0) }
	)

	func TestValidateRS256(t *testing.T) {
	idToken, pk := createRS256JWT(t)
	tests := []struct {
	name string
	keyID string
	n *big.Int
	e int
	nowFunc func() time.Time
	wantErr bool
	}{
	{
	name: "works",
	keyID: keyID,
	n: pk.N,
	e: pk.E,
	nowFunc: beforeExp,
	wantErr: false,
	},
	{
	name: "no matching key",
	keyID: "5678",
	n: pk.N,
	e: pk.E,
	nowFunc: beforeExp,
	wantErr: true,
	},
	{
	name: "sig does not match",
	keyID: keyID,
	n: new(big.Int).SetBytes([]byte("42")),
	e: 42,
	nowFunc: beforeExp,
	wantErr: true,
	},
	{
	name: "token expired",
	keyID: keyID,
	n: pk.N,
	e: pk.E,
	nowFunc: afterExp,
	wantErr: true,
	},
	}

	for _, tt := range tests {
	t.Run(tt.name, func(t *testing.T) {
	client := &http.Client{
	Transport: RoundTripFn(func(req http.Request) http.Response {
	cr := certResponse{
	Keys: []jwk{
	{
	Kid: tt.keyID,
	N: base64.RawURLEncoding.EncodeToString(tt.n.Bytes()),
	E: base64.RawURLEncoding.EncodeToString(new(big.Int).SetInt64(int64(tt.e)).Bytes()),
	},
	},
	}
	b, err := json.Marshal(&cr)
	if err != nil {
	t.Fatalf("unable to marshal response: %v", err)
	}
	return &http.Response{
	StatusCode: 200,
	Body: io.NopCloser(bytes.NewReader(b)),
	Header: make(http.Header),
	}
	}),
	}
	oldNow := now
	defer func() { now = oldNow }()
	now = tt.nowFunc

	v, err := NewValidator(&ValidatorOptions{
	Client: client,
	})
	if err != nil {
	t.Fatalf("NewValidator(...) = %q, want nil", err)
	}
	payload, err := v.Validate(context.Background(), idToken, testAudience)
	if tt.wantErr && err != nil {
	// Got the error we wanted.
	return
	}
	if !tt.wantErr && err != nil {
	t.Fatalf("Validate(ctx, %s, %s): got err %q, want nil", idToken, testAudience, err)
	}
	if tt.wantErr && err == nil {
	t.Fatalf("Validate(ctx, %s, %s): got nil err, want err", idToken, testAudience)
	}
	if payload == nil {
	t.Fatalf("Got nil payload, err: %v", err)
	}
	if payload.Audience != testAudience {
	t.Fatalf("Validate(ctx, %s, %s): got %v, want %v", idToken, testAudience, payload.Audience, testAudience)
	}
	if len(payload.Claims) == 0 {
	t.Fatalf("Validate(ctx, %s, %s): missing Claims map. payload.Claims = %+v", idToken, testAudience, payload.Claims)
	}
	if got, ok := payload.Claims["aud"]; !ok {
	t.Fatalf("Validate(ctx, %s, %s): missing aud claim. payload.Claims = %+v", idToken, testAudience, payload.Claims)
	} else {
	got, ok := got.(string)
	if !ok {
	t.Fatalf("Validate(ctx, %s, %s): aud wasn't a string. payload.Claims = %+v", idToken, testAudience, payload.Claims)
	}
	if got != testAudience {
	t.Fatalf("Validate(ctx, %s, %s): Payload[aud] want %v got %v", idToken, testAudience, testAudience, got)
	}
	}
	})
	}
	}

	func TestValidateES256(t *testing.T) {
	idToken, pk := createES256JWT(t)
	tests := []struct {
	name string
	keyID string
	x *big.Int
	y *big.Int
	nowFunc func() time.Time
	wantErr bool
	}{
	{
	name: "works",
	keyID: keyID,
	x: pk.X,
	y: pk.Y,
	nowFunc: beforeExp,
	wantErr: false,
	},
	{
	name: "no matching key",
	keyID: "5678",
	x: pk.X,
	y: pk.Y,
	nowFunc: beforeExp,
	wantErr: true,
	},
	{
	name: "sig does not match",
	keyID: keyID,
	x: new(big.Int),
	y: new(big.Int),
	nowFunc: beforeExp,
	wantErr: true,
	},
	{
	name: "token expired",
	keyID: keyID,
	x: pk.X,
	y: pk.Y,
	nowFunc: afterExp,
	wantErr: true,
	},
	}
	for _, tt := range tests {
	t.Run(tt.name, func(t *testing.T) {
	client := &http.Client{
	Transport: RoundTripFn(func(req http.Request) http.Response {
	cr := certResponse{
	Keys: []jwk{
	{
	Kid: tt.keyID,
	X: base64.RawURLEncoding.EncodeToString(tt.x.Bytes()),
	Y: base64.RawURLEncoding.EncodeToString(tt.y.Bytes()),
	},
	},
	}
	b, err := json.Marshal(&cr)
	if err != nil {
	t.Fatalf("unable to marshal response: %v", err)
	}
	return &http.Response{
	StatusCode: 200,
	Body: io.NopCloser(bytes.NewReader(b)),
	Header: make(http.Header),
	}
	}),
	}
	oldNow := now
	defer func() { now = oldNow }()
	now = tt.nowFunc

	v, err := NewValidator(&ValidatorOptions{
	Client: client,
	})
	if err != nil {
	t.Fatalf("NewValidator(...) = %q, want nil", err)
	}
	payload, err := v.Validate(context.Background(), idToken, testAudience)
	if !tt.wantErr && err != nil {
	t.Fatalf("Validate(ctx, %s, %s) = %q, want nil", idToken, testAudience, err)
	}
	if !tt.wantErr && payload.Audience != testAudience {
	t.Fatalf("got %v, want %v", payload.Audience, testAudience)
	}
	})
	}
	}

	func TestParsePayload(t *testing.T) {
	idToken, _ := createRS256JWT(t)
	tests := []struct {
	name string
	token string
	wantPayloadAudience string
	wantErr bool
	}{{
	name: "valid token",
	token: idToken,
	wantPayloadAudience: testAudience,
	}, {
	name: "unparseable token",
	token: "aaa.bbb.ccc",
	wantErr: true,
	}}

	for _, tt := range tests {
	t.Run(tt.name, func(t *testing.T) {
	payload, err := ParsePayload(tt.token)
	gotErr := err != nil
	if gotErr != tt.wantErr {
	t.Errorf("ParsePayload(%q) got error %v, wantErr = %v", tt.token, err, tt.wantErr)
	}
	if tt.wantPayloadAudience != "" {
	if payload == nil \|\| payload.Audience != tt.wantPayloadAudience {
	t.Errorf("ParsePayload(%q) got payload %+v, want payload with audience = %q", tt.token, payload, tt.wantPayloadAudience)
	}
	}
	})
	}
	}

	func createES256JWT(t *testing.T) (string, ecdsa.PublicKey) {
	t.Helper()
	header, claims := commonToken(t, "ES256")
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
	t.Fatalf("unable to generate key: %v", err)
	}
	signedContent := header + "." + claims
	hashed := sha256.Sum256([]byte(signedContent))
	hash := hashed[:]
	r, s, err := ecdsa.Sign(rand.Reader, privateKey, hash)
	if err != nil {
	t.Fatalf("unable to sign content: %v", err)
	}
	rb := r.Bytes()
	lPadded := make([]byte, es256KeySize)
	copy(lPadded[es256KeySize-len(rb):], rb)
	var sig []byte
	sig = append(sig, lPadded...)
	sig = append(sig, s.Bytes()...)
	signature := base64.RawURLEncoding.EncodeToString(sig)
	return fmt.Sprintf("%s.%s.%s", header, claims, signature), privateKey.PublicKey
	}

	func createRS256JWT(t *testing.T) (string, rsa.PublicKey) {
	t.Helper()
	header, claims := commonToken(t, jwt.HeaderAlgRSA256)
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
	t.Fatalf("unable to generate key: %v", err)
	}
	signedContent := header + "." + claims
	hashed := sha256.Sum256([]byte(signedContent))
	hash := hashed[:]
	sig, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hash)
	if err != nil {
	t.Fatalf("unable to sign content: %v", err)
	}
	signature := base64.RawURLEncoding.EncodeToString(sig)
	return fmt.Sprintf("%s.%s.%s", header, claims, signature), privateKey.PublicKey
	}

	// returns header and claims
	func commonToken(t *testing.T, alg string) (string, string) {
	t.Helper()
	header := jwt.Header{
	KeyID: keyID,
	Algorithm: alg,
	Type: jwt.HeaderType,
	}
	payload := Payload{
	Issuer: "example.com",
	Audience: testAudience,
	Expires: expiry,
	}

	hb, err := json.Marshal(&header)
	if err != nil {
	t.Fatalf("unable to marshall header: %v", err)
	}
	pb, err := json.Marshal(&payload)
	if err != nil {
	t.Fatalf("unable to marshall payload: %v", err)
	}
	eb := base64.RawURLEncoding.EncodeToString(hb)
	ep := base64.RawURLEncoding.EncodeToString(pb)
	return eb, ep
	}

	type RoundTripFn func(req http.Request) http.Response

	func (f RoundTripFn) RoundTrip(req http.Request) (http.Response, error) { return f(req), nil }