@@ -15,11 +15,14 @@ const (
1515 csrfTokenLength = 32
1616 csrfCookieName = "csrf_token"
1717 csrfHeaderName = "X-CSRF-Token"
18+ csrfContextKey = "CSRFToken"
19+ cleanupInterval = 1 * time .Hour
20+ tokenExpiryTime = 12 * time .Hour
1821)
1922
2023var (
2124 tokenStore = make (map [string ]time.Time )
22- tokenStoreMu sync.Mutex
25+ tokenStoreMu sync.RWMutex
2326 timeNow = time .Now
2427 csrfLogger * slog.Logger
2528)
@@ -28,47 +31,54 @@ func initCSRFLogger(logger *slog.Logger) {
2831 csrfLogger = logger
2932}
3033
31- func generateCSRFToken () ( string , error ) {
34+ func generateCSRFToken () string {
3235 b := make ([]byte , csrfTokenLength )
33- _ , err := rand .Read (b )
34- if err != nil {
35- return "" , err
36+ if _ , err := rand .Read (b ); err != nil {
37+ if csrfLogger != nil {
38+ csrfLogger .Error ("Failed to generate CSRF token" , "error" , err )
39+ }
40+ return ""
3641 }
37- return base64 .StdEncoding .EncodeToString (b ), nil
42+ return base64 .StdEncoding .EncodeToString (b )
3843}
3944
40- func setCSRFToken (r * http.Request , w http.ResponseWriter ) (string , error ) {
41- token , err := generateCSRFToken ()
42- if err != nil {
43- return "" , err
45+ func setCSRFToken (w http.ResponseWriter , r * http.Request ) (string , error ) {
46+ token := generateCSRFToken ()
47+ if token == "" {
48+ return "" , errors . New ( "Failed to generate CSRF token" )
4449 }
4550
4651 isSecure := r .TLS != nil || r .Header .Get ("X-Forwarded-Proto" ) == "https"
4752
4853 http .SetCookie (w , & http.Cookie {
4954 Name : csrfCookieName ,
5055 Value : token ,
56+ Path : "/" ,
5157 HttpOnly : true ,
52- Secure : isSecure , // Set to true if using HTTPS
58+ Secure : isSecure ,
5359 SameSite : http .SameSiteStrictMode ,
60+ MaxAge : int (tokenExpiryTime .Seconds ()),
5461 })
5562
56- if csrfLogger != nil {
57- csrfLogger .DebugContext (r .Context (), "set csrf cookie" )
58- }
59-
6063 tokenStoreMu .Lock ()
61- tokenStore [token ] = timeNow ().Add (24 * time . Hour ) // Token expires in 24 hours
64+ tokenStore [token ] = timeNow ().Add (tokenExpiryTime )
6265 tokenStoreMu .Unlock ()
6366
6467 return token , nil
6568}
6669
70+ func GetCSRFToken (r * http.Request ) string {
71+ if token , ok := r .Context ().Value (csrfContextKey ).(string ); ok {
72+ return token
73+ }
74+ return ""
75+ }
76+
6777func validateCSRFToken (r * http.Request ) error {
6878 cookie , err := r .Cookie (csrfCookieName )
6979 if err != nil {
7080 if csrfLogger != nil {
71- csrfLogger .DebugContext ( r . Context (), "csrf error "slog . String ( "message "err . Error ()) )
81+ csrfLogger .Debug ( "CSRF cookie not found ""error "err )
7282 }
7383 return errors .New ("CSRF cookie not found" )
7484 }
@@ -78,59 +88,99 @@ func validateCSRFToken(r *http.Request) error {
7888 token = r .FormValue ("csrf_token" )
7989 if token == "" {
8090 if csrfLogger != nil {
81- csrfLogger .DebugContext ( r . Context (), "csrf token not found in header or form"
91+ csrfLogger .Debug ( "CSRF token not found in header or form"
8292 }
83- return errors .New ("CSRF token not found in header or form " )
93+ return errors .New ("CSRF token not found" )
8494 }
8595 }
8696
8797 if cookie .Value != token {
8898 if csrfLogger != nil {
89- csrfLogger .DebugContext ( r . Context (), "csrf token mismatch"
99+ csrfLogger .Debug ( "CSRF token mismatch"
90100 }
91101 return errors .New ("CSRF token mismatch" )
92102 }
93103
94- tokenStoreMu .Lock ()
95- defer tokenStoreMu .Unlock ()
96-
104+ tokenStoreMu .RLock ()
97105 expiry , exists := tokenStore [token ]
106+ tokenStoreMu .RUnlock ()
107+
98108 if ! exists {
99109 if csrfLogger != nil {
100- csrfLogger .DebugContext ( r . Context (), "csrf token not found in store"
110+ csrfLogger .Debug ( "CSRF token not found in store"
101111 }
102112 return errors .New ("CSRF token not found in store" )
103113 }
104114
105115 if timeNow ().After (expiry ) {
106- delete (tokenStore , token )
107116 if csrfLogger != nil {
108- csrfLogger .DebugContext ( r . Context (), "csrf token expired"
117+ csrfLogger .Debug ( "CSRF token expired"
109118 }
110119 return errors .New ("CSRF token expired" )
111120 }
112121
122+ if csrfLogger != nil {
123+ csrfLogger .Debug ("CSRF validation successful" )
124+ }
113125 return nil
114126}
115127
116128func CSRFMiddleware (next http.HandlerFunc ) http.HandlerFunc {
117129 return func (w http.ResponseWriter , r * http.Request ) {
118- if r .Method == "GET" || r .Method == "HEAD" || r .Method == "OPTIONS" {
119- token , err := setCSRFToken (r , w )
130+ var token string
131+ var err error
132+
133+ if r .Method == http .MethodGet || r .Method == http .MethodHead {
134+ token , err = setCSRFToken (w , r )
120135 if err != nil {
121- http .Error (w , "Failed to set CSRF token" , http .StatusInternalServerError )
136+ http .Error (w , http . StatusText ( http . StatusInternalServerError ) , http .StatusInternalServerError )
122137 return
123138 }
124139 w .Header ().Set (csrfHeaderName , token )
125-
126- ctx := context .WithValue (r .Context (), "CSRFToken" , token )
127- next .ServeHTTP (w , r .WithContext (ctx ))
128140 } else {
129141 if err := validateCSRFToken (r ); err != nil {
130142 http .Error (w , "CSRF validation failed" , http .StatusForbidden )
131143 return
132144 }
145+ token = r .Header .Get (csrfHeaderName )
146+ if token == "" {
147+ token = r .FormValue ("csrf_token" )
148+ }
149+ }
150+
151+ ctx := context .WithValue (r .Context (), csrfContextKey , token )
152+ next .ServeHTTP (w , r .WithContext (ctx ))
153+ }
154+ }
155+
156+ func cleanupExpiredTokens () {
157+ tokenStoreMu .Lock ()
158+ defer tokenStoreMu .Unlock ()
159+ now := timeNow ()
160+ for token , expiry := range tokenStore {
161+ if now .After (expiry ) {
162+ delete (tokenStore , token )
133163 }
134- next .ServeHTTP (w , r )
135164 }
165+
166+ if csrfLogger != nil {
167+ csrfLogger .Debug ("Cleaned up expired CSRF tokens" )
168+ }
169+ }
170+
171+ func startCleanupRoutine (ctx context.Context ) {
172+ ticker := time .NewTicker (cleanupInterval )
173+ defer ticker .Stop ()
174+ for {
175+ select {
176+ case <- ticker .C :
177+ cleanupExpiredTokens ()
178+ case <- ctx .Done ():
179+ return
180+ }
181+ }
182+ }
183+
184+ func init () {
185+ go startCleanupRoutine (context .Background ())
136186}
0 commit comments