From a7978890c7bfcba41651b3ea4e05514e7fa69889 Mon Sep 17 00:00:00 2001
From: "m.nabokikh" <maksim.nabokikh@flant.com>
Date: Mon, 18 Jan 2021 11:10:00 +0400
Subject: [PATCH] Add Cache-control headers to token responses

Signed-off-by: m.nabokikh <maksim.nabokikh@flant.com>
---
 server/deviceflowhandlers.go |  4 ++++
 server/handlers.go           |  4 ++++
 server/server_test.go        | 16 ++++++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/server/deviceflowhandlers.go b/server/deviceflowhandlers.go
index e0932c54..20b0a19b 100644
--- a/server/deviceflowhandlers.go
+++ b/server/deviceflowhandlers.go
@@ -140,6 +140,10 @@ func (s *Server) handleDeviceCode(w http.ResponseWriter, r *http.Request) {
 			PollInterval:            pollIntervalSeconds,
 		}
 
+		// Device Authorization Response can contain cache control header according to
+		// https://tools.ietf.org/html/rfc8628#section-3.2
+		w.Header().Set("Cache-Control", "no-store")
+
 		enc := json.NewEncoder(w)
 		enc.SetEscapeHTML(false)
 		enc.SetIndent("", "   ")
diff --git a/server/handlers.go b/server/handlers.go
index ec056a76..be9cf4e8 100644
--- a/server/handlers.go
+++ b/server/handlers.go
@@ -1476,6 +1476,10 @@ func (s *Server) writeAccessToken(w http.ResponseWriter, resp *accessTokenRespon
 	}
 	w.Header().Set("Content-Type", "application/json")
 	w.Header().Set("Content-Length", strconv.Itoa(len(data)))
+
+	// Token response must include cache headers https://tools.ietf.org/html/rfc6749#section-5.1
+	w.Header().Set("Cache-Control", "no-store")
+	w.Header().Set("Pragma", "no-cache")
 	w.Write(data)
 }
 
diff --git a/server/server_test.go b/server/server_test.go
index fa73743a..31753ff6 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -395,6 +395,12 @@ func makeOAuth2Tests(clientID string, clientSecret string, now func() time.Time)
 						}
 						return fmt.Errorf("unexpected response: %s", dump)
 					}
+					if resp.Header.Get("Cache-Control") != "no-store" {
+						return fmt.Errorf("cache-control header doesn't included in token response")
+					}
+					if resp.Header.Get("Pragma") != "no-cache" {
+						return fmt.Errorf("pragma header doesn't included in token response")
+					}
 					return nil
 				},
 			},
@@ -423,6 +429,12 @@ func makeOAuth2Tests(clientID string, clientSecret string, now func() time.Time)
 						}
 						return fmt.Errorf("unexpected response: %s", dump)
 					}
+					if resp.Header.Get("Cache-Control") != "no-store" {
+						return fmt.Errorf("cache-control header doesn't included in token response")
+					}
+					if resp.Header.Get("Pragma") != "no-cache" {
+						return fmt.Errorf("pragma header doesn't included in token response")
+					}
 					return nil
 				},
 			},
@@ -701,6 +713,7 @@ func TestOAuth2CodeFlow(t *testing.T) {
 						checkErrorResponse(err, t, tc)
 						return
 					}
+
 					if err != nil {
 						t.Errorf("failed to exchange code for token: %v", err)
 						return
@@ -1515,6 +1528,9 @@ func TestOAuth2DeviceFlow(t *testing.T) {
 			if resp.StatusCode != http.StatusOK {
 				t.Errorf("%v - Unexpected Response Type.  Expected 200 got  %v.  Response: %v", tc.name, resp.StatusCode, string(responseBody))
 			}
+			if resp.Header.Get("Cache-Control") != "no-store" {
+				t.Errorf("Cache-Control header doesn't exist in Device Code Response")
+			}
 
 			// Parse the code response
 			var deviceCode deviceCodeResponse