new-api/payment/provider/wxpay_test.go

package provider

import (
	"context"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/payment"
	"github.com/QuantumNous/new-api/setting"
)

func testWxpayKeyPair(t *testing.T) (*rsa.PrivateKey, string, string) {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate rsa key: %v", err)
	}
	privateDER, err := x509.MarshalPKCS8PrivateKey(key)
	if err != nil {
		t.Fatalf("marshal private key: %v", err)
	}
	publicDER, err := x509.MarshalPKIXPublicKey(&key.PublicKey)
	if err != nil {
		t.Fatalf("marshal public key: %v", err)
	}
	return key,
		string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privateDER})),
		string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: publicDER}))
}

func testWxpayProvider(t *testing.T) (*Wxpay, *rsa.PrivateKey) {
	t.Helper()
	key, privatePEM, publicPEM := testWxpayKeyPair(t)
	provider, err := NewWxpay(setting.WxpaySetting{
		Enabled:           true,
		AppID:             "wx-app-123",
		MchID:             "mch-123",
		PrivateKey:        privatePEM,
		MerchantSerialNo:  "merchant-serial",
		APIv3Key:          "12345678901234567890123456789012",
		PlatformPublicKey: publicPEM,
		PlatformSerialNo:  "platform-serial",
		GatewayURL:        "https://example.test",
	})
	if err != nil {
		t.Fatalf("new wxpay: %v", err)
	}
	return provider, key
}

func TestWxpayCreatePaymentNativeBuildsSignedRequest(t *testing.T) {
	provider, _ := testWxpayProvider(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v3/pay/transactions/native" {
			t.Fatalf("path = %q", r.URL.Path)
		}
		auth := r.Header.Get("Authorization")
		if !strings.HasPrefix(auth, wxpayAuthSchema+" ") {
			t.Fatalf("missing wxpay authorization header: %q", auth)
		}
		var req wxpayPrepayRequest
		if err := common.DecodeJson(r.Body, &req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		if req.OutTradeNo != "ORDER-1" {
			t.Fatalf("out_trade_no = %q", req.OutTradeNo)
		}
		if req.Amount.Total != 1234 {
			t.Fatalf("amount.total = %d", req.Amount.Total)
		}
		body, err := common.Marshal(wxpayNativePrepayResponse{CodeURL: "weixin://wxpay/bizpayurl?pr=test"})
		if err != nil {
			t.Fatalf("marshal response: %v", err)
		}
		_, _ = w.Write(body)
	}))
	defer server.Close()
	provider.setting.GatewayURL = server.URL

	resp, err := provider.CreatePayment(context.Background(), &payment.PaymentRequest{
		OrderID:     "ORDER-1",
		Amount:      1234,
		Description: "Recharge",
		NotifyURL:   "https://merchant.test/api/payment/webhook/wxpay",
		TradeType:   "NATIVE",
	})
	if err != nil {
		t.Fatalf("create payment: %v", err)
	}
	if resp.QRCode != "weixin://wxpay/bizpayurl?pr=test" {
		t.Fatalf("qrcode = %q", resp.QRCode)
	}
}

func TestWxpayCreatePaymentH5ReturnsRedirectURL(t *testing.T) {
	provider, _ := testWxpayProvider(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v3/pay/transactions/h5" {
			t.Fatalf("path = %q", r.URL.Path)
		}
		var req wxpayPrepayRequest
		if err := common.DecodeJson(r.Body, &req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		if req.SceneInfo == nil || req.SceneInfo.PayerClientIP != "203.0.113.10" {
			t.Fatalf("payer client ip = %#v", req.SceneInfo)
		}
		body, err := common.Marshal(wxpayH5PrepayResponse{H5URL: "https://wxpay.test/h5"})
		if err != nil {
			t.Fatalf("marshal response: %v", err)
		}
		_, _ = w.Write(body)
	}))
	defer server.Close()
	provider.setting.GatewayURL = server.URL

	resp, err := provider.CreatePayment(context.Background(), &payment.PaymentRequest{
		OrderID:   "ORDER-H5",
		Amount:    1234,
		NotifyURL: "https://merchant.test/api/payment/webhook/wxpay",
		ReturnURL: "https://merchant.test/return",
		ClientIP:  "203.0.113.10",
		TradeType: "H5",
	})
	if err != nil {
		t.Fatalf("create h5 payment: %v", err)
	}
	if !strings.Contains(resp.PaymentURL, "redirect_url=https%3A%2F%2Fmerchant.test%2Freturn") {
		t.Fatalf("payment url missing redirect_url: %q", resp.PaymentURL)
	}
}

func TestWxpayVerifyNotificationDecryptsAndMapsSuccess(t *testing.T) {
	provider, key := testWxpayProvider(t)
	body := signedWxpayNotificationBody(t, provider, wxpayTransaction{
		AppID:         "wx-app-123",
		MchID:         "mch-123",
		OutTradeNo:    "ORDER-2",
		TransactionID: "4200000001",
		TradeState:    wxpayTradeSuccess,
		SuccessTime:   "2026-05-20T12:34:56+08:00",
		Amount: wxpayAmount{
			Total:    1050,
			Currency: wxpayCurrency,
		},
	})
	headers := signWxpayNotificationHeaders(t, key, body)

	status, err := provider.VerifyNotification(context.Background(), body, headers)
	if err != nil {
		t.Fatalf("verify notification: %v", err)
	}
	if status.ProviderOrderID != "ORDER-2" {
		t.Fatalf("provider order id = %q", status.ProviderOrderID)
	}
	if status.Status != payment.OrderStatusSuccess {
		t.Fatalf("status = %q", status.Status)
	}
	if status.Amount != 1050 {
		t.Fatalf("amount = %d", status.Amount)
	}
	if status.PaidAt == nil {
		t.Fatal("paid_at is nil")
	}
}

func TestWxpayVerifyNotificationRejectsMismatchedMchID(t *testing.T) {
	provider, key := testWxpayProvider(t)
	body := signedWxpayNotificationBody(t, provider, wxpayTransaction{
		MchID:      "other-mch",
		OutTradeNo: "ORDER-3",
		TradeState: wxpayTradeSuccess,
		Amount: wxpayAmount{
			Total:    100,
			Currency: wxpayCurrency,
		},
	})
	headers := signWxpayNotificationHeaders(t, key, body)

	if _, err := provider.VerifyNotification(context.Background(), body, headers); err == nil {
		t.Fatal("expected mchid mismatch error")
	}
}

func TestMapWxpayTradeStatus(t *testing.T) {
	tests := []struct {
		status string
		want   string
	}{
		{wxpayTradeSuccess, payment.OrderStatusSuccess},
		{wxpayTradeClosed, payment.OrderStatusClosed},
		{wxpayTradeRevoked, payment.OrderStatusClosed},
		{wxpayTradePayError, payment.OrderStatusFailed},
		{wxpayTradeNotPay, payment.OrderStatusPending},
		{"UNKNOWN", payment.OrderStatusFailed},
	}
	for _, tt := range tests {
		if got := mapWxpayTradeStatus(tt.status); got != tt.want {
			t.Fatalf("mapWxpayTradeStatus(%q) = %q, want %q", tt.status, got, tt.want)
		}
	}
}

func signedWxpayNotificationBody(t *testing.T, provider *Wxpay, tx wxpayTransaction) []byte {
	t.Helper()
	txBody, err := common.Marshal(tx)
	if err != nil {
		t.Fatalf("marshal tx: %v", err)
	}
	block, err := aes.NewCipher([]byte(provider.setting.APIv3Key))
	if err != nil {
		t.Fatalf("aes cipher: %v", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		t.Fatalf("gcm: %v", err)
	}
	nonce := "nonce-123456"
	associatedData := "transaction"
	ciphertext := gcm.Seal(nil, []byte(nonce), txBody, []byte(associatedData))
	body, err := common.Marshal(wxpayNotification{
		ID:           "notify-1",
		EventType:    "TRANSACTION.SUCCESS",
		ResourceType: "encrypt-resource",
		Resource: wxpayResource{
			Algorithm:      "AEAD_AES_256_GCM",
			Ciphertext:     base64.StdEncoding.EncodeToString(ciphertext),
			AssociatedData: associatedData,
			Nonce:          nonce,
		},
	})
	if err != nil {
		t.Fatalf("marshal notification: %v", err)
	}
	return body
}

func signWxpayNotificationHeaders(t *testing.T, key *rsa.PrivateKey, body []byte) map[string]string {
	t.Helper()
	timestamp := "1779251696"
	nonce := "notify-nonce"
	message := timestamp + "\n" + nonce + "\n" + string(body) + "\n"
	digest := sha256.Sum256([]byte(message))
	signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatalf("sign notification: %v", err)
	}
	return map[string]string{
		"Wechatpay-Timestamp": timestamp,
		"Wechatpay-Nonce":     nonce,
		"Wechatpay-Signature": base64.StdEncoding.EncodeToString(signature),
		"Wechatpay-Serial":    "platform-serial",
	}
}