使用 Go 构建安全会话管理器

go
type Session struct {
    createdAt      time.Time
    lastActivityAt time.Time
    id             string
    data           map[string]any
}

go
type SessionStore interface {
    read(id string) (*Session, error)
    write(session *Session) error
    destroy(id string) error
    gc(idleExpiration, absoluteExpiration time.Duration) error
}

go
type SessionManager struct {
    store              SessionStore
    idleExpiration     time.Duration
    absoluteExpiration time.Duration
    cookieName         string
}

go
import (
    "crypto/rand"
    "encoding/base64"
    "io"
)
func generateSessionId() string {
    id := make([]byte, 32)
    _, err := io.ReadFull(rand.Reader, id)
    if err != nil {
        panic("failed to generate session id")
    }
    return base64.RawURLEncoding.EncodeToString(id)
}

cOCMc1RV1m_uFZhkllsMLXPsJgFGc2YQ9GbdPy1hqHyA-

go
func newSession() *Session {
    return &Session{
        id:             generateSessionId(),
        data:           make(map[string]any),
        createdAt:      time.Now(),
        lastActivityAt: time.Now(),
    }
}

go
func (s *Session) Get(key string) any {
    return s.data[key]
}
func (s *Session) Put(key string, value any) {
    s.data[key] = value
}
func (s *Session) Delete(key string) {
    delete(s.data, key)
}

go
func NewSessionManager(
    store SessionStore,
    gcInterval,
    idleExpiration,
    absoluteExpiration time.Duration,
    cookieName string) *SessionManager {
    m := &SessionManager{
        store:              store,
        idleExpiration:     idleExpiration,
        absoluteExpiration: absoluteExpiration,
        cookieName:         cookieName,
    }
    go m.gc(gcInterval)
    return m
}

go
func (m *SessionManager) gc(d time.Duration) {
    ticker := time.NewTicker(d)
    for range ticker.C {
        m.store.gc(m.idleExpiration, m.absoluteExpiration)
    }
}

go
func (m *SessionManager) validate(session *Session) bool {
    if time.Since(session.createdAt) > m.absoluteExpiration ||
        time.Since(session.lastActivityAt) > m.idleExpiration {
        
        // 从存储中删除会话
        err := m.store.destroy(session.id)
        if err != nil {
            panic(err)
        }
        return false
    }
    return true
}

go
func (m *SessionManager) start(r *http.Request) (*Session, *http.Request) {
    var session *Session
    // 从 cookie 中读取
    cookie, err := r.Cookie(m.cookieName)
    if err == nil {
        session, err = m.store.read(cookie.Value)
        if err != nil {
            log.Printf("Failed to read session from store: %v", err)
        }
    }
    // 生成新会话
    if session == nil || !m.validate(session) {
        session = newSession()
    }
    // 将会话附加到上下文
    ctx := context.WithValue(r.Context(), sessionContextKey{}, session)
    r = r.WithContext(ctx)
    return session, r
}

go
func (m *SessionManager) save(session *Session) error {
    session.lastActivityAt = time.Now()
    err := m.store.write(session)
    if err != nil {
        return err
    }
    return nil
}

go
func (m *SessionManager) migrate(session *Session) error {
    session.mu.Lock()
    defer session.mu.Unlock()
    err := m.store.destroy(session.id)
    if err != nil {
        return err
    }
    session.id = generateSessionId()
    return nil
}

go
func (m *SessionManager) Handle(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        
    })
}

``go
// 启动会话
session, rws := m.start(r)
// 创建一个新的响应写入器
sw := &sessionResponseWriter{
    ResponseWriter: w,
    sessionManager: m,
    request:        rws,
}
// 添加必要的头信息
w.Header().Add("Vary", "Cookie")
w.Header().Add("Cache-Control", no-cache="Set-Cookie")
// 调用下一个处理器并传递新的响应写入器和新的请求
next.ServeHTTP(sw, rws)
// 保存会话
m.save(session)
// 如果尚未写入，将会话 cookie 写入响应
writeCookieIfNecessary(sw)
`

go
type sessionResponseWriter struct {
    http.ResponseWriter
    sessionManager *SessionManager
    request        *http.Request
    done           bool
}
func (w *sessionResponseWriter) Write(b []byte) (int, error) {
    writeCookieIfNecessary(w)
    return w.ResponseWriter.Write(b)
}
func (w *sessionResponseWriter) WriteHeader(code int) {
    writeCookieIfNecessary(w)
    w.ResponseWriter.WriteHeader(code)
}
func (w *sessionResponseWriter) Unwrap() http.ResponseWriter {
    return w.ResponseWriter
}

go
func writeCookieIfNecessary(w *sessionResponseWriter) {
    if w.done {
        return
    }
    session, ok := r.Context().Value(sessionContextKey{}).(*Session)
    if !ok {
        panic("session not found in request context")
    }
    cookie := &http.Cookie{
        Name:     w.sessionManager.cookieName,
        Value:    session.id,
        Domain:   "mywebsite.com",
        HttpOnly: true,
        Path:     "/",
        Secure:   true,
        SameSite: http.SameSiteLaxMode,
        Expires:  time.Now().Add(w.sessionManager.idleExpiration),
        MaxAge:   int(w.sessionManager.idleExpiration / time.Second),
    }
    http.SetCookie(w.ResponseWriter, cookie)
    w.done = true
}

go
func main() {
    sessionManager := sess.NewSessionManager(
        sess.NewInMemorySessionStore(),
        30 * time.Minute,
        1 * time.Hour,
        12 * time.Hour,
        "session",
    )
}

go
mux.HandleFunc("/projects/switch/some-project-id", func(w http.ResponseWriter, r *http.Request) {
    session := sess.GetSession(r)
    
    session.Put("current_project", "some-project-id")
})
mux.HandleFunc("/project", func(w http.ResponseWriter, r *http.Request) {
    session := sess.GetSession(r)
    
    currentProject := session.Get("current_project")
})