new-api/payment/provider/wxpay.go

package provider

import (
	"bytes"
	"context"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

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

const (
	wxpayDefaultGateway = "https://api.mch.weixin.qq.com"
	wxpayAuthSchema     = "WECHATPAY2-SHA256-RSA256"
	wxpayCurrency       = "CNY"

	wxpayTradeSuccess  = "SUCCESS"
	wxpayTradeClosed   = "CLOSED"
	wxpayTradeRevoked  = "REVOKED"
	wxpayTradePayError = "PAYERROR"
	wxpayTradeNotPay   = "NOTPAY"
	wxpayTradeUserPay  = "USERPAYING"
	wxpayTradeAccept   = "ACCEPT"
)

type Wxpay struct {
	setting    setting.WxpaySetting
	privateKey *rsa.PrivateKey
	publicKey  *rsa.PublicKey
	httpClient *http.Client
}

func NewWxpay(cfg setting.WxpaySetting) (*Wxpay, error) {
	if !cfg.Enabled {
		return nil, errors.New("wxpay is not enabled")
	}
	cfg.AppID = strings.TrimSpace(cfg.AppID)
	cfg.MchID = strings.TrimSpace(cfg.MchID)
	cfg.PrivateKey = strings.TrimSpace(cfg.PrivateKey)
	cfg.MerchantSerialNo = strings.TrimSpace(cfg.MerchantSerialNo)
	cfg.APIv3Key = strings.TrimSpace(cfg.APIv3Key)
	cfg.PlatformPublicKey = strings.TrimSpace(cfg.PlatformPublicKey)
	cfg.PlatformCertificate = strings.TrimSpace(cfg.PlatformCertificate)
	cfg.PlatformSerialNo = strings.TrimSpace(cfg.PlatformSerialNo)
	cfg.GatewayURL = strings.TrimRight(strings.TrimSpace(cfg.GatewayURL), "/")
	if cfg.GatewayURL == "" {
		cfg.GatewayURL = wxpayDefaultGateway
	}
	if cfg.AppID == "" {
		return nil, errors.New("wxpay app_id is required")
	}
	if cfg.MchID == "" {
		return nil, errors.New("wxpay mch_id is required")
	}
	if cfg.PrivateKey == "" {
		return nil, errors.New("wxpay private_key is required")
	}
	if cfg.MerchantSerialNo == "" {
		return nil, errors.New("wxpay merchant_serial_no is required")
	}
	if len(cfg.APIv3Key) != 32 {
		return nil, errors.New("wxpay api_v3_key must be 32 bytes")
	}
	privateKey, err := parseRSAPrivateKey(cfg.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("parse wxpay private_key: %w", err)
	}
	publicKey, err := parseWxpayPlatformPublicKey(cfg)
	if err != nil {
		return nil, err
	}
	return &Wxpay{
		setting:    cfg,
		privateKey: privateKey,
		publicKey:  publicKey,
		httpClient: http.DefaultClient,
	}, nil
}

func (w *Wxpay) Name() string {
	return "wxpay"
}

func (w *Wxpay) CreatePayment(ctx context.Context, req *payment.PaymentRequest) (*payment.PaymentResponse, error) {
	if req == nil {
		return nil, errors.New("missing payment request")
	}
	if strings.TrimSpace(req.OrderID) == "" {
		return nil, errors.New("missing order id")
	}
	if req.Amount <= 0 {
		return nil, errors.New("invalid payment amount")
	}
	notifyURL := nonEmpty(req.NotifyURL, w.setting.NotifyURL)
	if notifyURL == "" {
		return nil, errors.New("wxpay notify_url is required")
	}

	tradeType := strings.ToUpper(strings.TrimSpace(req.TradeType))
	if tradeType == "" {
		tradeType = "NATIVE"
	}
	switch tradeType {
	case "NATIVE":
		return w.prepayNative(ctx, req, notifyURL)
	case "H5":
		return w.prepayH5(ctx, req, notifyURL)
	case "JSAPI":
		return w.prepayJSAPI(ctx, req, notifyURL)
	default:
		return nil, fmt.Errorf("unsupported wxpay trade type: %s", tradeType)
	}
}

func (w *Wxpay) prepayNative(ctx context.Context, req *payment.PaymentRequest, notifyURL string) (*payment.PaymentResponse, error) {
	body, err := common.Marshal(wxpayPrepayRequest{
		AppID:       w.setting.AppID,
		MchID:       w.setting.MchID,
		Description: nonEmpty(req.Description, "Top up"),
		OutTradeNo:  req.OrderID,
		NotifyURL:   notifyURL,
		Amount: wxpayAmount{
			Total:    req.Amount,
			Currency: wxpayCurrency,
		},
	})
	if err != nil {
		return nil, fmt.Errorf("marshal wxpay native request: %w", err)
	}
	respBody, err := w.requestJSON(ctx, http.MethodPost, "/v3/pay/transactions/native", nil, body)
	if err != nil {
		return nil, err
	}
	var resp wxpayNativePrepayResponse
	if err := common.Unmarshal(respBody, &resp); err != nil {
		return nil, fmt.Errorf("decode wxpay native response: %w", err)
	}
	return &payment.PaymentResponse{
		ProviderOrderID: req.OrderID,
		QRCode:          resp.CodeURL,
		RawResponse:     string(respBody),
	}, nil
}

func (w *Wxpay) prepayH5(ctx context.Context, req *payment.PaymentRequest, notifyURL string) (*payment.PaymentResponse, error) {
	if strings.TrimSpace(req.ClientIP) == "" {
		return nil, errors.New("wxpay h5 requires client ip")
	}
	body, err := common.Marshal(wxpayPrepayRequest{
		AppID:       w.setting.AppID,
		MchID:       w.setting.MchID,
		Description: nonEmpty(req.Description, "Top up"),
		OutTradeNo:  req.OrderID,
		NotifyURL:   notifyURL,
		Amount: wxpayAmount{
			Total:    req.Amount,
			Currency: wxpayCurrency,
		},
		SceneInfo: &wxpaySceneInfo{
			PayerClientIP: req.ClientIP,
			H5Info: &wxpayH5Info{
				Type:    "Wap",
				AppName: w.setting.H5AppName,
				AppURL:  w.setting.H5AppURL,
			},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("marshal wxpay h5 request: %w", err)
	}
	respBody, err := w.requestJSON(ctx, http.MethodPost, "/v3/pay/transactions/h5", nil, body)
	if err != nil {
		return nil, err
	}
	var resp wxpayH5PrepayResponse
	if err := common.Unmarshal(respBody, &resp); err != nil {
		return nil, fmt.Errorf("decode wxpay h5 response: %w", err)
	}
	return &payment.PaymentResponse{
		ProviderOrderID: req.OrderID,
		PaymentURL:      appendWxpayRedirectURL(resp.H5URL, nonEmpty(req.ReturnURL, w.setting.ReturnURL)),
		RawResponse:     string(respBody),
	}, nil
}

func (w *Wxpay) prepayJSAPI(ctx context.Context, req *payment.PaymentRequest, notifyURL string) (*payment.PaymentResponse, error) {
	if strings.TrimSpace(req.OpenID) == "" {
		return nil, errors.New("wxpay jsapi requires openid")
	}
	body, err := common.Marshal(wxpayPrepayRequest{
		AppID:       w.setting.AppID,
		MchID:       w.setting.MchID,
		Description: nonEmpty(req.Description, "Top up"),
		OutTradeNo:  req.OrderID,
		NotifyURL:   notifyURL,
		Amount: wxpayAmount{
			Total:    req.Amount,
			Currency: wxpayCurrency,
		},
		Payer: &wxpayPayer{OpenID: req.OpenID},
	})
	if err != nil {
		return nil, fmt.Errorf("marshal wxpay jsapi request: %w", err)
	}
	respBody, err := w.requestJSON(ctx, http.MethodPost, "/v3/pay/transactions/jsapi", nil, body)
	if err != nil {
		return nil, err
	}
	return &payment.PaymentResponse{
		ProviderOrderID: req.OrderID,
		RawResponse:     string(respBody),
	}, nil
}

func (w *Wxpay) QueryOrder(ctx context.Context, providerOrderID string) (*payment.OrderStatus, error) {
	providerOrderID = strings.TrimSpace(providerOrderID)
	if providerOrderID == "" {
		return nil, errors.New("missing provider order id")
	}
	path := "/v3/pay/transactions/out-trade-no/" + url.PathEscape(providerOrderID)
	respBody, err := w.requestJSON(ctx, http.MethodGet, path, url.Values{"mchid": []string{w.setting.MchID}}, nil)
	if err != nil {
		return nil, err
	}
	var tx wxpayTransaction
	if err := common.Unmarshal(respBody, &tx); err != nil {
		return nil, fmt.Errorf("decode wxpay query response: %w", err)
	}
	return wxpayOrderStatusFromTransaction(&tx, w.setting.MchID)
}

func (w *Wxpay) VerifyNotification(ctx context.Context, body []byte, headers map[string]string) (*payment.OrderStatus, error) {
	if len(body) == 0 {
		return nil, errors.New("missing wxpay notification body")
	}
	if err := w.verifyNotificationSignature(body, headers); err != nil {
		return nil, err
	}

	var notification wxpayNotification
	if err := common.Unmarshal(body, &notification); err != nil {
		return nil, fmt.Errorf("decode wxpay notification: %w", err)
	}
	plain, err := w.decryptResource(notification.Resource)
	if err != nil {
		return nil, err
	}
	var tx wxpayTransaction
	if err := common.Unmarshal(plain, &tx); err != nil {
		return nil, fmt.Errorf("decode wxpay transaction: %w", err)
	}
	return wxpayOrderStatusFromTransaction(&tx, w.setting.MchID)
}

func (w *Wxpay) Refund(ctx context.Context, providerOrderID string, amount int) error {
	providerOrderID = strings.TrimSpace(providerOrderID)
	if providerOrderID == "" {
		return errors.New("missing provider order id")
	}
	if amount <= 0 {
		return errors.New("invalid refund amount")
	}
	status, err := w.QueryOrder(ctx, providerOrderID)
	if err != nil {
		return err
	}
	body, err := common.Marshal(wxpayRefundRequest{
		OutTradeNo:  providerOrderID,
		OutRefundNo: fmt.Sprintf("%s-refund-%d", providerOrderID, time.Now().UnixNano()),
		Amount: wxpayRefundAmount{
			Refund:   amount,
			Total:    status.Amount,
			Currency: wxpayCurrency,
		},
	})
	if err != nil {
		return fmt.Errorf("marshal wxpay refund request: %w", err)
	}
	_, err = w.requestJSON(ctx, http.MethodPost, "/v3/refund/domestic/refunds", nil, body)
	return err
}

func (w *Wxpay) requestJSON(ctx context.Context, method string, path string, query url.Values, body []byte) ([]byte, error) {
	target := w.setting.GatewayURL + path
	canonicalURL := path
	if len(query) > 0 {
		encodedQuery := query.Encode()
		target += "?" + encodedQuery
		canonicalURL += "?" + encodedQuery
	}
	req, err := http.NewRequestWithContext(ctx, method, target, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("build wxpay request: %w", err)
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	req.Header.Set("Accept", "application/json")
	auth, err := w.authorizationHeader(method, canonicalURL, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", auth)
	client := w.httpClient
	if client == nil {
		client = http.DefaultClient
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request wxpay gateway: %w", err)
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read wxpay response: %w", err)
	}
	if resp.StatusCode >= http.StatusBadRequest {
		return nil, fmt.Errorf("wxpay gateway http status %d: %s", resp.StatusCode, string(respBody))
	}
	return respBody, nil
}

func (w *Wxpay) authorizationHeader(method string, canonicalURL string, body []byte) (string, error) {
	timestamp := strconv.FormatInt(time.Now().Unix(), 10)
	nonce := common.GetRandomString(32)
	signature, err := w.signMessage(wxpayBuildMessage(method, canonicalURL, timestamp, nonce, string(body)))
	if err != nil {
		return "", err
	}
	return fmt.Sprintf(`%s mchid="%s",nonce_str="%s",timestamp="%s",serial_no="%s",signature="%s"`,
		wxpayAuthSchema,
		w.setting.MchID,
		nonce,
		timestamp,
		w.setting.MerchantSerialNo,
		signature,
	), nil
}

func (w *Wxpay) signMessage(message string) (string, error) {
	digest := sha256.Sum256([]byte(message))
	signature, err := rsa.SignPKCS1v15(rand.Reader, w.privateKey, crypto.SHA256, digest[:])
	if err != nil {
		return "", fmt.Errorf("sign wxpay request: %w", err)
	}
	return base64.StdEncoding.EncodeToString(signature), nil
}

func (w *Wxpay) verifyNotificationSignature(body []byte, headers map[string]string) error {
	timestamp := getWxpayHeader(headers, "Wechatpay-Timestamp")
	nonce := getWxpayHeader(headers, "Wechatpay-Nonce")
	signature := getWxpayHeader(headers, "Wechatpay-Signature")
	serial := getWxpayHeader(headers, "Wechatpay-Serial")
	if timestamp == "" || nonce == "" || signature == "" {
		return errors.New("missing wxpay notification signature headers")
	}
	if w.setting.PlatformSerialNo != "" && serial != w.setting.PlatformSerialNo {
		return errors.New("wxpay platform serial mismatch")
	}
	rawSignature, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return fmt.Errorf("decode wxpay notification signature: %w", err)
	}
	message := timestamp + "\n" + nonce + "\n" + string(body) + "\n"
	digest := sha256.Sum256([]byte(message))
	if err := rsa.VerifyPKCS1v15(w.publicKey, crypto.SHA256, digest[:], rawSignature); err != nil {
		return fmt.Errorf("verify wxpay notification signature: %w", err)
	}
	return nil
}

func (w *Wxpay) decryptResource(resource wxpayResource) ([]byte, error) {
	if resource.Ciphertext == "" || resource.Nonce == "" {
		return nil, errors.New("missing wxpay notification resource")
	}
	ciphertext, err := base64.StdEncoding.DecodeString(resource.Ciphertext)
	if err != nil {
		return nil, fmt.Errorf("decode wxpay resource ciphertext: %w", err)
	}
	block, err := aes.NewCipher([]byte(w.setting.APIv3Key))
	if err != nil {
		return nil, fmt.Errorf("create wxpay aes cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("create wxpay gcm: %w", err)
	}
	plain, err := gcm.Open(nil, []byte(resource.Nonce), ciphertext, []byte(resource.AssociatedData))
	if err != nil {
		return nil, fmt.Errorf("decrypt wxpay notification resource: %w", err)
	}
	return plain, nil
}

func wxpayBuildMessage(method string, canonicalURL string, timestamp string, nonce string, body string) string {
	return strings.ToUpper(method) + "\n" + canonicalURL + "\n" + timestamp + "\n" + nonce + "\n" + body + "\n"
}

func wxpayOrderStatusFromTransaction(tx *wxpayTransaction, expectedMchID string) (*payment.OrderStatus, error) {
	if tx == nil {
		return nil, errors.New("missing wxpay transaction")
	}
	if tx.MchID != expectedMchID {
		return nil, errors.New("wxpay mchid mismatch")
	}
	status := &payment.OrderStatus{
		ProviderOrderID: nonEmpty(tx.OutTradeNo, tx.TransactionID),
		Status:          mapWxpayTradeStatus(tx.TradeState),
		Amount:          tx.Amount.Total,
		Currency:        nonEmpty(tx.Amount.Currency, wxpayCurrency),
	}
	if paidAt := parseWxpayTime(tx.SuccessTime); paidAt != nil {
		status.PaidAt = paidAt
	}
	return status, nil
}

func mapWxpayTradeStatus(status string) string {
	switch strings.ToUpper(strings.TrimSpace(status)) {
	case wxpayTradeSuccess:
		return payment.OrderStatusSuccess
	case wxpayTradeClosed, wxpayTradeRevoked:
		return payment.OrderStatusClosed
	case wxpayTradePayError:
		return payment.OrderStatusFailed
	case wxpayTradeNotPay, wxpayTradeUserPay, wxpayTradeAccept:
		return payment.OrderStatusPending
	default:
		return payment.OrderStatusFailed
	}
}

func parseWxpayTime(value string) *time.Time {
	if strings.TrimSpace(value) == "" {
		return nil
	}
	parsed, err := time.Parse(time.RFC3339, value)
	if err != nil {
		return nil
	}
	return &parsed
}

func getWxpayHeader(headers map[string]string, name string) string {
	for key, value := range headers {
		if strings.EqualFold(key, name) {
			return strings.TrimSpace(value)
		}
	}
	return ""
}

func parseWxpayPlatformPublicKey(cfg setting.WxpaySetting) (*rsa.PublicKey, error) {
	if cfg.PlatformPublicKey != "" {
		publicKey, err := parseRSAPublicKey(cfg.PlatformPublicKey)
		if err != nil {
			return nil, fmt.Errorf("parse wxpay platform_public_key: %w", err)
		}
		return publicKey, nil
	}
	if cfg.PlatformCertificate == "" {
		return nil, errors.New("wxpay platform_public_key or platform_certificate is required")
	}
	block, _ := pem.Decode([]byte(strings.TrimSpace(strings.ReplaceAll(cfg.PlatformCertificate, `\n`, "\n"))))
	if block == nil || block.Type != "CERTIFICATE" {
		return nil, errors.New("invalid wxpay platform_certificate")
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse wxpay platform_certificate: %w", err)
	}
	publicKey, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("wxpay platform_certificate public key is not RSA")
	}
	return publicKey, nil
}

func appendWxpayRedirectURL(h5URL string, returnURL string) string {
	h5URL = strings.TrimSpace(h5URL)
	returnURL = strings.TrimSpace(returnURL)
	if h5URL == "" || returnURL == "" {
		return h5URL
	}
	sep := "&"
	if !strings.Contains(h5URL, "?") {
		sep = "?"
	}
	return h5URL + sep + "redirect_url=" + url.QueryEscape(returnURL)
}

type wxpayPrepayRequest struct {
	AppID       string          `json:"appid"`
	MchID       string          `json:"mchid"`
	Description string          `json:"description"`
	OutTradeNo  string          `json:"out_trade_no"`
	NotifyURL   string          `json:"notify_url"`
	Amount      wxpayAmount     `json:"amount"`
	SceneInfo   *wxpaySceneInfo `json:"scene_info,omitempty"`
	Payer       *wxpayPayer     `json:"payer,omitempty"`
}

type wxpayAmount struct {
	Total    int    `json:"total"`
	Currency string `json:"currency"`
}

type wxpaySceneInfo struct {
	PayerClientIP string       `json:"payer_client_ip,omitempty"`
	H5Info        *wxpayH5Info `json:"h5_info,omitempty"`
}

type wxpayH5Info struct {
	Type    string `json:"type"`
	AppName string `json:"app_name,omitempty"`
	AppURL  string `json:"app_url,omitempty"`
}

type wxpayPayer struct {
	OpenID string `json:"openid"`
}

type wxpayNativePrepayResponse struct {
	CodeURL string `json:"code_url"`
}

type wxpayH5PrepayResponse struct {
	H5URL string `json:"h5_url"`
}

type wxpayTransaction struct {
	AppID         string      `json:"appid"`
	MchID         string      `json:"mchid"`
	OutTradeNo    string      `json:"out_trade_no"`
	TransactionID string      `json:"transaction_id"`
	TradeState    string      `json:"trade_state"`
	SuccessTime   string      `json:"success_time"`
	Amount        wxpayAmount `json:"amount"`
}

type wxpayNotification struct {
	ID           string        `json:"id"`
	CreateTime   string        `json:"create_time"`
	EventType    string        `json:"event_type"`
	ResourceType string        `json:"resource_type"`
	Resource     wxpayResource `json:"resource"`
}

type wxpayResource struct {
	Algorithm      string `json:"algorithm"`
	Ciphertext     string `json:"ciphertext"`
	AssociatedData string `json:"associated_data"`
	Nonce          string `json:"nonce"`
}

type wxpayRefundRequest struct {
	OutTradeNo  string            `json:"out_trade_no"`
	OutRefundNo string            `json:"out_refund_no"`
	Amount      wxpayRefundAmount `json:"amount"`
}

type wxpayRefundAmount struct {
	Refund   int    `json:"refund"`
	Total    int    `json:"total"`
	Currency string `json:"currency"`
}

var _ payment.Provider = (*Wxpay)(nil)