From 571024182dc8effda8c1ce370a1f3f91be456c6f Mon Sep 17 00:00:00 2001
From: Eric Chiang <eric.chiang.m@gmail.com>
Date: Thu, 25 Aug 2016 16:18:09 -0700
Subject: [PATCH] *: set response types supported in discovery based on server
 config

---
 cmd/dex/serve.go   |  2 ++
 server/handlers.go | 40 +++++++++++++++++++++++-----------------
 server/server.go   |  7 ++++++-
 3 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/cmd/dex/serve.go b/cmd/dex/serve.go
index 932f6c53..be27f12b 100644
--- a/cmd/dex/serve.go
+++ b/cmd/dex/serve.go
@@ -92,6 +92,8 @@ func serve(cmd *cobra.Command, args []string) error {
 		Issuer:     c.Issuer,
 		Connectors: connectors,
 		Storage:    s,
+
+		SupportedResponseTypes: c.OAuth2.ResponseTypes,
 	}
 
 	serv, err := server.New(serverConfig)
diff --git a/server/handlers.go b/server/handlers.go
index 7c9bba58..3aa0b368 100644
--- a/server/handlers.go
+++ b/server/handlers.go
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/url"
 	"path"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -72,32 +73,37 @@ type discovery struct {
 	Claims        []string `json:"claims_supported"`
 }
 
-func (s *Server) handleDiscovery(w http.ResponseWriter, r *http.Request) {
-	// TODO(ericchiang): Cache this
+func (s *Server) discoveryHandler() (http.HandlerFunc, error) {
 	d := discovery{
-		Issuer:        s.issuerURL.String(),
-		Auth:          s.absURL("/auth"),
-		Token:         s.absURL("/token"),
-		Keys:          s.absURL("/keys"),
-		ResponseTypes: []string{"code"},
-		Subjects:      []string{"public"},
-		IDTokenAlgs:   []string{string(jose.RS256)},
-		Scopes:        []string{"openid", "email", "profile", "offline_access"},
-		AuthMethods:   []string{"client_secret_basic"},
+		Issuer:      s.issuerURL.String(),
+		Auth:        s.absURL("/auth"),
+		Token:       s.absURL("/token"),
+		Keys:        s.absURL("/keys"),
+		Subjects:    []string{"public"},
+		IDTokenAlgs: []string{string(jose.RS256)},
+		Scopes:      []string{"openid", "email", "profile", "offline_access"},
+		AuthMethods: []string{"client_secret_basic"},
 		Claims: []string{
 			"aud", "email", "email_verified", "exp",
 			"iat", "iss", "locale", "name", "sub",
 		},
 	}
+
+	for responseType := range s.supportedResponseTypes {
+		d.ResponseTypes = append(d.ResponseTypes, responseType)
+	}
+	sort.Strings(d.ResponseTypes)
+
 	data, err := json.MarshalIndent(d, "", "  ")
 	if err != nil {
-		log.Printf("failed to marshal discovery data: %v", err)
-		http.Error(w, "Internal server error", http.StatusInternalServerError)
-		return
+		return nil, fmt.Errorf("failed to marshal discovery data: %v", err)
 	}
-	w.Header().Set("Content-Type", "application/json")
-	w.Header().Set("Content-Length", strconv.Itoa(len(data)))
-	w.Write(data)
+
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.Header().Set("Content-Length", strconv.Itoa(len(data)))
+		w.Write(data)
+	}, nil
 }
 
 // handleAuthorization handles the OAuth2 auth endpoint.
diff --git a/server/server.go b/server/server.go
index bb488f54..ef2697db 100644
--- a/server/server.go
+++ b/server/server.go
@@ -136,8 +136,13 @@ func newServer(c Config, rotationStrategy rotationStrategy) (*Server, error) {
 	}
 	r.NotFoundHandler = http.HandlerFunc(s.notFound)
 
+	discoveryHandler, err := s.discoveryHandler()
+	if err != nil {
+		return nil, err
+	}
+	handleFunc("/.well-known/openid-configuration", discoveryHandler)
+
 	// TODO(ericchiang): rate limit certain paths based on IP.
-	handleFunc("/.well-known/openid-configuration", s.handleDiscovery)
 	handleFunc("/token", s.handleToken)
 	handleFunc("/keys", s.handlePublicKeys)
 	handleFunc("/auth", s.handleAuthorization)