diff --git a/max_age.go b/max_age.go index 3032285..603a0b5 100644 --- a/max_age.go +++ b/max_age.go @@ -5,6 +5,9 @@ package traefik_session_max_age import ( "context" "net/http" + "net/textproto" + "strconv" + "strings" ) // Config the plugin configuration. @@ -29,13 +32,30 @@ func (rww responseWriterWrapper) Header() http.Header { } func (rww responseWriterWrapper) WriteHeader(code int) { - if rww.cookieName != "" { - res := http.Response{Header: rww.ResponseWriter.Header()} - cookies := res.Cookies() - for _, cookie := range cookies { - if cookie.Name == rww.cookieName { - cookie.MaxAge = rww.maxAge - http.SetCookie(rww.ResponseWriter, cookie) + cookies := rww.ResponseWriter.Header()["Set-Cookie"] + if rww.cookieName != "" && len(cookies) > 0 { + for i, line := range cookies { + parts := strings.Split(textproto.TrimString(line), ";") + if len(parts) == 1 && parts[0] == "" { + continue + } + parts[0] = textproto.TrimString(parts[0]) + name, _, ok := strings.Cut(parts[0], "=") + if !ok { + continue + } + buf := make([]byte, 0, 19) + name = textproto.TrimString(name) + if name == rww.cookieName { + var b strings.Builder + b.WriteString(line) + if rww.maxAge > 0 { + b.WriteString("; Max-Age=") + b.Write(strconv.AppendInt(buf, int64(rww.maxAge), 10)) + } else if rww.maxAge < 0 { + b.WriteString("; Max-Age=0") + } + cookies[i] = b.String() } } }