dobotai-platform/internal/httpapi/server.go

package httpapi

import (
	"context"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"

	"gitea.zheromiku.top/Weknora/dobotai-platform/internal/audit"
	"gitea.zheromiku.top/Weknora/dobotai-platform/internal/core"
)

type QAService interface {
	Ask(context.Context, core.QARequest) (core.QAResult, error)
}

type AuditService interface {
	ListTurns(context.Context, audit.ListTurnsRequest) (audit.ListTurnsResult, error)
	GetTurn(context.Context, string) (audit.TurnDetail, error)
}

type BuildInfo struct {
	Version   string `json:"version"`
	CommitSHA string `json:"commit_sha"`
	BuildTime string `json:"build_time"`
}

type Options struct {
	AuditService AuditService
	AdminToken   string
	BuildInfo    BuildInfo
}

type Server struct {
	service      QAService
	ingressToken string
	auditService AuditService
	adminToken   string
	buildInfo    BuildInfo
	mux          *http.ServeMux
}

func New(service QAService, ingressToken string) (*Server, error) {
	return NewWithOptions(service, ingressToken, Options{})
}

func NewWithOptions(service QAService, ingressToken string, opts Options) (*Server, error) {
	if service == nil {
		return nil, errors.New("QA service is required")
	}
	if strings.TrimSpace(ingressToken) == "" {
		return nil, errors.New("DOBOTAI_INGRESS_TOKEN is required")
	}
	adminToken := strings.TrimSpace(opts.AdminToken)
	hasAuditService := opts.AuditService != nil
	hasAdminToken := adminToken != ""
	if hasAuditService != hasAdminToken {
		return nil, errors.New("audit service and DOBOTAI_ADMIN_TOKEN must be configured together")
	}
	if hasAdminToken && constantTimeTokenEqual(adminToken, ingressToken) {
		return nil, errors.New("DOBOTAI_ADMIN_TOKEN must differ from DOBOTAI_INGRESS_TOKEN")
	}
	buildInfo := opts.BuildInfo
	if buildInfo.Version == "" {
		buildInfo.Version = "dev"
	}
	if buildInfo.CommitSHA == "" {
		buildInfo.CommitSHA = "unknown"
	}
	if buildInfo.BuildTime == "" {
		buildInfo.BuildTime = "unknown"
	}

	s := &Server{
		service: service, ingressToken: ingressToken,
		auditService: opts.AuditService, adminToken: adminToken,
		buildInfo: buildInfo,
		mux:       http.NewServeMux(),
	}
	s.mux.HandleFunc("GET /healthz", s.health)
	s.mux.HandleFunc("GET /version", s.version)
	s.mux.HandleFunc("POST /v1/qa", s.authorizeToken(s.ingressToken, s.ask))
	if hasAuditService {
		s.mux.HandleFunc("GET /v1/admin/audit/turns", s.authorizeToken(s.adminToken, s.listAuditTurns))
		s.mux.HandleFunc("GET /v1/admin/audit/turns/{id}", s.authorizeToken(s.adminToken, s.getAuditTurn))
	}
	return s, nil
}

func (s *Server) Handler() http.Handler { return s.mux }

func (s *Server) authorizeToken(expected string, next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/v1/admin/audit/") {
			w.Header().Set("Cache-Control", "no-store")
		}
		header := r.Header.Get("Authorization")
		if !strings.HasPrefix(header, "Bearer ") {
			writeError(w, http.StatusUnauthorized, "unauthorized")
			return
		}
		provided := strings.TrimPrefix(header, "Bearer ")
		if provided == "" || strings.ContainsAny(provided, " \t\r\n") || !constantTimeTokenEqual(provided, expected) {
			writeError(w, http.StatusUnauthorized, "unauthorized")
			return
		}
		next(w, r)
	}
}

func constantTimeTokenEqual(provided, expected string) bool {
	return len(provided) == len(expected) && subtle.ConstantTimeCompare([]byte(provided), []byte(expected)) == 1
}

func (s *Server) health(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

func (s *Server) version(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	writeJSON(w, http.StatusOK, s.buildInfo)
}

func (s *Server) ask(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, 1<<20)
	defer r.Body.Close()
	var request struct {
		TenantKey         string                `json:"tenant_key"`
		Channel           string                `json:"channel"`
		ConversationType  core.ConversationType `json:"conversation_type"`
		ExternalChatID    string                `json:"external_chat_id"`
		ExternalUserID    string                `json:"external_user_id"`
		ExternalMessageID string                `json:"external_message_id"`
		SenderName        string                `json:"sender_name"`
		MessageType       string                `json:"message_type"`
		Query             string                `json:"query"`
	}
	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
		if errors.Is(err, io.EOF) {
			writeError(w, http.StatusBadRequest, "request body is required")
		} else {
			writeError(w, http.StatusBadRequest, "invalid JSON request")
		}
		return
	}
	result, err := s.service.Ask(r.Context(), core.QARequest{TenantKey: request.TenantKey, Channel: request.Channel, ConversationType: request.ConversationType, ExternalChatID: request.ExternalChatID, ExternalUserID: request.ExternalUserID, ExternalMessageID: request.ExternalMessageID, SenderName: request.SenderName, MessageType: request.MessageType, Query: request.Query})
	if err != nil {
		if errors.Is(err, core.ErrTurnInProgress) {
			writeError(w, http.StatusConflict, "message is already processing")
			return
		}
		writeError(w, http.StatusBadGateway, "question processing failed")
		return
	}
	writeJSON(w, http.StatusOK, result)
}

func (s *Server) listAuditTurns(w http.ResponseWriter, r *http.Request) {
	allowed := map[string]bool{"limit": true, "cursor": true, "channel": true, "status": true}
	for key, values := range r.URL.Query() {
		if !allowed[key] || len(values) != 1 {
			writeError(w, http.StatusBadRequest, "invalid audit query")
			return
		}
	}
	var request audit.ListTurnsRequest
	if raw, ok := singleQueryValue(r, "limit"); ok {
		limit, err := strconv.Atoi(raw)
		if err != nil || limit < 1 || limit > 100 {
			writeError(w, http.StatusBadRequest, "invalid audit query")
			return
		}
		request.Limit = limit
	}
	request.Cursor = r.URL.Query().Get("cursor")
	request.Channel = r.URL.Query().Get("channel")
	request.Status = audit.TurnStatus(r.URL.Query().Get("status"))
	if len(request.Cursor) > 512 || len(request.Channel) > 128 || !validAuditStatus(request.Status) {
		writeError(w, http.StatusBadRequest, "invalid audit query")
		return
	}
	result, err := s.auditService.ListTurns(r.Context(), request)
	if err != nil {
		writeAuditError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, result)
}

func (s *Server) getAuditTurn(w http.ResponseWriter, r *http.Request) {
	if len(r.URL.Query()) != 0 {
		writeError(w, http.StatusBadRequest, "invalid audit query")
		return
	}
	result, err := s.auditService.GetTurn(r.Context(), r.PathValue("id"))
	if err != nil {
		writeAuditError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, result)
}

func singleQueryValue(r *http.Request, key string) (string, bool) {
	values, ok := r.URL.Query()[key]
	if !ok || len(values) != 1 {
		return "", false
	}
	return values[0], true
}

func validAuditStatus(status audit.TurnStatus) bool {
	switch status {
	case "", audit.StatusProcessing, audit.StatusSucceeded, audit.StatusFailed, audit.StatusCancelled:
		return true
	default:
		return false
	}
}

func writeAuditError(w http.ResponseWriter, err error) {
	switch {
	case errors.Is(err, audit.ErrInvalidArgument):
		writeError(w, http.StatusBadRequest, "invalid audit query")
	case errors.Is(err, audit.ErrNotFound):
		writeError(w, http.StatusNotFound, "audit turn not found")
	default:
		writeError(w, http.StatusInternalServerError, "audit query failed")
	}
}

func ValidateLoopbackListenAddress(address string) error {
	host, port, err := net.SplitHostPort(strings.TrimSpace(address))
	if err != nil || host == "" || port == "" {
		return errors.New("audit API requires a loopback host and explicit port")
	}
	portNumber, err := strconv.Atoi(port)
	if err != nil || portNumber < 1 || portNumber > 65535 {
		return errors.New("audit API requires a valid port")
	}
	if strings.EqualFold(host, "localhost") {
		return nil
	}
	ip := net.ParseIP(host)
	if ip == nil || !ip.IsLoopback() {
		return fmt.Errorf("audit API listen host %q is not loopback", host)
	}
	return nil
}

func writeError(w http.ResponseWriter, status int, message string) {
	writeJSON(w, status, map[string]string{"error": message})
}

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