From 88af2c73e51f74a3d199eb5572d9df90f4ad47a1 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Mon, 19 Feb 2024 10:26:17 +0100
Subject: [PATCH] Set secure cookie if reverse proxy indicates HTTPS

---
 middleware.go | 21 +++++++++++++++++----
 user.go       |  4 ++--

diff --git a/middleware.go b/middleware.go
index 909b5293f02b070318aef33b2cbb5e2924d2cd15..099af5d5dd5b9bf35f8d3e8059147f758d546ea9 100644
--- a/middleware.go
+++ b/middleware.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"html/template"
+	"mime"
 	"net/http"
 )
 
@@ -40,25 +41,37 @@ 	ctx = context.WithValue(ctx, contextKeyTemplate, tpl)
 	return ctx
 }
 
-func setLoginTokenCookie(w http.ResponseWriter, token *AccessToken, secret string) {
+func setLoginTokenCookie(w http.ResponseWriter, req *http.Request, token *AccessToken, secret string) {
 	http.SetCookie(w, &http.Cookie{
 		Name:     "sinwon-token",
 		Value:    MarshalSecret(token.ID, secret),
 		HttpOnly: true,
 		SameSite: http.SameSiteStrictMode,
-		// TODO: Secure
+		Secure:   isForwardedHTTPS(req),
 	})
 }
 
-func unsetLoginTokenCookie(w http.ResponseWriter) {
+func unsetLoginTokenCookie(w http.ResponseWriter, req *http.Request) {
 	http.SetCookie(w, &http.Cookie{
 		Name:     "sinwon-token",
 		HttpOnly: true,
 		SameSite: http.SameSiteStrictMode,
+		Secure:   isForwardedHTTPS(req),
 		MaxAge:   -1,
 	})
 }
 
+func isForwardedHTTPS(req *http.Request) bool {
+	if forwarded := req.Header.Get("Forwarded"); forwarded != "" {
+		_, params, _ := mime.ParseMediaType("_; " + forwarded)
+		return params["proto"] == "https"
+	}
+	if forwardedProto := req.Header.Get("X-Forwarded-Proto"); forwardedProto != "" {
+		return forwardedProto == "https"
+	}
+	return false
+}
+
 func loginTokenMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		cookie, _ := req.Cookie("sinwon-token")
@@ -72,7 +85,7 @@ 		db := dbFromContext(ctx)
 		tokenID, tokenSecret, _ := UnmarshalSecret[*AccessToken](cookie.Value)
 		token, err := db.FetchAccessToken(ctx, tokenID)
 		if err == errNoDBRows || (err == nil && !token.VerifySecret(tokenSecret)) {
-			unsetLoginTokenCookie(w)
+			unsetLoginTokenCookie(w, req)
 			next.ServeHTTP(w, req)
 			return
 		} else if err != nil {
diff --git a/user.go b/user.go
index a0b3dd05253ba208c19f2c686c5981a30992f945..7ee8bc93dbed941e4cc64c79c1f40c324ff70fd5 100644
--- a/user.go
+++ b/user.go
@@ -100,12 +100,12 @@ 		httpError(w, fmt.Errorf("failed to create access token: %v", err))
 		return
 	}
 
-	setLoginTokenCookie(w, &token, secret)
+	setLoginTokenCookie(w, req, &token, secret)
 	http.Redirect(w, req, redirectURI.String(), http.StatusFound)
 }
 
 func logout(w http.ResponseWriter, req *http.Request) {
-	unsetLoginTokenCookie(w)
+	unsetLoginTokenCookie(w, req)
 	http.Redirect(w, req, "/login", http.StatusFound)
 }
 

-- 
2.48.1