usdtman/cmd/server/main.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"sync"
	"usdtman"

	"github.com/gorilla/websocket"
)

var (
	uman          *usdtman.USDTMan
	paymentEvents []usdtman.USDTPayment
	paymentLock   sync.RWMutex
	upgrader      = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}
	clients     = make(map[*websocket.Conn]bool)
	clientsLock sync.RWMutex
)

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "8084"
	}

	apiKey := os.Getenv("TRON_API_KEY")
	if apiKey == "" {
		apiKey = "da1e77dc-b35b-4458-846a-5a551b9df4b2"
	}

	// 初始化 USDTMan (无初始地址)
	uman = usdtman.NewUSDTMan([]string{}, apiKey)

	// 设置收款回调
	uman.OnPaymentComplete(func(payment *usdtman.USDTPayment) {
		log.Printf("💰 收到 USDT: %s -> %.6f USDT (确认数: %d, TxID: %s)",
			payment.From, payment.Amount, payment.Confirmations, payment.TxID)

		paymentLock.Lock()
		paymentEvents = append(paymentEvents, *payment)
		if len(paymentEvents) > 100 {
			paymentEvents = paymentEvents[len(paymentEvents)-100:]
		}
		paymentLock.Unlock()

		broadcastPayment(payment)
	})

	http.HandleFunc("/start", startMonitor)
	http.HandleFunc("/stop", stopMonitor)
	http.HandleFunc("/add-address", addAddress)
	http.HandleFunc("/remove-address", removeAddress)
	http.HandleFunc("/list-addresses", listAddresses)
	http.HandleFunc("/payments", getPayments)
	http.HandleFunc("/ws", handleWebSocket)
	http.HandleFunc("/", serveIndex)

	log.Printf("🚀 USDTMan Server running on :%s", port)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}

func jsonResponse(w http.ResponseWriter, success bool, message string, data interface{}) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"success": success,
		"message": message,
		"data":    data,
	})
}

func startMonitor(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		jsonResponse(w, false, "Method not allowed", nil)
		return
	}

	if err := uman.Start(); err != nil {
		jsonResponse(w, false, fmt.Sprintf("启动失败: %v", err), nil)
		return
	}

	jsonResponse(w, true, "监听已启动", nil)
}

func stopMonitor(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		jsonResponse(w, false, "Method not allowed", nil)
		return
	}

	uman.Stop()
	jsonResponse(w, true, "监听已停止", nil)
}

func addAddress(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		jsonResponse(w, false, "Method not allowed", nil)
		return
	}

	var req struct {
		Address string `json:"address"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		jsonResponse(w, false, "invalid request", nil)
		return
	}

	if req.Address == "" {
		jsonResponse(w, false, "address is required", nil)
		return
	}

	uman.AddAddress(req.Address)
	jsonResponse(w, true, "地址已添加", map[string]interface{}{
		"address": req.Address,
	})
}

func removeAddress(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		jsonResponse(w, false, "Method not allowed", nil)
		return
	}

	var req struct {
		Address string `json:"address"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		jsonResponse(w, false, "invalid request", nil)
		return
	}

	uman.RemoveAddress(req.Address)
	jsonResponse(w, true, "地址已移除", map[string]interface{}{
		"address": req.Address,
	})
}

func listAddresses(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		jsonResponse(w, false, "Method not allowed", nil)
		return
	}

	addresses := uman.GetAddresses()
	jsonResponse(w, true, "success", map[string]interface{}{
		"addresses": addresses,
		"count":     len(addresses),
	})
}

func getPayments(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		jsonResponse(w, false, "Method not allowed", nil)
		return
	}

	paymentLock.RLock()
	events := make([]usdtman.USDTPayment, len(paymentEvents))
	copy(events, paymentEvents)
	paymentLock.RUnlock()

	jsonResponse(w, true, "success", map[string]interface{}{
		"payments": events,
		"count":    len(events),
	})
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket upgrade error:", err)
		return
	}

	clientsLock.Lock()
	clients[conn] = true
	clientsLock.Unlock()

	defer func() {
		clientsLock.Lock()
		delete(clients, conn)
		clientsLock.Unlock()
		conn.Close()
	}()

	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
}

func broadcastPayment(payment *usdtman.USDTPayment) {
	message := map[string]interface{}{
		"type":          "usdt_payment",
		"address":       payment.Address,
		"amount":        payment.Amount,
		"from":          payment.From,
		"txId":          payment.TxID,
		"block":         payment.BlockNumber,
		"time":          payment.Timestamp,
		"confirmations": payment.Confirmations,
	}

	clientsLock.RLock()
	defer clientsLock.RUnlock()

	for conn := range clients {
		conn.WriteJSON(message)
	}
}

func serveIndex(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}