diff --git a/contrib/token-server/main.go b/contrib/token-server/main.go
index 138793c7..8f9029ea 100644
--- a/contrib/token-server/main.go
+++ b/contrib/token-server/main.go
@@ -245,7 +245,7 @@ func (ts *tokenServer) getToken(ctx context.Context, w http.ResponseWriter, r *h
 		// Get response context.
 		ctx, w = dcontext.WithResponseWriter(ctx, w)
 
-		challenge.SetHeaders(w)
+		challenge.SetHeaders(r, w)
 		handleError(ctx, errcode.ErrorCodeUnauthorized.WithDetail(challenge.Error()), w)
 
 		dcontext.GetResponseLogger(ctx).Info("get token authentication challenge")
diff --git a/registry/auth/htpasswd/access_test.go b/registry/auth/htpasswd/access_test.go
index 7a3d411e..0bfc427e 100644
--- a/registry/auth/htpasswd/access_test.go
+++ b/registry/auth/htpasswd/access_test.go
@@ -50,7 +50,7 @@ func TestBasicAccessController(t *testing.T) {
 		if err != nil {
 			switch err := err.(type) {
 			case auth.Challenge:
-				err.SetHeaders(w)
+				err.SetHeaders(r, w)
 				w.WriteHeader(http.StatusUnauthorized)
 				return
 			default:
diff --git a/registry/auth/silly/access_test.go b/registry/auth/silly/access_test.go
index 0a5103e6..19824494 100644
--- a/registry/auth/silly/access_test.go
+++ b/registry/auth/silly/access_test.go
@@ -21,7 +21,7 @@ func TestSillyAccessController(t *testing.T) {
 		if err != nil {
 			switch err := err.(type) {
 			case auth.Challenge:
-				err.SetHeaders(w)
+				err.SetHeaders(r, w)
 				w.WriteHeader(http.StatusUnauthorized)
 				return
 			default: