XDRIVE transport: Support two-step "template" operations (hrefPath) (#7188)

https://github.com/XTLS/Xray-core/pull/7188#issuecomment-6250146420
This commit is contained in:
genkauzen
2026-10-21 10:55:36 +00:00
committed by RPRX
parent 3653dcb4fc
commit 00c0277787
2 changed files with 153 additions and 10 deletions
+43 -10
View File
@@ -33,6 +33,14 @@ type opTemplate struct {
Headers map[string]string `json:"headers"`
Body string `json:"body"`
NamesRegex string `json:"namesRegex"`
// HrefPath makes the operation two-step, as in the Yandex Disk REST API: the first request
// answers with a JSON body, the URL found at HrefPath is then called with HrefMethod (GET by
// default), carrying the payload. The URL is used as given, with neither auth nor headers.
HrefPath string `json:"hrefPath"`
HrefMethod string `json:"hrefMethod"`
presigned bool
}
type retryTemplate struct {
@@ -262,6 +270,11 @@ func (s *templateStorage) do(ctx context.Context, op *opTemplate, vars map[strin
backoff := driveInitialBackoff
var lastErr error
url := op.URL
if !op.presigned {
url = subst(url, vars)
}
for attempt := 0; attempt < driveMaxAttempts; attempt++ {
if attempt > 0 {
select {
@@ -275,10 +288,14 @@ func (s *templateStorage) do(ctx context.Context, op *opTemplate, vars map[strin
}
}
authHeaders, err := s.authHeaders(ctx, vars)
if err != nil {
lastErr = err
continue
var authHeaders map[string]string
if !op.presigned {
headers, err := s.authHeaders(ctx, vars)
if err != nil {
lastErr = err
continue
}
authHeaders = headers
}
method := op.Method
@@ -290,7 +307,7 @@ func (s *templateStorage) do(ctx context.Context, op *opTemplate, vars map[strin
if body != nil {
reader = bytes.NewReader(body)
}
req, err := http.NewRequestWithContext(ctx, method, subst(op.URL, vars), reader)
req, err := http.NewRequestWithContext(ctx, method, url, reader)
if err != nil {
return 0, nil, errors.New("failed to build request").Base(err)
}
@@ -325,7 +342,7 @@ func (s *templateStorage) do(ctx context.Context, op *opTemplate, vars map[strin
continue
}
if resp.StatusCode == http.StatusUnauthorized && s.tmpl.Auth.Type == "oauth2" {
if resp.StatusCode == http.StatusUnauthorized && s.tmpl.Auth.Type == "oauth2" && !op.presigned {
s.invalidateToken()
lastErr = errors.New("the service rejected the token")
continue
@@ -339,6 +356,22 @@ func (s *templateStorage) do(ctx context.Context, op *opTemplate, vars map[strin
return 0, nil, lastErr
}
// run performs op, following its href when the operation is two-step.
func (s *templateStorage) run(ctx context.Context, op *opTemplate, vars map[string]string, body []byte) (int, []byte, error) {
if op.HrefPath == "" {
return s.do(ctx, op, vars, body)
}
status, payload, err := s.do(ctx, op, vars, nil)
if err != nil || status < 200 || status >= 300 {
return status, payload, err
}
href := jsonString(payload, op.HrefPath)
if href == "" {
return 0, nil, errors.New("the response has no href at ", op.HrefPath)
}
return s.do(ctx, &opTemplate{Method: op.HrefMethod, URL: href, presigned: true}, vars, body)
}
func (s *templateStorage) Put(ctx context.Context, name string, data []byte) error {
vars := s.baseVars()
vars["name"] = s.storedName(name)
@@ -349,7 +382,7 @@ func (s *templateStorage) Put(ctx context.Context, name string, data []byte) err
body = []byte(subst(s.tmpl.Put.Body, vars))
}
status, payload, err := s.do(ctx, &s.tmpl.Put, vars, body)
status, payload, err := s.run(ctx, &s.tmpl.Put, vars, body)
if err != nil {
return err
}
@@ -363,7 +396,7 @@ func (s *templateStorage) Get(ctx context.Context, name string) ([]byte, error)
vars := s.baseVars()
vars["name"] = s.storedName(name)
status, payload, err := s.do(ctx, &s.tmpl.Get, vars, nil)
status, payload, err := s.run(ctx, &s.tmpl.Get, vars, nil)
if err != nil {
return nil, err
}
@@ -381,7 +414,7 @@ func (s *templateStorage) Delete(ctx context.Context, name string) error {
vars := s.baseVars()
vars["name"] = s.storedName(name)
status, payload, err := s.do(ctx, &s.tmpl.Delete, vars, nil)
status, payload, err := s.run(ctx, &s.tmpl.Delete, vars, nil)
if err != nil {
return err
}
@@ -396,7 +429,7 @@ func (s *templateStorage) List(ctx context.Context, prefix string) ([]Entry, err
flat := s.storedName(prefix)
vars["prefix"] = flat
status, payload, err := s.do(ctx, &s.tmpl.List, vars, nil)
status, payload, err := s.run(ctx, &s.tmpl.List, vars, nil)
if err != nil {
return nil, err
}
+110
View File
@@ -306,3 +306,113 @@ func TestTemplateConcurrencyDefault(t *testing.T) {
t.Fatalf("default inflight cap is %d, want %d", cap(storage.inflight), driveMaxInflight)
}
}
// hrefStore mimics the Yandex Disk REST API: upload and download answer with an href
// to a presigned blob URL, which must be called without the API credentials.
type hrefStore struct {
*fakeStore
blobAuth string
}
func newHrefStore(t *testing.T) *hrefStore {
t.Helper()
store := &hrefStore{fakeStore: &fakeStore{objects: make(map[string][]byte)}}
store.server = httptest.NewServer(http.HandlerFunc(store.handle))
t.Cleanup(store.server.Close)
return store
}
func (s *hrefStore) handle(w http.ResponseWriter, r *http.Request) {
if strings.HasPrefix(r.URL.Path, "/blob/") {
s.mu.Lock()
if auth := r.Header.Get("Authorization"); auth != "" {
s.blobAuth = auth
}
s.mu.Unlock()
r.URL.Path = "/folder/" + strings.TrimPrefix(r.URL.Path, "/blob/")
s.fakeStore.handle(w, r)
return
}
if r.Header.Get("Authorization") != "OAuth tok" {
w.WriteHeader(http.StatusUnauthorized)
return
}
name := r.URL.Query().Get("path")
switch r.URL.Path {
case "/upload":
case "/download":
s.mu.Lock()
_, ok := s.objects[name]
s.mu.Unlock()
if !ok {
w.WriteHeader(http.StatusNotFound)
return
}
default:
w.WriteHeader(http.StatusNotFound)
return
}
json.NewEncoder(w).Encode(map[string]interface{}{"href": s.server.URL + "/blob/" + name})
}
func hrefSettings(store *hrefStore) *internet.MemoryStreamConfig {
base := store.server.URL
settings := templateSettings(store.fakeStore, map[string]interface{}{
"type": "static", "header": map[string]interface{}{"Authorization": "OAuth {secret0}"},
}, []string{"tok"})
cfg := settings.ProtocolSettings.(*Config)
var tmpl map[string]interface{}
json.Unmarshal([]byte(cfg.Template), &tmpl)
tmpl["put"] = map[string]interface{}{"url": base + "/upload?path={name}", "hrefPath": "href", "hrefMethod": "PUT"}
tmpl["get"] = map[string]interface{}{"url": base + "/download?path={name}", "hrefPath": "href"}
tmpl["list"] = map[string]interface{}{"method": "PROPFIND", "url": base + "/blob/", "namesRegex": `<d:href>/folder/([^<]+)</d:href>`}
tmpl["delete"] = map[string]interface{}{"method": "DELETE", "url": base + "/blob/{name}"}
raw, _ := json.Marshal(tmpl)
cfg.Template = string(raw)
return settings
}
func TestTemplateTwoStep(t *testing.T) {
store := newHrefStore(t)
settings := hrefSettings(store)
storage, err := newTemplateStorage(settings, settings.ProtocolSettings.(*Config))
if err != nil {
t.Fatalf("newTemplateStorage: %v", err)
}
ctx := context.Background()
if err := storage.Put(ctx, "streams/abc/c2s/000000000.seg", []byte("hello")); err != nil {
t.Fatalf("Put: %v", err)
}
data, err := storage.Get(ctx, "streams/abc/c2s/000000000.seg")
if err != nil {
t.Fatalf("Get: %v", err)
}
if string(data) != "hello" {
t.Fatalf("Get returned %q, want hello", data)
}
if _, err := storage.Get(ctx, "streams/abc/c2s/000000009.seg"); err != errNotFound {
t.Fatalf("Get of a missing object returned %v, want errNotFound", err)
}
if store.blobAuth != "" {
t.Fatalf("the presigned href received credentials %q", store.blobAuth)
}
}
func TestTemplateTwoStepTransport(t *testing.T) {
store := newHrefStore(t)
client, server, cleanup := pairWith(t, hrefSettings(store))
defer cleanup()
if _, err := client.Write([]byte("ping")); err != nil {
t.Fatalf("client write: %v", err)
}
expectRead(t, server, "ping")
if _, err := server.Write([]byte("pong")); err != nil {
t.Fatalf("server write: %v", err)
}
expectRead(t, client, "pong")
}