Go语言Web开发：构建灵活的Per-Handler中间件并安全传递请求数据

func checkCSRF(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        // 执行CSRF检查逻辑
        // 如果检查失败，返回HTTP 403
        // 否则，调用下一个处理器
        next.ServeHTTP(w, r)
    }
}

type CSRFHandlerFunc func(w http.ResponseWriter, r *http.Request, token string)

func csrfCheckMiddleware(next CSRFHandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        // 模拟生成CSRF token
        token := "generated_csrf_token_xyz"
        // 调用下一个处理器，并传递token
        next(w, r, token)
    }
}

package main

import (
    "context"
    "fmt"
    "net/http"
    "log"
)

// 定义一个自定义的context key类型，避免与其他包的key冲突
type contextKey string

const csrfTokenKey contextKey = "csrfToken"

// checkCSRFMiddleware 是一个中间件，用于检查CSRF并设置token到context
func checkCSRFMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 模拟生成或验证CSRF token
        // 在实际应用中，这里会包含从session获取、验证或生成token的逻辑
        token := "generated_csrf_token_123"

        // 将token放入请求的context中
        // context.WithValue 返回一个新的context，其中包含给定的键值对
        ctx := context.WithValue(r.Context(), csrfTokenKey, token)

        // 使用新的context更新请求
        r = r.WithContext(ctx)

        // 调用下一个处理器
        next.ServeHTTP(w, r)
    })
}

// checkExistingDataMiddleware 是另一个中间件，模拟检查现有数据
func checkExistingDataMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 模拟检查是否存在某些数据
        hasData := true // 假设检查结果为true

        if !hasData {
            http.Redirect(w, r, "/no-data", http.StatusFound)
            return
        }
        // 如果需要，也可以将hasData放入context
        // ctx := context.WithValue(r.Context(), contextKey("hasData"), hasData)
        // r = r.WithContext(ctx)

        next.ServeHTTP(w, r)
    })
}

// previewHandler 是最终的处理函数，它从context中获取CSRF token
func previewHandler(w http.ResponseWriter, r *http.Request) {
    // 从context中获取CSRF token
    token, ok := r.Context().Value(csrfTokenKey).(string)
    if !ok {
        http.Error(w, "CSRF token not found or invalid in context", http.StatusInternalServerError)
        return
    }
    fmt.Fprintf(w, "欢迎来到预览页面！您的CSRF Token: %s\n", token)
    // 在这里可以使用token渲染模板或进行其他操作
}

func noDataHandler(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "没有数据可供预览！\n")
}

func main() {
    mux := http.NewServeMux()

    // 链式调用中间件和处理器
    // 中间件的调用顺序是从外到内，执行顺序是从外到内再到外
    // checkCSRFMiddleware -> checkExistingDataMiddleware -> previewHandler
    mux.Handle("/preview", checkCSRFMiddleware(checkExistingDataMiddleware(http.HandlerFunc(previewHandler))))
    mux.HandleFunc("/no-data", noDataHandler)

    log.Println("服务器正在监听 :8080")
    log.Fatal(http.ListenAndServe(":8080", mux))
}

package main

import (
    "net/http"
    "sync"
    "fmt"
    "log"
)

// requestContext 是一个用于存储请求特定数据的映射
var requestContext = struct {
    sync.RWMutex
    data map[*http.Request]map[interface{}]interface{}
}{
    data: make(map[*http.Request]map[interface{}]interface{}),
}

// SetContextValue 将值与请求关联
func SetContextValue(r *http.Request, key, value interface{}) {
    requestContext.Lock()
    defer requestContext.Unlock()
    if requestContext.data[r] == nil {
        requestContext.data[r] = make(map[interface{}]interface{})
    }
    requestContext.data[r][key] = value
}

// GetContextValue 从请求中获取关联的值
func GetContextValue(r *http.Request, key interface{}) interface{} {
    requestContext.RLock()
    defer requestContext.RUnlock()
    if rData, ok := requestContext.data[r]; ok {
        return rData[key]
    }
    return nil
}

// ClearContext 清理与请求关联的所有数据
// 这一步至关重要，必须在请求处理结束后调用
func ClearContext(r *http.Request) {
    requestContext.Lock()
    defer requestContext.Unlock()
    delete(requestContext.data, r)
}

// customContextMiddleware 是一个使用自定义上下文的中间件
func customContextMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        defer ClearContext(r) // 确保在请求结束时清理上下文

        SetContextValue(r, "myCustomData", "Hello from custom context!")
        next.ServeHTTP(w, r)
    })
}

// customHandler 使用自定义上下文数据
func customHandler(w http.ResponseWriter, r *http.Request) {
    data := GetContextValue(r, "myCustomData")
    if data != nil {
        fmt.Fprintf(w, "Custom data: %s\n", data)
    } else {
        fmt.Fprintf(w, "No custom data found.\n")
    }
}

func main() {
    mux := http.NewServeMux()
    mux.Handle("/custom", customContextMiddleware(http.HandlerFunc(customHandler)))
    log.Println("Custom context server listening on :8081")
    log.Fatal(http.ListenAndServe(":8081", mux))
}