pangolin/server/internal/httpapi/subscription.go

package httpapi

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/go-chi/chi/v5"

	"github.com/wangjia/pangolin/server/internal/apierr"
	"github.com/wangjia/pangolin/server/internal/auth"
	"github.com/wangjia/pangolin/server/internal/nodes"
)

// SubscriptionAPI serves the per-user subscription URL used by the web
// user-center ("订阅导入"). The URL points at GET /sub/{token}, which renders the
// user's sing-box client config so it can be imported into a client app.
//
//	GET  /v1/me/subscription        → {"url": "<subBase>/sub/<token>"}
//	POST /v1/me/subscription/reset  → rotates the token, returns the new url
//	GET  /sub/{token}               → public; sing-box client config for that user
type SubscriptionAPI struct {
	db        *sql.DB
	nodes     nodes.NodeStore
	deriveKey string
	subBase   string // e.g. "https://sub.example.com" (no trailing slash); may be empty
}

// NewSubscriptionAPI creates a SubscriptionAPI. subBase is the public base URL
// the subscription link is built from; when empty the request scheme+host is used.
func NewSubscriptionAPI(db *sql.DB, store nodes.NodeStore, deriveKey, subBase string) *SubscriptionAPI {
	return &SubscriptionAPI{db: db, nodes: store, deriveKey: deriveKey, subBase: subBase}
}

type subscriptionResponse struct {
	URL string `json:"url"`
}

// GetSubscription handles GET /v1/me/subscription.
func (s *SubscriptionAPI) GetSubscription(w http.ResponseWriter, r *http.Request) {
	uid, ok := auth.UserIDFromContext(r.Context())
	if !ok {
		apierr.WriteJSON(w, http.StatusUnauthorized, apierr.ErrUnauthorized)
		return
	}
	token, apiErr := s.ensureToken(r.Context(), uid)
	if apiErr != nil {
		apierr.WriteJSON(w, statusFor(apiErr), apiErr)
		return
	}
	writeJSON(w, http.StatusOK, subscriptionResponse{URL: s.urlFor(r, token)})
}

// ResetSubscription handles POST /v1/me/subscription/reset — rotates the token.
func (s *SubscriptionAPI) ResetSubscription(w http.ResponseWriter, r *http.Request) {
	uid, ok := auth.UserIDFromContext(r.Context())
	if !ok {
		apierr.WriteJSON(w, http.StatusUnauthorized, apierr.ErrUnauthorized)
		return
	}
	token, err := newSubToken()
	if err != nil {
		apierr.WriteJSON(w, http.StatusInternalServerError, apierr.ErrInternal)
		return
	}
	if _, err := s.db.ExecContext(r.Context(),
		`UPDATE users SET sub_token = ? WHERE id = ? AND status = 'active'`, token, uid); err != nil {
		apierr.WriteJSON(w, http.StatusInternalServerError, apierr.ErrInternal)
		return
	}
	writeJSON(w, http.StatusOK, subscriptionResponse{URL: s.urlFor(r, token)})
}

// ServeSub handles GET /sub/{token} — public, no JWT. Renders the sing-box client
// config for the user owning the token, against the first active node.
func (s *SubscriptionAPI) ServeSub(w http.ResponseWriter, r *http.Request) {
	token := chi.URLParam(r, "token")
	if token == "" {
		apierr.WriteJSON(w, http.StatusNotFound, apierr.ErrNotFound)
		return
	}

	var dpUUID string
	err := s.db.QueryRowContext(r.Context(),
		`SELECT dp_uuid FROM users WHERE sub_token = ? AND status = 'active'`, token).Scan(&dpUUID)
	if err == sql.ErrNoRows {
		apierr.WriteJSON(w, http.StatusNotFound, apierr.ErrNotFound)
		return
	} else if err != nil {
		apierr.WriteJSON(w, http.StatusInternalServerError, apierr.ErrInternal)
		return
	}

	node, apiErr := s.firstActiveNode(r.Context())
	if apiErr != nil {
		apierr.WriteJSON(w, statusFor(apiErr), apiErr)
		return
	}

	cfg, err := BuildClientConfig(node, dpUUID, s.deriveKey)
	if err != nil {
		apierr.WriteJSON(w, http.StatusInternalServerError, apierr.ErrInternal)
		return
	}
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	_, _ = w.Write(cfg)
}

// ── helpers ──────────────────────────────────────────────────────────────────

// ensureToken returns the user's sub_token, generating and persisting one on
// first access.
func (s *SubscriptionAPI) ensureToken(ctx context.Context, uid int64) (string, *apierr.Error) {
	var token sql.NullString
	if err := s.db.QueryRowContext(ctx,
		`SELECT sub_token FROM users WHERE id = ? AND status = 'active'`, uid).Scan(&token); err == sql.ErrNoRows {
		return "", apierr.ErrNotFound
	} else if err != nil {
		return "", apierr.ErrInternal
	}
	if token.Valid && token.String != "" {
		return token.String, nil
	}
	t, err := newSubToken()
	if err != nil {
		return "", apierr.ErrInternal
	}
	if _, err := s.db.ExecContext(ctx,
		`UPDATE users SET sub_token = ? WHERE id = ?`, t, uid); err != nil {
		return "", apierr.ErrInternal
	}
	return t, nil
}

func (s *SubscriptionAPI) firstActiveNode(ctx context.Context) (*nodes.NodeRow, *apierr.Error) {
	uuids, err := s.nodes.ActiveNodeUUIDs(ctx)
	if err != nil || len(uuids) == 0 {
		return nil, apierr.ErrNotFound
	}
	node, err := s.nodes.NodeByUUID(ctx, uuids[0])
	if err != nil || node == nil {
		return nil, apierr.ErrNotFound
	}
	return node, nil
}

func (s *SubscriptionAPI) urlFor(r *http.Request, token string) string {
	base := s.subBase
	if base == "" {
		scheme := "https"
		if r.TLS == nil && r.Header.Get("X-Forwarded-Proto") == "" {
			scheme = "http"
		}
		base = scheme + "://" + r.Host
	}
	return base + "/sub/" + token
}

// newSubToken returns a random UUIDv4 string (fits users.sub_token CHAR(36)).
func newSubToken() (string, error) {
	var b [16]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "", err
	}
	b[6] = (b[6] & 0x0f) | 0x40 // version 4
	b[8] = (b[8] & 0x3f) | 0x80 // variant 10
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]), nil
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

// statusFor maps the small set of apierr values used here to HTTP status codes.
func statusFor(e *apierr.Error) int {
	switch e {
	case apierr.ErrNotFound:
		return http.StatusNotFound
	case apierr.ErrUnauthorized:
		return http.StatusUnauthorized
	default:
		return http.StatusInternalServerError
	}
}