ai_site/gateway/internal/proxy/director.go

package proxy

import (
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"
)

type Director struct {
	Platform *url.URL
	AI       *url.URL
	Timeout  time.Duration
}

func New(platformURL, aiURL string, timeoutSec int) (*Director, error) {
	p, err := url.Parse(platformURL)
	if err != nil {
		return nil, err
	}
	a, err := url.Parse(aiURL)
	if err != nil {
		return nil, err
	}
	if timeoutSec <= 0 {
		timeoutSec = 60
	}
	return &Director{Platform: p, AI: a, Timeout: time.Duration(timeoutSec) * time.Second}, nil
}

func (d *Director) Handler() http.HandlerFunc {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   10 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}

	return func(w http.ResponseWriter, r *http.Request) {
		target, path := d.route(r.URL.Path)
		proxy := httputil.NewSingleHostReverseProxy(target)
		proxy.Transport = transport
		proxy.FlushInterval = 100 * time.Millisecond
		proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {
			rw.Header().Set("Content-Type", "application/json")
			rw.WriteHeader(http.StatusBadGateway)
			_, _ = rw.Write([]byte(`{"code":502,"message":"upstream unavailable: ` + escape(err.Error()) + `"}`))
		}
		proxy.ModifyResponse = func(resp *http.Response) error {
			// 网关已统一加 CORS；去掉上游重复头，避免浏览器报 Failed to fetch
			for _, h := range []string{
				"Access-Control-Allow-Origin",
				"Access-Control-Allow-Credentials",
				"Access-Control-Allow-Headers",
				"Access-Control-Allow-Methods",
				"Access-Control-Expose-Headers",
				"Access-Control-Max-Age",
			} {
				resp.Header.Del(h)
			}
			resp.Header.Set("X-Gateway", "aijianzhan-gateway")
			return nil
		}

		// 重写路径
		r.URL.Path = path
		r.Host = target.Host
		r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))
		r.Header.Set("X-Forwarded-Proto", "http")
		if r.Header.Get("X-Request-Id") == "" {
			r.Header.Set("X-Request-Id", newRequestID())
		}

		// 超时上下文
		ctx := r.Context()
		proxy.ServeHTTP(w, r.WithContext(ctx))
	}
}

func (d *Director) route(path string) (*url.URL, string) {
	// AI：/ai/* → 上游去掉 /ai 前缀
	if strings.HasPrefix(path, "/ai/") || path == "/ai" {
		next := strings.TrimPrefix(path, "/ai")
		if next == "" {
			next = "/"
		}
		return d.AI, next
	}
	// 生成蓝图：统一 /api/v1/apps/generate（兼容旧冒号路径，避免重复实现）
	if path == "/api/v1/apps/generate" || strings.HasPrefix(path, "/api/v1/apps/generate?") {
		return d.AI, path
	}
	if strings.HasPrefix(path, "/api/v1/apps:generate") || strings.Contains(path, "/apps%3Agenerate") {
		return d.AI, "/api/v1/apps/generate"
	}
	// 其余 /api 走中台
	return d.Platform, path
}

func escape(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `"`, `\"`)
	if len(s) > 200 {
		s = s[:200]
	}
	return s
}

func newRequestID() string {
	return strings.ReplaceAll(time.Now().UTC().Format("20060102T150405.000000000"), ".", "")
}

// Health 聚合探测
func (d *Director) Health(w http.ResponseWriter, _ *http.Request) {
	type st struct {
		Name   string `json:"name"`
		OK     bool   `json:"ok"`
		Detail string `json:"detail,omitempty"`
	}
	client := &http.Client{Timeout: 2 * time.Second}
	out := []st{
		probe(client, "platform", d.Platform.String()+"/api/v1/auth/login"),
		probe(client, "ai", d.AI.String()+"/health"),
	}
	allOK := true
	for _, s := range out {
		if !s.OK {
			allOK = false
			break
		}
	}
	code := http.StatusOK
	if !allOK {
		code = http.StatusServiceUnavailable
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_, _ = io.WriteString(w, `{"gateway":"ok","upstreams":[`)
	for i, s := range out {
		if i > 0 {
			_, _ = io.WriteString(w, ",")
		}
		ok := "false"
		if s.OK {
			ok = "true"
		}
		_, _ = io.WriteString(w, `{"name":"`+s.Name+`","ok":`+ok+`}`)
	}
	_, _ = io.WriteString(w, `]}`)
}

func probe(client *http.Client, name, rawURL string) struct {
	Name   string `json:"name"`
	OK     bool   `json:"ok"`
	Detail string `json:"detail,omitempty"`
} {
	// login 用 OPTIONS/GET 可能 405；用短超时 HEAD/GET health 风格
	req, _ := http.NewRequest(http.MethodGet, rawURL, nil)
	if name == "platform" {
		// 未登录会 401/405 都说明服务活着；连接失败才算挂
		req, _ = http.NewRequest(http.MethodPost, rawURL, strings.NewReader(`{}`))
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := client.Do(req)
	if err != nil {
		return struct {
			Name   string `json:"name"`
			OK     bool   `json:"ok"`
			Detail string `json:"detail,omitempty"`
		}{Name: name, OK: false, Detail: err.Error()}
	}
	defer resp.Body.Close()
	return struct {
		Name   string `json:"name"`
		OK     bool   `json:"ok"`
		Detail string `json:"detail,omitempty"`
	}{Name: name, OK: true}
}