story/internal/onboarding/captcha_policy.go

package onboarding

import (
	"context"
	"data-terminal-cloud/internal/audit"
	"encoding/json"
	"net"
	"net/http"
	"net/netip"
	"strings"
)

type captchaRequestKey struct{}
type captchaRequestInfo struct{ client, ip string }

func captchaRequest(ctx context.Context) captchaRequestInfo {
	if v, ok := ctx.Value(captchaRequestKey{}).(captchaRequestInfo); ok {
		return v
	}
	return captchaRequestInfo{client: "saas"}
}

// Walk the proxy chain from the trusted peer, never trust the leftmost value
// merely because a client supplied X-Forwarded-For.
func (s *Service) clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	ip, err := netip.ParseAddr(host)
	if err != nil {
		return ""
	}
	trusted := func(ip netip.Addr) bool {
		for _, cidr := range s.options.TrustedProxyCIDRs {
			prefix, e := netip.ParsePrefix(cidr)
			if e == nil && prefix.Contains(ip) {
				return true
			}
		}
		return false
	}
	if trusted(ip) {
		parts := strings.Split(r.Header.Get("X-Forwarded-For"), ",")
		for i := len(parts) - 1; i >= 0 && trusted(ip); i-- {
			next, e := netip.ParseAddr(strings.TrimSpace(parts[i]))
			if e != nil {
				break
			}
			ip = next
		}
	}
	return ip.String()
}

type captchaPolicy struct {
	SaaS    string `json:"saas_mode"`
	Mini    string `json:"mini_mode"`
	Version int64  `json:"version"`
}

func (p captchaPolicy) mode(client string) string {
	if client == "mini_program" {
		return p.Mini
	}
	return p.SaaS
}
func (s *Service) readCaptchaPolicy(ctx context.Context, q queryer) (captchaPolicy, error) {
	var p captchaPolicy
	err := q.QueryRow(ctx, `SELECT saas_mode,mini_mode,version FROM login_captcha_policy WHERE singleton`).Scan(&p.SaaS, &p.Mini, &p.Version)
	return p, err
}
func (s *Service) captchaAppID(client string) string {
	if s.options.Captcha == nil {
		return ""
	}
	return s.options.Captcha.AppID(client)
}
func (s *Service) captchaPolicyResult(p captchaPolicy) map[string]any {
	return map[string]any{"saas_mode": p.SaaS, "mini_mode": p.Mini, "version": p.Version, "saas_configured": s.captchaAppID("saas") != "", "mini_configured": s.captchaAppID("mini_program") != "", "saas_disabled": s.options.SaaSCaptchaDisabled}
}
func (s *Service) captchaPolicyHTTP(r *http.Request, requestID string) (any, error) {
	if r.URL.RawQuery != "" {
		return nil, errAccess
	}
	p, err := s.operatorView().browserPrincipal(r, r.Method != http.MethodGet)
	if err != nil {
		return nil, err
	}
	ctx := r.Context()
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback(ctx)
	a, err := operatorAccess(ctx, tx, p, r.Header.Get("X-Tenant-ID"), false)
	if err != nil {
		return nil, err
	}
	// This affects all login entry points, so ordinary operators cannot change it.
	if a.Role != "ops_admin" {
		return nil, errAccess
	}
	policy, err := s.readCaptchaPolicy(ctx, tx)
	if err != nil {
		return nil, err
	}
	if r.Method == http.MethodGet {
		return s.captchaPolicyResult(policy), nil
	}
	raw, err := readBody(r)
	if err != nil {
		return nil, err
	}
	var in captchaPolicy
	if err = StrictJSON(raw, &in); err != nil {
		return nil, err
	}
	if !oneOf(in.SaaS, "image", "slide", "tencent") || !oneOf(in.Mini, "image", "tencent") || in.Version < 1 {
		return nil, fail(400, "VALIDATION_FAILED", "请选择有效的验证码方式")
	}
	if (in.SaaS == "tencent" && s.captchaAppID("saas") == "") || (in.Mini == "tencent" && s.captchaAppID("mini_program") == "") {
		return nil, fail(409, "CAPTCHA_NOT_CONFIGURED", "请先在服务器填写对应端的腾讯云验证码配置并重启服务")
	}
	result, err := tx.Exec(ctx, `UPDATE login_captcha_policy SET saas_mode=$1,mini_mode=$2,version=version+1,updated_at=now() WHERE singleton AND version=$3`, in.SaaS, in.Mini, in.Version)
	if err != nil {
		return nil, err
	}
	if result.RowsAffected() != 1 {
		return nil, fail(409, "CAPTCHA_POLICY_CONFLICT", "配置已被更新，请刷新后重试")
	}
	after := in
	after.Version++
	err = s.audit.Record(ctx, tx, audit.Event{TenantID: &a.TenantID, ActorType: "user", ActorPublicID: &a.PublicID, Action: "login.captcha.policy.update", ResourceType: "login_captcha_policy", RequestID: requestID, Result: "success", Detail: map[string]any{"before": policy, "after": after}})
	if err != nil {
		return nil, err
	}
	if err = tx.Commit(ctx); err != nil {
		return nil, err
	}
	in.Version++
	return s.captchaPolicyResult(in), nil
}

type storedCaptcha struct {
	ID         string `json:"id"`
	Mode       string `json:"mode"`
	Client     string `json:"client"`
	Version    int64  `json:"version"`
	AppID      string `json:"app_id,omitempty"`
	AnswerHash string `json:"answer_hash,omitempty"`
	Solution   []byte `json:"slide_solution,omitempty"`
	ProofHash  string `json:"slide_proof_hash,omitempty"`
	ExpiresAt  int64  `json:"expires_at_ms,omitempty"`
}

func captchaJSON(v storedCaptcha) string { raw, _ := json.Marshal(v); return string(raw) }