From 0eff4f45af966d390663d3aeb2264bec39c7757a Mon Sep 17 00:00:00 2001
From: lzeng <lzeng@fortinet.com>
Date: Mon, 21 Oct 2024 14:33:20 -0700
Subject: [PATCH] Improve set-cookie for example plugin

---
 ac/httpac.go                   |  8 ++++++++
 ac/tokenstore.go               |  2 +-
 server/httpserver.go           |  4 ++--
 server/plugins/example/main.go | 24 +++++++++++++-----------
 4 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/ac/httpac.go b/ac/httpac.go
index f929ed2e..efd9328c 100644
--- a/ac/httpac.go
+++ b/ac/httpac.go
@@ -5,6 +5,7 @@ import (
 	"encoding/base64"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"path/filepath"
 	"sync"
@@ -147,6 +148,13 @@ func (ha *HttpAC) initRouter() {
 			return
 		}
 
+		if token, err = url.QueryUnescape(token); err != nil {
+			err = common.ErrUrlPathInvalid
+			log.Error("token error: %v", err)
+			ctx.String(http.StatusOK, "{\"errMsg\": \"token error: %v\"}", err)
+			return
+		}
+
 		req := &common.HttpRefreshRequest{
 			Token: token,
 			SrcIp: ctx.Query("srcip"),
diff --git a/ac/tokenstore.go b/ac/tokenstore.go
index f81e5ca2..40cb40aa 100644
--- a/ac/tokenstore.go
+++ b/ac/tokenstore.go
@@ -37,7 +37,7 @@ func (a *UdpAC) GenerateAccessToken(entry *AccessEntry) string {
 	a.tokenStoreMutex.Lock()
 	defer a.tokenStoreMutex.Unlock()
 
-	entry.ExpireTime = time.Now().Add(time.Duration(entry.OpenTime) * time.Second)
+	entry.ExpireTime = time.Now().Add(time.Duration(entry.OpenTime+5) * time.Second) // keep token for additional 5 seconds in case a request is received late
 	tokenMap, found := a.tokenStore[token[0:1]]
 	if found {
 		tokenMap[token] = entry
diff --git a/server/httpserver.go b/server/httpserver.go
index 9ff63f9a..8b6e4d50 100644
--- a/server/httpserver.go
+++ b/server/httpserver.go
@@ -287,8 +287,8 @@ func corsMiddleware() gin.HandlerFunc {
 		// HTTP headers for CORS
 		c.Writer.Header().Set("Access-Control-Allow-Origin", "*")                   // allow cross-origin resource sharing
 		c.Writer.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS, POST") // methods
-		c.Writer.Header().Set("Access-Control-Expose-Headers", "Content-Type, Content-Length")
-		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Authorization, X-NHP-Ver")
+		c.Writer.Header().Set("Access-Control-Expose-Headers", "Content-Type, Content-Length, Set-Cookie")
+		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Authorization, X-NHP-Ver, Cookie")
 		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
 		c.Writer.Header().Set("Access-Control-Max-Age", "300")
 		// NHP headers
diff --git a/server/plugins/example/main.go b/server/plugins/example/main.go
index 8fba9fbc..19dcbe16 100644
--- a/server/plugins/example/main.go
+++ b/server/plugins/example/main.go
@@ -247,27 +247,29 @@ func authRegular(ctx *gin.Context, req *common.HttpKnockRequest, res *common.Res
 		}
 
 		// set cookies
+		// note that a dot in domain prefix used to make a difference, but now it doesn't (RFC6265).
+		// The cookie will be sent to any subdomain of the specified domain, with or without the leading dot.
 		singleHost := len(ackMsg.ACTokens) == 1
 		for resName, token := range ackMsg.ACTokens {
 			if singleHost {
 				ctx.SetCookie(
-					"nhp-token",      // Name
-					token,            // Value
-					-1,               // MaxAge
-					"/",              // Path
-					res.CookieDomain, // Domain
-					true,             // Secure
-					true)             // HttpOnly
+					"nhp-token",            // Name
+					url.QueryEscape(token), // Value
+					int(res.OpenTime),      // MaxAge - use the knock interval time
+					"/",                    // Path
+					res.CookieDomain,       // Domain
+					true,                   // Secure - if true, this cookie will only be sent on https, not http
+					true)                   // HttpOnly - if true, this cookie will only be sent on http(s)
 			} else {
 				domain := strings.Split(ackMsg.ResourceHost[resName], ":")[0]
 				ctx.SetCookie(
 					"nhp-token"+"/"+resName, // Name
-					token,                   // Value
-					-1,                      // MaxAge
+					url.QueryEscape(token),  // Value
+					int(res.OpenTime),       // MaxAge - use the knock interval time
 					"/",                     // Path
 					domain,                  // Domain
-					true,                    // Secure
-					true)                    // HttpOnly
+					true,                    // Secure - if true, this cookie will only be sent on https, not http
+					true)                    // HttpOnly - if true, this cookie will only be sent on http(s)
 			}
 			log.Info("ctx.SetCookie.")
 		}