126 lines
4.2 KiB
Go
126 lines
4.2 KiB
Go
package providers
|
|
|
|
import (
|
|
"context"
|
|
"crypto/hmac"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type Volcengine struct {
|
|
config Config
|
|
client HTTPClient
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewVolcengine(c Config, client HTTPClient, now func() time.Time) *Volcengine {
|
|
if client == nil {
|
|
client = http.DefaultClient
|
|
}
|
|
if now == nil {
|
|
now = time.Now
|
|
}
|
|
if c.Region == "" {
|
|
c.Region = "cn-north-1"
|
|
}
|
|
if c.Service == "" {
|
|
c.Service = "cv"
|
|
}
|
|
return &Volcengine{c, client, now}
|
|
}
|
|
func (v *Volcengine) Submit(ctx context.Context, r Request) (Result, error) {
|
|
p := map[string]any{"req_key": requestModel(r, v.config.Model), "prompt": r.Prompt, "image_urls": r.InputURLs}
|
|
for _, key := range []string{"scale", "width", "height", "min_ratio", "max_ratio", "force_single"} {
|
|
if value, exists := r.Settings[key]; exists && value != nil && value != "" {
|
|
p[key] = value
|
|
}
|
|
}
|
|
return v.call(ctx, "CVSync2AsyncSubmitTask", p)
|
|
}
|
|
func (v *Volcengine) Query(ctx context.Context, id string) (Result, error) {
|
|
return v.QueryModel(ctx, id, v.config.Model)
|
|
}
|
|
func (v *Volcengine) QueryModel(ctx context.Context, id, model string) (Result, error) {
|
|
queryOptions, _ := json.Marshal(map[string]any{
|
|
"return_url": true,
|
|
"logo_info": map[string]any{"add_logo": false, "position": 0, "language": 0, "opacity": 1},
|
|
})
|
|
return v.call(ctx, "CVSync2AsyncGetResult", map[string]any{"req_key": requestModel(Request{Model: model}, v.config.Model), "task_id": id, "req_json": string(queryOptions)})
|
|
}
|
|
func (v *Volcengine) call(ctx context.Context, action string, payload any) (Result, error) {
|
|
body, _ := json.Marshal(payload)
|
|
endpoint, err := url.Parse(v.config.BaseURL)
|
|
if err != nil {
|
|
return Result{}, errors.New("invalid provider base URL")
|
|
}
|
|
q := endpoint.Query()
|
|
q.Set("Action", action)
|
|
q.Set("Version", "2022-08-31")
|
|
endpoint.RawQuery = canonicalQuery(q)
|
|
date := v.now().UTC()
|
|
xdate := date.Format("20060102T150405Z")
|
|
short := xdate[:8]
|
|
hash := sha(body)
|
|
headers := "content-type:application/json\nhost:" + endpoint.Host + "\nx-content-sha256:" + hash + "\nx-date:" + xdate + "\n"
|
|
signed := "content-type;host;x-content-sha256;x-date"
|
|
canonicalPath := endpoint.EscapedPath()
|
|
if canonicalPath == "" {
|
|
canonicalPath = "/"
|
|
}
|
|
canonical := "POST\n" + canonicalPath + "\n" + endpoint.RawQuery + "\n" + headers + "\n" + signed + "\n" + hash
|
|
scope := short + "/" + v.config.Region + "/" + v.config.Service + "/request"
|
|
stringToSign := "HMAC-SHA256\n" + xdate + "\n" + scope + "\n" + sha([]byte(canonical))
|
|
key := hmacBytes([]byte(v.config.SecretAccessKey), short)
|
|
key = hmacBytes(key, v.config.Region)
|
|
key = hmacBytes(key, v.config.Service)
|
|
key = hmacBytes(key, "request")
|
|
signature := hex.EncodeToString(hmacBytes(key, stringToSign))
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, endpoint.String(), strings.NewReader(string(body)))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-Content-Sha256", hash)
|
|
req.Header.Set("X-Date", xdate)
|
|
req.Header.Set("Authorization", "HMAC-SHA256 Credential="+v.config.AccessKeyID+"/"+scope+", SignedHeaders="+signed+", Signature="+signature)
|
|
resp, err := v.client.Do(req)
|
|
if err != nil {
|
|
return Result{}, &ProviderError{Operation: "volcengine request"}
|
|
}
|
|
defer resp.Body.Close()
|
|
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 || !json.Valid(raw) {
|
|
return Result{}, &ProviderError{Operation: "volcengine request", Status: resp.StatusCode}
|
|
}
|
|
r := record(raw)
|
|
d := object(r["data"])
|
|
out := []string{}
|
|
collectURLs(d, &out)
|
|
return Result{TaskID: stringValue(r["task_id"], d["task_id"]), Status: status(first(r["status"], d["status"])), OutputURLs: out, Raw: raw}, nil
|
|
}
|
|
func sha(b []byte) string { x := sha256.Sum256(b); return hex.EncodeToString(x[:]) }
|
|
func hmacBytes(k []byte, s string) []byte {
|
|
h := hmac.New(sha256.New, k)
|
|
_, _ = h.Write([]byte(s))
|
|
return h.Sum(nil)
|
|
}
|
|
func canonicalQuery(q url.Values) string {
|
|
keys := make([]string, 0, len(q))
|
|
for k := range q {
|
|
keys = append(keys, k)
|
|
}
|
|
sort.Strings(keys)
|
|
parts := []string{}
|
|
for _, k := range keys {
|
|
for _, v := range q[k] {
|
|
parts = append(parts, url.QueryEscape(k)+"="+url.QueryEscape(v))
|
|
}
|
|
}
|
|
return strings.Join(parts, "&")
|
|
}
|