单元测试

2026-08-14 23:37:31 +08:00
parent 170bbdca76
commit 04c7a148e1
4 changed files with 138 additions and 29 deletions
@@ -36,6 +36,12 @@ import (
 // See https://www.iana.org/assignments/media-types/media-types.xhtml
 var allowedMimeTypes = []string{"application/", "audio/", "font/", "image/", "text/", "video/"}

+// RequestPrincipal is a tenant-scoped identity authenticated for HTTP/gRPC file access.
+type RequestPrincipal struct {
+	TenantID types.TenantID
+	UID      types.Uid
+}
+
 func largeFileServeHTTP(wrt http.ResponseWriter, req *http.Request) {
 	now := types.TimeNow()
 	enc := json.NewEncoder(wrt)
@@ -86,7 +92,7 @@ func largeFileServeHTTP(wrt http.ResponseWriter, req *http.Request) {

 	// Check authorization: either auth information or SID must be present
 	authMethod, secret := getHttpAuth(req)
-	tenantID, uid, challenge, err := authFileRequest(authMethod, secret, req.FormValue("sid"), getRemoteAddr(req))
+	principal, challenge, err := authFileRequest(authMethod, secret, req.FormValue("sid"), getRemoteAddr(req))
 	if err != nil {
 		writeHttpResponse(decodeStoreError(err, "", now, nil), err)
 		return
@@ -97,14 +103,14 @@ func largeFileServeHTTP(wrt http.ResponseWriter, req *http.Request) {
 		return
 	}

-	if uid.IsZero() {
+	if principal.UID.IsZero() {
 		// Not authenticated
 		writeHttpResponse(ErrAuthRequired("", "", now, now), errors.New("user not authenticated"))
 		return
 	}

 	// Check if media handler redirects or adds headers.
-	headers, statusCode, err := mh.Headers(tenantID, req.Method, req.URL, req.Header, true)
+	headers, statusCode, err := mh.Headers(principal.TenantID, req.Method, req.URL, req.Header, true)
 	if err != nil {
 		writeHttpResponse(decodeStoreError(err, "", now, nil), err)
 		return
@@ -128,11 +134,11 @@ func largeFileServeHTTP(wrt http.ResponseWriter, req *http.Request) {
 				},
 			})
 		}
-		logs.Info.Println("media serve: completed with status", statusCode, "uid=", uid)
+		logs.Info.Println("media serve: completed with status", statusCode, "uid=", principal.UID)
 		return
 	}

-	fd, rsc, err := mh.Download(tenantID, req.URL.String())
+	fd, rsc, err := mh.Download(principal.TenantID, req.URL.String())
 	if err != nil {
 		writeHttpResponse(decodeStoreError(err, "", now, nil), err)
 		return
@@ -140,7 +146,7 @@ func largeFileServeHTTP(wrt http.ResponseWriter, req *http.Request) {

 	defer rsc.Close()

-	if err = authorizeLargeFileDownload(tenantID, uid, req.URL.String(), fd); err != nil {
+	if err = authorizeLargeFileDownload(principal.TenantID, principal.UID, req.URL.String(), fd); err != nil {
 		writeHttpResponse(decodeStoreError(err, "", now, nil), err)
 		return
 	}
@@ -167,7 +173,7 @@ func largeFileServeHTTP(wrt http.ResponseWriter, req *http.Request) {

 	http.ServeContent(wrt, req, "", fd.UpdatedAt, rsc)

-	logs.Info.Println("media serve: OK, uid=", uid)
+	logs.Info.Println("media serve: OK, uid=", principal.UID)
 }

 // largeFileReceiveHTTP receives files from client over HTTP(S) and passes them to the configured media handler.
@@ -228,7 +234,7 @@ func largeFileReceiveHTTP(wrt http.ResponseWriter, req *http.Request) {
 	msgID := req.FormValue("id")
 	// Check authorization: either auth information or SID must be present
 	authMethod, secret := getHttpAuth(req)
-	tenantID, uid, challenge, err := authFileRequest(authMethod, secret, req.FormValue("sid"), getRemoteAddr(req))
+	principal, challenge, err := authFileRequest(authMethod, secret, req.FormValue("sid"), getRemoteAddr(req))
 	if err != nil {
 		writeHttpResponse(decodeStoreError(err, msgID, now, nil), err)
 		return
@@ -237,14 +243,14 @@ func largeFileReceiveHTTP(wrt http.ResponseWriter, req *http.Request) {
 		writeHttpResponse(InfoChallenge(msgID, now, challenge), nil)
 		return
 	}
-	if uid.IsZero() && req.FormValue("topic") != "newacc" {
+	if principal.UID.IsZero() && req.FormValue("topic") != "newacc" {
 		// Not authenticated and not signup.
 		writeHttpResponse(ErrAuthRequired(msgID, "", now, now), nil)
 		return
 	}

 	// Check if uploads are handled elsewhere.
-	headers, statusCode, err := mh.Headers(tenantID, req.Method, req.URL, req.Header, true)
+	headers, statusCode, err := mh.Headers(principal.TenantID, req.Method, req.URL, req.Header, true)
 	if err != nil {
 		logs.Info.Println("media upload: headers check failed", err)
 		writeHttpResponse(decodeStoreError(err, "", now, nil), err)
@@ -316,8 +322,8 @@ func largeFileReceiveHTTP(wrt http.ResponseWriter, req *http.Request) {
 		ObjHeader: types.ObjHeader{
 			Id: store.Store.GetUidString(),
 		},
-		TenantID: tenantID,
-		User:     uid.String(),
+		TenantID: principal.TenantID,
+		User:     principal.UID.String(),
 		MimeType: mimeType,
 	}
 	fdef.InitTimes()
@@ -330,12 +336,12 @@ func largeFileReceiveHTTP(wrt http.ResponseWriter, req *http.Request) {
 	url, size, err := mh.Upload(fdef, file)
 	if err != nil {
 		logs.Info.Println("media upload: failed", file, "key", fdef.Location, err)
-		store.Files.FinishUpload(tenantID, fdef, false, 0)
+		store.Files.FinishUpload(principal.TenantID, fdef, false, 0)
 		writeHttpResponse(decodeStoreError(err, msgID, now, nil), err)
 		return
 	}

-	fdef, err = store.Files.FinishUpload(tenantID, fdef, true, size)
+	fdef, err = store.Files.FinishUpload(principal.TenantID, fdef, true, size)
 	if err != nil {
 		logs.Info.Println("media upload: failed to finalize", file, "key", fdef.Location, err)
 		// Best effort cleanup.
@@ -373,7 +379,7 @@ func (*grpcNodeServer) LargeFileServe(req *pbx.FileDownReq, stream pbx.Node_Larg
 	if p, ok := peer.FromContext(stream.Context()); ok {
 		remoteAddr = p.Addr.String()
 	}
-	tenantID, uid, challenge, err := authFileRequest(authMethod, secret, "", remoteAddr)
+	principal, challenge, err := authFileRequest(authMethod, secret, "", remoteAddr)
 	if err != nil {
 		writeResponse(decodeStoreError(err, msgID, now, nil), err)
 		return nil
@@ -384,7 +390,7 @@ func (*grpcNodeServer) LargeFileServe(req *pbx.FileDownReq, stream pbx.Node_Larg
 		return nil
 	}

-	if uid.IsZero() {
+	if principal.UID.IsZero() {
 		// Not authenticated
 		writeResponse(ErrAuthRequired(msgID, "", now, now), errors.New("user not authenticated"))
 		return nil
@@ -393,7 +399,7 @@ func (*grpcNodeServer) LargeFileServe(req *pbx.FileDownReq, stream pbx.Node_Larg
 	// Check if media handler redirects or adds headers.
 	mh := store.Store.GetMediaHandler()
 	url, _ := url.Parse(req.Uri)
-	headers, statusCode, err := mh.Headers(tenantID, http.MethodGet, url, http.Header{}, true)
+	headers, statusCode, err := mh.Headers(principal.TenantID, http.MethodGet, url, http.Header{}, true)
 	if err != nil {
 		writeResponse(decodeStoreError(err, "", now, nil), err)
 		return nil
@@ -406,11 +412,11 @@ func (*grpcNodeServer) LargeFileServe(req *pbx.FileDownReq, stream pbx.Node_Larg
 		resp.Text = http.StatusText(statusCode)
 		resp.RedirUrl = headers.Get("Location")
 		stream.Send(&resp)
-		logs.Info.Println("media serve: completed with status", statusCode, "uid=", uid)
+		logs.Info.Println("media serve: completed with status", statusCode, "uid=", principal.UID)
 		return nil
 	}

-	fd, rsc, err := mh.Download(tenantID, req.GetUri())
+	fd, rsc, err := mh.Download(principal.TenantID, req.GetUri())
 	if err != nil {
 		writeResponse(decodeStoreError(err, msgID, now, nil), err)
 		return nil
@@ -418,7 +424,7 @@ func (*grpcNodeServer) LargeFileServe(req *pbx.FileDownReq, stream pbx.Node_Larg

 	defer rsc.Close()

-	if err = authorizeLargeFileDownload(tenantID, uid, req.GetUri(), fd); err != nil {
+	if err = authorizeLargeFileDownload(principal.TenantID, principal.UID, req.GetUri(), fd); err != nil {
 		writeResponse(decodeStoreError(err, msgID, now, nil), err)
 		return nil
 	}
@@ -437,7 +443,7 @@ func (*grpcNodeServer) LargeFileServe(req *pbx.FileDownReq, stream pbx.Node_Larg
 		if err == nil {
 			resp.Content = resp.Content[:n]
 			if err = stream.Send(&resp); err != nil {
-				logs.Info.Println("media serve: failed, uid=", uid, err)
+				logs.Info.Println("media serve: failed, uid=", principal.UID, err)
 				break
 			}
 			continue
@@ -449,7 +455,7 @@ func (*grpcNodeServer) LargeFileServe(req *pbx.FileDownReq, stream pbx.Node_Larg
 		}
 		break
 	}
-	logs.Info.Println("media serve: ", result, ", uid=", uid)
+	logs.Info.Println("media serve: ", result, ", uid=", principal.UID)
 	return err
 }

@@ -522,7 +528,7 @@ func (*grpcNodeServer) LargeFileReceive(stream pbx.Node_LargeFileReceiveServer)
 	if p, ok := peer.FromContext(stream.Context()); ok {
 		remoteAddr = p.Addr.String()
 	}
-	tenantID, uid, challenge, err := authFileRequest(authMethod, secret, "", remoteAddr)
+	principal, challenge, err := authFileRequest(authMethod, secret, "", remoteAddr)
 	if err != nil {
 		writeResponse(decodeStoreError(err, msgID, now, nil), err)
 		return nil
@@ -533,14 +539,14 @@ func (*grpcNodeServer) LargeFileReceive(stream pbx.Node_LargeFileReceiveServer)
 		return nil
 	}

-	if uid.IsZero() {
+	if principal.UID.IsZero() {
 		// Not authenticated
 		writeResponse(ErrAuthRequired(msgID, "", now, now), errors.New("user not authenticated"))
 		return nil
 	}

 	// Check if uploads are handled elsewhere.
-	headers, statusCode, err := mh.Headers(tenantID, http.MethodPost, nil, http.Header{}, false)
+	headers, statusCode, err := mh.Headers(principal.TenantID, http.MethodPost, nil, http.Header{}, false)
 	if err != nil {
 		logs.Info.Println("media upload: headers check failed", err)
 		writeResponse(decodeStoreError(err, "", now, nil), nil)
@@ -555,7 +561,7 @@ func (*grpcNodeServer) LargeFileReceive(stream pbx.Node_LargeFileReceiveServer)
 			Text:     http.StatusText(statusCode),
 			RedirUrl: headers.Get("Location"),
 		})
-		logs.Info.Println("media upload: completed with status", statusCode, "uid=", uid, err)
+		logs.Info.Println("media upload: completed with status", statusCode, "uid=", principal.UID, err)
 		return err
 	}

@@ -571,8 +577,8 @@ func (*grpcNodeServer) LargeFileReceive(stream pbx.Node_LargeFileReceiveServer)
 		ObjHeader: types.ObjHeader{
 			Id: store.Store.GetUidString(),
 		},
-		TenantID: tenantID,
-		User:     uid.String(),
+		TenantID: principal.TenantID,
+		User:     principal.UID.String(),
 		MimeType: mimeType,
 	}
 	fdef.InitTimes()
@@ -606,7 +612,7 @@ func (*grpcNodeServer) LargeFileReceive(stream pbx.Node_LargeFileReceiveServer)
 	}
 	if err != nil {
 		logs.Info.Println("media upload: failed", req.Meta.Name, "key", fdef.Location, err)
-		store.Files.FinishUpload(tenantID, fdef, false, 0)
+		store.Files.FinishUpload(principal.TenantID, fdef, false, 0)
 		writeResponse(decodeStoreError(err, msgID, now, nil), nil)
 		return nil
 	}
@@ -1,15 +1,84 @@
 package main

 import (
+	"encoding/base64"
+	"net/http/httptest"
 	"testing"
 	"time"

 	"github.com/golang/mock/gomock"
+	"github.com/tinode/chat/server/auth"
+	"github.com/tinode/chat/server/auth/mock_auth"
 	"github.com/tinode/chat/server/store"
 	"github.com/tinode/chat/server/store/mock_store"
 	"github.com/tinode/chat/server/store/types"
 )

+func TestAuthFileRequestFromSessionRequiresTenant(t *testing.T) {
+	oldSessionStore := globals.sessionStore
+	globals.sessionStore = NewSessionStore(time.Minute)
+	defer func() { globals.sessionStore = oldSessionStore }()
+
+	sess, _ := globals.sessionStore.NewSession(httptest.NewRecorder(), "file-test-sid")
+	sess.uid = types.Uid(42)
+
+	if _, _, _, err := authFileRequest("", "", sess.sid, ""); err != types.ErrFailed {
+		t.Fatalf("session without tenant: got %v, want %v", err, types.ErrFailed)
+	}
+
+	sess.tenantID = types.TenantID(7)
+	principal, challenge, _, err := authFileRequest("", "", sess.sid, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if challenge != nil {
+		t.Fatal("unexpected challenge")
+	}
+	if principal.TenantID != sess.tenantID || principal.UID != sess.uid {
+		t.Fatalf("principal mismatch: got %+v, want tenant=%d uid=%s", principal, sess.tenantID, sess.uid)
+	}
+}
+
+func TestAuthFileRequestFromTokenReturnsPrincipal(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	oldStore := store.Store
+	defer func() { store.Store = oldStore }()
+
+	mockStore := mock_store.NewMockPersistentStorageInterface(ctrl)
+	tokenAuth := mock_auth.NewMockAuthHandler(ctrl)
+	mockStore.EXPECT().GetLogicalAuthHandler("token").Return(tokenAuth)
+	tokenAuth.EXPECT().IsInitialized().Return(true)
+	tokenAuth.EXPECT().Authenticate(auth.AuthContext{RemoteAddr: "203.0.113.10"}, []byte("token-1")).
+		Return(&auth.Rec{TenantID: 7, Uid: 42, AuthLevel: auth.LevelAuth}, nil, nil)
+	store.Store = mockStore
+
+	principal, challenge, _, err := authFileRequest("token", base64.StdEncoding.EncodeToString([]byte("token-1")), "", "203.0.113.10")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if challenge != nil {
+		t.Fatal("unexpected challenge")
+	}
+	if principal.TenantID != 7 || principal.UID != 42 {
+		t.Fatalf("principal mismatch: %#v", principal)
+	}
+}
+
+func TestAuthFileRequestAllowsUnauthenticatedNewAccountPath(t *testing.T) {
+	principal, challenge, _, err := authFileRequest("", "", "", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if challenge != nil {
+		t.Fatal("unexpected challenge")
+	}
+	if !principal.TenantID.IsZero() || !principal.UID.IsZero() {
+		t.Fatalf("principal should be empty for unauthenticated request: %+v", principal)
+	}
+}
+
 func TestAuthorizeLargeFileDownloadDenied(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -31,6 +100,19 @@ func TestAuthorizeLargeFileDownloadDenied(t *testing.T) {
 	}
 }

+func TestAuthorizeLargeFileDownloadRejectsCrossTenant(t *testing.T) {
+	fd := &types.FileDef{
+		ObjHeader: types.ObjHeader{Id: types.Uid(100).String(), UpdatedAt: time.Date(2026, 8, 13, 9, 0, 0, 0, time.UTC)},
+		TenantID:  8,
+		Status:    types.UploadCompleted,
+		MimeType:  "text/plain",
+		Location:  "notes.txt",
+	}
+	if err := authorizeLargeFileDownload(7, 42, "", fd); err != types.ErrNotFound {
+		t.Fatalf("expected not found for cross-tenant file, got %v", err)
+	}
+}
+
 func TestContentDispositionAttachmentSanitizesName(t *testing.T) {
 	fd := &types.FileDef{
 		ObjHeader: types.ObjHeader{Id: types.Uid(100).String()},
@@ -838,6 +838,26 @@ func TestServeIMVerificationCodesCreatesChallenge(t *testing.T) {
 	}
 }

+func TestServeIMVerificationCodesRequiresTenantCode(t *testing.T) {
+	req := httptest.NewRequest(http.MethodPost, "/v1/im/auth/verification-codes",
+		bytes.NewBufferString(`{"scene":"register","phone":"+8613800138000"}`))
+	req.Header.Set("X-Tinode-APIKey", makeTestAPIKey(t))
+	rec := httptest.NewRecorder()
+
+	serveIMVerificationCodes(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status: got %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
+	}
+	resp := decodeIMHTTPResponse(t, rec)
+	if resp.Reason != imReasonMalformedRequest {
+		t.Fatalf("reason: got %q, want %q", resp.Reason, imReasonMalformedRequest)
+	}
+	if resp.Details == nil {
+		t.Fatal("expected malformed request details")
+	}
+}
+
 func TestServeIMVerificationCodesRateLimited(t *testing.T) {
 	oldTenants := store.Tenants
 	oldVerifications := store.Verifications