From 40c6d73ba0d839c636fbb0a713a20dcebb3dbc17 Mon Sep 17 00:00:00 2001
From: arekkas <aeneas@ory.am>
Date: Thu, 14 Dec 2017 21:52:27 +0100
Subject: [PATCH] Replaces key discovery with well-known feature

Closes #43
---
 director/director.go    |  7 +++++--
 rsakey/handler.go       | 18 +++++++++---------
 rsakey/manager.go       |  1 +
 rsakey/manager_hydra.go |  4 ++++
 rsakey/manager_local.go |  4 ++++
 5 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/director/director.go b/director/director.go
index 6705a2e33d..6a73f8ac87 100644
--- a/director/director.go
+++ b/director/director.go
@@ -113,7 +113,10 @@ func (d *Director) Director(r *http.Request) {
 		return
 	}
 
-	token, err := jwt.NewWithClaims(jwt.SigningMethodRS256, access.ToClaims()).SignedString(privateKey)
+	token := jwt.NewWithClaims(jwt.SigningMethodRS256, access.ToClaims())
+	token.Header["kid"] = d.KeyManager.PublicKeyID()
+
+	signed, err := token.SignedString(privateKey)
 	if err != nil {
 		d.Logger.
 			WithError(errors.WithStack(err)).
@@ -125,5 +128,5 @@ func (d *Director) Director(r *http.Request) {
 
 	r.URL.Scheme = d.TargetURL.Scheme
 	r.URL.Host = d.TargetURL.Host
-	*r = *r.WithContext(context.WithValue(r.Context(), requestAllowed, token))
+	*r = *r.WithContext(context.WithValue(r.Context(), requestAllowed, signed))
 }
diff --git a/rsakey/handler.go b/rsakey/handler.go
index 7f1adf415b..aac94c7635 100644
--- a/rsakey/handler.go
+++ b/rsakey/handler.go
@@ -14,21 +14,21 @@ type Handler struct {
 }
 
 func (h *Handler) SetRoutes(r *httprouter.Router) {
-	r.GET("/keys/id-token.public", h.GetPublicKey)
+	r.GET("/.well-known/jwks.json", h.WellKnown)
 }
 
-func (h *Handler) GetPublicKey(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+func (h *Handler) WellKnown(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
 	key, err := h.M.PublicKey()
 	if err != nil {
 		h.H.WriteError(w, r, err)
 		return
 	}
 
-	jwk := &jose.JSONWebKey{
-		Key:       key,
-		KeyID:     "id-token.public",
-		Algorithm: h.M.Algorithm(),
-	}
-
-	h.H.Write(w, r, jwk)
+	h.H.Write(w, r, &jose.JSONWebKeySet{
+		Keys: []jose.JSONWebKey{{
+			Key:       key,
+			KeyID:     h.M.PublicKeyID(),
+			Algorithm: h.M.Algorithm(),
+		}},
+	})
 }
diff --git a/rsakey/manager.go b/rsakey/manager.go
index 0c8a2300b5..1663087957 100644
--- a/rsakey/manager.go
+++ b/rsakey/manager.go
@@ -6,5 +6,6 @@ type Manager interface {
 	Refresh() error
 	PrivateKey() (*rsa.PrivateKey, error)
 	PublicKey() (*rsa.PublicKey, error)
+	PublicKeyID() string
 	Algorithm() string
 }
diff --git a/rsakey/manager_hydra.go b/rsakey/manager_hydra.go
index 354a6f9d5c..2a958401d9 100644
--- a/rsakey/manager_hydra.go
+++ b/rsakey/manager_hydra.go
@@ -76,6 +76,10 @@ func (m *HydraManager) PrivateKey() (*rsa.PrivateKey, error) {
 	return m.key, nil
 }
 
+func (m *HydraManager) PublicKeyID() string {
+	return m.Set + ":public"
+}
+
 func (m *HydraManager) Algorithm() string {
 	return "RS256"
 }
diff --git a/rsakey/manager_local.go b/rsakey/manager_local.go
index 347de7473a..316ab6103a 100644
--- a/rsakey/manager_local.go
+++ b/rsakey/manager_local.go
@@ -44,6 +44,10 @@ func (m *LocalManager) PrivateKey() (*rsa.PrivateKey, error) {
 	return m.key, nil
 }
 
+func (m *LocalManager) PublicKeyID() string {
+	return "id-token:public"
+}
+
 func (m *LocalManager) Algorithm() string {
 	return "RS256"
 }