google: support key ID in JWTAccessTokenSourceFromJSON

Change-Id: I20ffede5bf81aa4990afb2820561d5633cdb43a8
Reviewed-on: https://go-review.googlesource.com/24440
Reviewed-by: Chris Broadfoot <cbro@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

google/google.go

@@ -88,16 +88,19 @@ func JWTConfigFromJSON(jsonKey []byte, scope ...string) (*jwt.Config, error) {
 	var key struct {
 		Email        string `json:"client_email"`
 		PrivateKey   string `json:"private_key"`
+		PrivateKeyID string `json:"private_key_id"`
 	}
 	if err := json.Unmarshal(jsonKey, &key); err != nil {
 		return nil, err
 	}
-	return &jwt.Config{
+	config := &jwt.Config{
 		Email:        key.Email,
 		PrivateKey:   []byte(key.PrivateKey),
+		PrivateKeyID: key.PrivateKeyID,
 		Scopes:       scope,
 		TokenURL:     JWTTokenURL,
-	}, nil
+	}
+	return config, nil
 }
 
 // ComputeTokenSource returns a token source that fetches access tokens

google/google_test.go

@@ -31,6 +31,14 @@ var installedJSONKey = []byte(`{
     }
 }`)
 
+var jwtJSONKey = []byte(`{
+  "private_key_id": "268f54e43a1af97cfc71731688434f45aca15c8b",
+  "private_key": "super secret key",
+  "client_email": "gopher@developer.gserviceaccount.com",
+  "client_id": "gopher.apps.googleusercontent.com",
+  "type": "service_account"
+}`)
+
 func TestConfigFromJSON(t *testing.T) {
 	conf, err := ConfigFromJSON(webJSONKey, "scope1", "scope2")
 	if err != nil {
@@ -65,3 +73,25 @@ func TestConfigFromJSON_Installed(t *testing.T) {
 		t.Errorf("ClientID = %q; want %q", got, want)
 	}
 }
+
+func TestJWTConfigFromJSON(t *testing.T) {
+	conf, err := JWTConfigFromJSON(jwtJSONKey, "scope1", "scope2")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := conf.Email, "gopher@developer.gserviceaccount.com"; got != want {
+		t.Errorf("Email = %q, want %q", got, want)
+	}
+	if got, want := string(conf.PrivateKey), "super secret key"; got != want {
+		t.Errorf("PrivateKey = %q, want %q", got, want)
+	}
+	if got, want := conf.PrivateKeyID, "268f54e43a1af97cfc71731688434f45aca15c8b"; got != want {
+		t.Errorf("PrivateKeyID = %q, want %q", got, want)
+	}
+	if got, want := strings.Join(conf.Scopes, ","), "scope1,scope2"; got != want {
+		t.Errorf("Scopes = %q; want %q", got, want)
+	}
+	if got, want := conf.TokenURL, "https://accounts.google.com/o/oauth2/token"; got != want {
+		t.Errorf("TokenURL = %q; want %q", got, want)
+	}
+}

google/jwt.go

@@ -36,6 +36,7 @@ func JWTAccessTokenSourceFromJSON(jsonKey []byte, audience string) (oauth2.Token
 		email:    cfg.Email,
 		audience: audience,
 		pk:       pk,
+		pkID:     cfg.PrivateKeyID,
 	}
 	tok, err := ts.Token()
 	if err != nil {
@@ -47,6 +48,7 @@ func JWTAccessTokenSourceFromJSON(jsonKey []byte, audience string) (oauth2.Token
 type jwtAccessTokenSource struct {
 	email, audience string
 	pk              *rsa.PrivateKey
+	pkID            string
 }
 
 func (ts *jwtAccessTokenSource) Token() (*oauth2.Token, error) {
@@ -62,6 +64,7 @@ func (ts *jwtAccessTokenSource) Token() (*oauth2.Token, error) {
 	hdr := &jws.Header{
 		Algorithm: "RS256",
 		Typ:       "JWT",
+		KeyID:     string(ts.pkID),
 	}
 	msg, err := jws.Encode(hdr, cs, ts.pk)
 	if err != nil {

google/jwt_test.go

@@ -0,0 +1,92 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package google
+
+import (
+	"bytes"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/json"
+	"encoding/pem"
+	"strings"
+	"testing"
+	"time"
+
+	"golang.org/x/oauth2/jws"
+)
+
+func TestJWTAccessTokenSourceFromJSON(t *testing.T) {
+	// Generate a key we can use in the test data.
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Encode the key and substitute into our example JSON.
+	enc := pem.EncodeToMemory(&pem.Block{
+		Type:  "PRIVATE KEY",
+		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
+	})
+	enc, err = json.Marshal(string(enc))
+	if err != nil {
+		t.Fatalf("json.Marshal: %v", err)
+	}
+	jsonKey := bytes.Replace(jwtJSONKey, []byte(`"super secret key"`), enc, 1)
+
+	ts, err := JWTAccessTokenSourceFromJSON(jsonKey, "audience")
+	if err != nil {
+		t.Fatalf("JWTAccessTokenSourceFromJSON: %v\nJSON: %s", err, string(jsonKey))
+	}
+
+	tok, err := ts.Token()
+	if err != nil {
+		t.Fatalf("Token: %v", err)
+	}
+
+	if got, want := tok.TokenType, "Bearer"; got != want {
+		t.Errorf("TokenType = %q, want %q", got, want)
+	}
+	if got := tok.Expiry; tok.Expiry.Before(time.Now()) {
+		t.Errorf("Expiry = %v, should not be expired", got)
+	}
+
+	err = jws.Verify(tok.AccessToken, &privateKey.PublicKey)
+	if err != nil {
+		t.Errorf("jws.Verify on AccessToken: %v", err)
+	}
+
+	claim, err := jws.Decode(tok.AccessToken)
+	if err != nil {
+		t.Fatalf("jws.Decode on AccessToken: %v", err)
+	}
+
+	if got, want := claim.Iss, "gopher@developer.gserviceaccount.com"; got != want {
+		t.Errorf("Iss = %q, want %q", got, want)
+	}
+	if got, want := claim.Sub, "gopher@developer.gserviceaccount.com"; got != want {
+		t.Errorf("Sub = %q, want %q", got, want)
+	}
+	if got, want := claim.Aud, "audience"; got != want {
+		t.Errorf("Aud = %q, want %q", got, want)
+	}
+
+	// Finally, check the header private key.
+	parts := strings.Split(tok.AccessToken, ".")
+	parts[0] += strings.Repeat("=", len(parts[0])%4) // Add padding.
+	hdrJSON, err := base64.URLEncoding.DecodeString(parts[0])
+	if err != nil {
+		t.Fatalf("base64 DecodeString: %v\nString: %q", err, parts[0])
+	}
+	var hdr jws.Header
+	if err := json.Unmarshal([]byte(hdrJSON), &hdr); err != nil {
+		t.Fatalf("json.Unmarshal: %v (%q)", err)
+	}
+
+	if got, want := hdr.KeyID, "268f54e43a1af97cfc71731688434f45aca15c8b"; got != want {
+		t.Errorf("Header KeyID = %q, want %q", got, want)
+	}
+}

jws/jws.go

@@ -92,6 +92,9 @@ type Header struct {
 
 	// Represents the token type.
 	Typ string `json:"typ"`
+
+	// The optional hint of which key is being used.
+	KeyID string `json:"kid,omitempty"`
 }
 
 func (h *Header) encode() (string, error) {

jwt/jwt.go

@@ -46,6 +46,10 @@ type Config struct {
 	//
 	PrivateKey []byte
 
+	// PrivateKeyID contains an optional hint indicating which key is being
+	// used.
+	PrivateKeyID string
+
 	// Subject is the optional user to impersonate.
 	Subject string