session_binding_test.go 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. //go:build unit
  2. package middleware
  3. import (
  4. "net/http/httptest"
  5. "strings"
  6. "testing"
  7. "github.com/Wei-Shaw/sub2api/internal/config"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
  9. "github.com/Wei-Shaw/sub2api/internal/service"
  10. "github.com/gin-gonic/gin"
  11. "github.com/stretchr/testify/require"
  12. )
  13. func TestSessionBindingContextFollowsForwardedIPSwitch(t *testing.T) {
  14. gin.SetMode(gin.TestMode)
  15. for _, tc := range []struct {
  16. name string
  17. trustForwarded bool
  18. trustedProxies []string
  19. wantIP string
  20. }{
  21. {name: "enabled switch takes over raw headers", trustForwarded: true, wantIP: "1.2.3.4"},
  22. {name: "disabled switch ignores untrusted headers", trustForwarded: false, wantIP: "127.0.0.1"},
  23. {name: "disabled switch uses configured Gin proxy", trustForwarded: false, trustedProxies: []string{"127.0.0.1"}, wantIP: "1.2.3.4"},
  24. } {
  25. t.Run(tc.name, func(t *testing.T) {
  26. cfg := &config.Config{}
  27. cfg.SetTrustForwardedIPForAPIKeyACL(tc.trustForwarded)
  28. r := gin.New()
  29. require.NoError(t, r.SetTrustedProxies(tc.trustedProxies))
  30. r.Use(SessionBindingContext(cfg))
  31. r.GET("/t", func(c *gin.Context) {
  32. binding := service.SessionBindingFromContext(c.Request.Context())
  33. require.NotNil(t, binding)
  34. require.Equal(t, tc.wantIP, binding.IP)
  35. require.Equal(t, "test-agent", binding.UserAgent)
  36. require.Equal(t, tc.wantIP, SecurityClientIP(c))
  37. c.Status(200)
  38. })
  39. w := httptest.NewRecorder()
  40. req := httptest.NewRequest("GET", "/t", nil)
  41. req.RemoteAddr = "127.0.0.1:54321"
  42. req.Header.Set("X-Real-IP", "1.2.3.4")
  43. req.Header.Set("User-Agent", "test-agent")
  44. r.ServeHTTP(w, req)
  45. require.Equal(t, 200, w.Code)
  46. })
  47. }
  48. }
  49. func TestSessionBindingContextSnapshotsForwardedModeAndHeaders(t *testing.T) {
  50. gin.SetMode(gin.TestMode)
  51. cfg := &config.Config{}
  52. cfg.SetForwardedClientIPSettings(true, []string{"X-Initial-IP"})
  53. r := gin.New()
  54. require.NoError(t, r.SetTrustedProxies(nil))
  55. r.Use(SessionBindingContext(cfg))
  56. r.GET("/t", func(c *gin.Context) {
  57. binding := service.SessionBindingFromContext(c.Request.Context())
  58. require.NotNil(t, binding)
  59. require.Equal(t, "1.2.3.4", binding.IP)
  60. cfg.SetForwardedClientIPSettings(false, []string{"X-Changed-IP"})
  61. require.Equal(t, "1.2.3.4", ip.GetSecurityClientIP(c, false))
  62. c.Status(200)
  63. })
  64. w := httptest.NewRecorder()
  65. req := httptest.NewRequest("GET", "/t", nil)
  66. req.RemoteAddr = "9.9.9.9:12345"
  67. req.Header.Set("X-Initial-IP", "1.2.3.4")
  68. req.Header.Set("X-Changed-IP", "4.4.4.4")
  69. req.Header.Set("X-Real-IP", "8.8.8.8")
  70. r.ServeHTTP(w, req)
  71. require.Equal(t, 200, w.Code)
  72. runtimeSettings := cfg.ForwardedClientIPSettings()
  73. require.False(t, runtimeSettings.TrustForwardedIP)
  74. require.Equal(t, []string{"X-Changed-IP"}, runtimeSettings.Headers)
  75. }
  76. func TestSessionBindingContextBoundsPersistedUserAgent(t *testing.T) {
  77. cfg := &config.Config{}
  78. r := gin.New()
  79. r.Use(SessionBindingContext(cfg))
  80. r.GET("/t", func(c *gin.Context) {
  81. binding := service.SessionBindingFromContext(c.Request.Context())
  82. require.Len(t, binding.UserAgent, maxPersistentUserAgentBytes)
  83. require.Equal(t, binding.UserAgent, c.Request.UserAgent())
  84. c.Status(200)
  85. })
  86. w := httptest.NewRecorder()
  87. req := httptest.NewRequest("GET", "/t", nil)
  88. req.Header.Set("User-Agent", strings.Repeat("u", 2048))
  89. r.ServeHTTP(w, req)
  90. require.Equal(t, 200, w.Code)
  91. }
  92. // 未经过 SessionBindingContext 注入时(异常挂载顺序/单测直调),回退 trusted_proxies 链,
  93. // 等价于开关关闭时的历史行为。
  94. func TestSecurityClientIPFallsBackWithoutInjectedBinding(t *testing.T) {
  95. gin.SetMode(gin.TestMode)
  96. r := gin.New()
  97. require.NoError(t, r.SetTrustedProxies(nil))
  98. r.GET("/t", func(c *gin.Context) {
  99. c.String(200, SecurityClientIP(c))
  100. })
  101. w := httptest.NewRecorder()
  102. req := httptest.NewRequest("GET", "/t", nil)
  103. req.RemoteAddr = "9.9.9.9:12345"
  104. req.Header.Set("X-Real-IP", "1.2.3.4")
  105. r.ServeHTTP(w, req)
  106. require.Equal(t, 200, w.Code)
  107. require.Equal(t, "9.9.9.9", w.Body.String())
  108. }
  109. func TestRequestSessionBindingPrefersInjectedBinding(t *testing.T) {
  110. gin.SetMode(gin.TestMode)
  111. cfg := &config.Config{}
  112. cfg.SetTrustForwardedIPForAPIKeyACL(true)
  113. r := gin.New()
  114. require.NoError(t, r.SetTrustedProxies([]string{"127.0.0.1"}))
  115. r.Use(SessionBindingContext(cfg))
  116. r.GET("/t", func(c *gin.Context) {
  117. issued := &service.SessionBinding{IP: "1.2.3.4", UserAgent: "test-agent"}
  118. require.Equal(t, issued.Hash(), requestSessionBinding(c).Hash())
  119. c.Status(200)
  120. })
  121. w := httptest.NewRecorder()
  122. req := httptest.NewRequest("GET", "/t", nil)
  123. req.RemoteAddr = "127.0.0.1:54321"
  124. req.Header.Set("X-Real-IP", "1.2.3.4")
  125. req.Header.Set("User-Agent", "test-agent")
  126. r.ServeHTTP(w, req)
  127. require.Equal(t, 200, w.Code)
  128. }