126 lines
2.6 KiB
Go
126 lines
2.6 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"3588AdminBackend/internal/config"
|
|
)
|
|
|
|
type AgentClient struct {
|
|
cfg *config.Config
|
|
client *http.Client
|
|
upload *http.Client
|
|
}
|
|
|
|
func NewAgentClient(cfg *config.Config) *AgentClient {
|
|
return &AgentClient{
|
|
cfg: cfg,
|
|
client: &http.Client{
|
|
Timeout: 3 * time.Second,
|
|
},
|
|
upload: &http.Client{Timeout: 120 * time.Second},
|
|
}
|
|
}
|
|
|
|
func (c *AgentClient) Do(method, ip string, port int, path string, body []byte) ([]byte, int, error) {
|
|
return c.DoStream(method, ip, port, path, bytes.NewReader(body), "", int64(len(body)))
|
|
}
|
|
|
|
func (c *AgentClient) DoStream(method, ip string, port int, path string, body io.Reader, contentType string, contentLength int64) ([]byte, int, error) {
|
|
url := fmt.Sprintf("http://%s:%d%s", ip, port, path)
|
|
var rc io.ReadCloser
|
|
if body == nil {
|
|
rc = http.NoBody
|
|
contentLength = 0
|
|
} else if r, ok := body.(io.ReadCloser); ok {
|
|
rc = r
|
|
} else {
|
|
rc = io.NopCloser(body)
|
|
}
|
|
|
|
ctx := context.Background()
|
|
req, err := http.NewRequestWithContext(ctx, method, url, rc)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
if contentLength >= 0 {
|
|
req.ContentLength = contentLength
|
|
}
|
|
|
|
if c.cfg.AgentToken != "" {
|
|
req.Header.Set("X-RK-Token", c.cfg.AgentToken)
|
|
}
|
|
|
|
if req.Body != nil && req.Body != http.NoBody {
|
|
ct := strings.TrimSpace(contentType)
|
|
if ct == "" {
|
|
ct = defaultContentTypeForPath(path)
|
|
}
|
|
if ct != "" {
|
|
req.Header.Set("Content-Type", ct)
|
|
}
|
|
}
|
|
|
|
client := c.client
|
|
if isLongOp(method, path) {
|
|
client = c.upload
|
|
}
|
|
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
respBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, resp.StatusCode, err
|
|
}
|
|
|
|
return respBody, resp.StatusCode, nil
|
|
}
|
|
|
|
func defaultContentTypeForPath(path string) string {
|
|
switch {
|
|
case strings.HasPrefix(path, "/v1/models/"):
|
|
return "application/octet-stream"
|
|
case strings.HasPrefix(path, "/v1/resources/"):
|
|
return "application/octet-stream"
|
|
case path == "/v1/face-gallery":
|
|
return "application/octet-stream"
|
|
default:
|
|
return "application/json"
|
|
}
|
|
}
|
|
|
|
func isLongOp(method, path string) bool {
|
|
if method == http.MethodPut && path == "/v1/config" {
|
|
return true
|
|
}
|
|
if method == http.MethodPost && path == "/v1/config/candidate/apply" {
|
|
return true
|
|
}
|
|
if method == http.MethodPost && path == "/v1/media-server/rollback" {
|
|
return true
|
|
}
|
|
if strings.HasPrefix(path, "/v1/config/ui/") {
|
|
return true
|
|
}
|
|
if strings.HasPrefix(path, "/v1/models/") {
|
|
return true
|
|
}
|
|
if strings.HasPrefix(path, "/v1/resources/") {
|
|
return true
|
|
}
|
|
if path == "/v1/face-gallery" {
|
|
return true
|
|
}
|
|
return false
|
|
}
|