优化结构

internal/server/middleware/auth_test.go

@@ -0,0 +1,82 @@
+package middleware
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"kra/internal/biz"
+
+	"github.com/gin-gonic/gin"
+)
+
+type tokenAuthenticatorStub struct {
+	claims *biz.AuthClaims
+	err    error
+	got    string
+}
+
+func (s *tokenAuthenticatorStub) AuthenticateToken(_ context.Context, token string) (*biz.TokenAuthentication, error) {
+	s.got = token
+	if s.err != nil {
+		return nil, s.err
+	}
+	return &biz.TokenAuthentication{Claims: s.claims}, nil
+}
+
+func TestAuthenticateWebSocketAcceptsQueryToken(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	stub := &tokenAuthenticatorStub{claims: &biz.AuthClaims{ID: 7}}
+	recorder := httptest.NewRecorder()
+	context, _ := gin.CreateTestContext(recorder)
+	context.Request = httptest.NewRequest(http.MethodGet, "/ws?token=query-token", nil)
+
+	if !AuthenticateWebSocket(context, stub) {
+		t.Fatalf("AuthenticateWebSocket() = false, body=%s", recorder.Body.String())
+	}
+	if stub.got != "query-token" {
+		t.Fatalf("authenticated token = %q, want query-token", stub.got)
+	}
+	if Claims(context) == nil || Claims(context).ID != 7 {
+		t.Fatalf("claims were not stored: %#v", Claims(context))
+	}
+}
+
+func TestHTTPAuthDoesNotAcceptQueryToken(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	stub := &tokenAuthenticatorStub{claims: &biz.AuthClaims{ID: 7}}
+	engine := gin.New()
+	called := false
+	engine.GET("/protected", Auth(stub), func(c *gin.Context) { called = true })
+
+	response := httptest.NewRecorder()
+	engine.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/protected?token=query-token", nil))
+	if response.Code != http.StatusUnauthorized || called {
+		t.Fatalf("status=%d called=%v body=%s", response.Code, called, response.Body.String())
+	}
+}
+
+func TestAuthenticateWebSocketSupportsBearerHeaderAndReportsInvalidToken(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	stub := &tokenAuthenticatorStub{claims: &biz.AuthClaims{ID: 7}}
+	context, _ := gin.CreateTestContext(httptest.NewRecorder())
+	context.Request = httptest.NewRequest(http.MethodGet, "/ws", nil)
+	context.Request.Header.Set("Authorization", "Bearer header-token")
+
+	if !AuthenticateWebSocket(context, stub) || stub.got != "header-token" {
+		t.Fatalf("Bearer token was not accepted: ok=%v token=%q", Claims(context) != nil, stub.got)
+	}
+
+	invalid := &tokenAuthenticatorStub{err: errors.New("invalid token")}
+	invalidRecorder := httptest.NewRecorder()
+	invalidContext, _ := gin.CreateTestContext(invalidRecorder)
+	invalidContext.Request = httptest.NewRequest(http.MethodGet, "/ws?token=bad-token", nil)
+	if AuthenticateWebSocket(invalidContext, invalid) {
+		t.Fatal("invalid token was accepted")
+	}
+	if invalidRecorder.Code != http.StatusUnauthorized {
+		t.Fatalf("invalid token status=%d, want %d", invalidRecorder.Code, http.StatusUnauthorized)
+	}
+}