reprogl

1
package session
2

3
import (
4
	"net/http"
5
	"time"
6

7
	"xelbot.com/reprogl/container"
8
)
9

10
type ResponseWriter struct {
11
	http.ResponseWriter
12
	sessionData *Store
13
	written     bool
14
}
15

16
func (sw *ResponseWriter) SetSessionData(d *Store) {
17
	sw.sessionData = d
18
}
19

20
func (sw *ResponseWriter) Write(b []byte) (int, error) {
21
	sw.CheckAndWrite()
22

23
	return sw.ResponseWriter.Write(b)
24
}
25

26
func (sw *ResponseWriter) WriteHeader(code int) {
27
	sw.CheckAndWrite()
28

29
	sw.ResponseWriter.WriteHeader(code)
30
}
31

32
func (sw *ResponseWriter) CheckAndWrite() {
33
	if !sw.written {
34
		var secureCookie *SecureCookie
35
		switch sw.sessionData.status {
36
		case Modified:
37
			expiry := time.Now().Add(maxAge)
38

39
			secureCookie = NewSecureCookie(
40
				container.GetConfig().SessionHashKey,
41
				container.GetConfig().SessionBlockKey,
42
			)
43

44
			sw.sessionData.mu.Lock()
45
			sw.sessionData.data.deadline = deadline(expiry)
46
			err := secureCookie.encode(sw.sessionData.data)
47
			sw.sessionData.mu.Unlock()
48

49
			if err != nil {
50
				panic(err)
51
			}
52

53
			writeCookie(sw, secureCookie, expiry)
54
		case Destroyed:
55
			secureCookie = NewSecureCookie(
56
				container.GetConfig().SessionHashKey,
57
				container.GetConfig().SessionBlockKey,
58
			)
59
			writeCookie(sw, secureCookie, time.Time{})
60
		}
61

62
		if len(sw.Header().Values("Set-Cookie")) > 0 {
63
			sw.Header().Set("Cache-Control", `no-cache="Set-Cookie"`)
64
		}
65
	}
66

67
	sw.written = true
68
}
69

70
func writeCookie(w http.ResponseWriter, c CookieInterface, expiry time.Time) {
71
	cookie := &http.Cookie{
72
		Name:     c.Name(),
73
		Value:    c.Value(),
74
		Path:     c.Path(),
75
		Secure:   true,
76
		HttpOnly: true,
77
		SameSite: http.SameSiteDefaultMode,
78
	}
79

80
	if expiry.IsZero() {
81
		cookie.Expires = time.Unix(1, 0)
82
		cookie.MaxAge = -1
83
	} else if c.Persist() {
84
		cookie.Expires = time.Unix(expiry.Unix()+1, 0)
85
		cookie.MaxAge = int(time.Until(expiry).Seconds() + 1)
86
	}
87

88
	http.SetCookie(w, cookie)
89
}
90