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

transport/internet/xdrive/template.go

@@ -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
 	}

transport/internet/xdrive/template_test.go

@@ -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")
+}