mirror of
https://github.com/mudler/LocalAI.git
synced 2026-11-06 04:23:15 -05:00
feat(addons): move connector file bytes by handle
A connector add-on that declares files.transfer no longer carries file bytes in the connector envelope. The host puts a random handle in the request of the op media (write) or send_file (read), and the add-on streams the file with two new host operations, connector.file.put and connector.file.get, in chunks of at most 256 KiB. The process that runs the add-on registers the handle when the call arrives and ends it when the call returns. A handle is bound to the add-on, so another add-on gets the same answer as for an unknown handle. An unfinished inbound file is dropped with its partial bytes, and a handle with no call for two minutes expires. Inbound bytes go to the blob store of the workspace, outbound bytes come from it, so a cluster needs the shared object store and no shared disk. The file is never held in memory. The host overwrites the size and hash of an inbound file with its own record. The change is additive: the field files.transfer needs manifest API 2, a kind without it keeps the inline path and its 1 MiB cap, and the host never sends it a handle. A call that carries a handle may run up to 15 minutes. The ceiling of a file also stays under the largest file of a workspace. Assisted-by: Claude:claude-sonnet-5-5 go test
This commit is contained in:
21 files changed
+1453
-14
No files matched your search
@@ -13,6 +13,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/addons"
|
||||
"github.com/mudler/LocalAI/core/services/addons/hostops"
|
||||
"github.com/mudler/LocalAI/core/services/studio/connectors"
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
)
|
||||
@@ -51,7 +52,7 @@ func (g *Gateway) connectorSource() ConnectorSource {
|
||||
func (g *Gateway) ConnectorKinds(ctx context.Context) ([]connectors.Kind, error) {
|
||||
if g.Forward != nil {
|
||||
var out []connectors.Kind
|
||||
err := g.Forward.connectorRequest(ctx, "GET", "/addons/connectors/kinds", nil, &out)
|
||||
err := g.Forward.connectorRequest(ctx, "GET", "/addons/connectors/kinds", nil, &out, connectorTimeout)
|
||||
return out, err
|
||||
}
|
||||
return g.localConnectorKinds(ctx)
|
||||
@@ -77,13 +78,17 @@ func (g *Gateway) ConnectorCall(ctx context.Context, s connectors.ConnectorCallS
|
||||
req.Token = ""
|
||||
}
|
||||
var out contract.ConnectorResponse
|
||||
err := g.Forward.connectorRequest(ctx, "POST", "/addons/connectors/call", connectorCall{s, req}, &out)
|
||||
err := g.Forward.connectorRequest(ctx, "POST", "/addons/connectors/call", connectorCall{s, req}, &out, req.CallTimeout())
|
||||
return out, err
|
||||
}
|
||||
return g.localConnectorCall(ctx, s, req)
|
||||
}
|
||||
func (g *Gateway) localConnectorCall(ctx context.Context, s connectors.ConnectorCallScope, req contract.ConnectorRequest) (contract.ConnectorResponse, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, connectorTimeout)
|
||||
timeout := connectorTimeout
|
||||
if req.Transfer != nil {
|
||||
timeout = req.CallTimeout()
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
source := g.connectorSource()
|
||||
if source == nil || g.ConnectorAuthority == nil {
|
||||
@@ -109,7 +114,18 @@ func (g *Gateway) localConnectorCall(ctx context.Context, s connectors.Connector
|
||||
}
|
||||
return contract.ConnectorResponse{}, ErrForbidden
|
||||
}
|
||||
if req.Transfer != nil && (kinds[i].Files == nil || !kinds[i].Files.Transfer) {
|
||||
// A handle only goes to an add-on that declared that it takes one.
|
||||
return contract.ConnectorResponse{}, ErrForbidden
|
||||
}
|
||||
end, err := g.beginTransfer(s, req)
|
||||
if err != nil {
|
||||
return contract.ConnectorResponse{}, err
|
||||
}
|
||||
out, err := source.Connector(ctx, s.AddonID, req)
|
||||
if end != nil {
|
||||
out = finishTransfer(end, out)
|
||||
}
|
||||
if err != nil {
|
||||
// A definite refusal of the add-on keeps its type; anything else could
|
||||
// not tell, so the caller treats it as an outage, with the cause kept.
|
||||
@@ -120,8 +136,8 @@ func (g *Gateway) localConnectorCall(ctx context.Context, s connectors.Connector
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (c *Carrier) connectorRequest(ctx context.Context, method, path string, body any, out any) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, connectorTimeout)
|
||||
func (c *Carrier) connectorRequest(ctx context.Context, method, path string, body any, out any, timeout time.Duration) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
var raw []byte
|
||||
if body != nil {
|
||||
@@ -136,7 +152,13 @@ func (c *Carrier) connectorRequest(ctx context.Context, method, path string, bod
|
||||
return ErrUnavailable
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := c.client.Do(req)
|
||||
// A call that moves a file runs longer than the client's own timeout; the
|
||||
// context carries the bound.
|
||||
client := c.client
|
||||
if timeout > connectorTimeout {
|
||||
client = c.long
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return ErrUnavailable
|
||||
}
|
||||
@@ -172,6 +194,10 @@ func (g *Gateway) serveConnectors(w http.ResponseWriter, r *http.Request) bool {
|
||||
typedError(w, ErrForbidden)
|
||||
return true
|
||||
}
|
||||
if call.Request.Transfer != nil {
|
||||
// The server's write timeout is sized for short calls.
|
||||
_ = http.NewResponseController(w).SetWriteDeadline(time.Now().Add(contract.ConnectorTransferTimeout + time.Minute))
|
||||
}
|
||||
out, err := g.localConnectorCall(r.Context(), call.Scope, call.Request)
|
||||
if err != nil {
|
||||
typedError(w, err)
|
||||
@@ -193,3 +219,51 @@ func decodeConnector(r io.Reader, out any) error {
|
||||
}
|
||||
return decode(bytes.NewReader(raw), out)
|
||||
}
|
||||
|
||||
// beginTransfer registers the handle of a call that carries one, on the
|
||||
// process that runs the add-on. The frontend made the handle and checked the
|
||||
// chat's access; this process binds it to the add-on of the call, so no other
|
||||
// add-on can use it, and to the content or the bound the scope names.
|
||||
func (g *Gateway) beginTransfer(s connectors.ConnectorCallScope, req contract.ConnectorRequest) (func() hostops.TransferOutcome, error) {
|
||||
if req.Transfer == nil {
|
||||
return nil, nil
|
||||
}
|
||||
grant := hostops.TransferGrant{Handle: req.Transfer.Handle, AddonID: s.AddonID, BindingID: s.BindingID}
|
||||
switch req.Op {
|
||||
case contract.ConnectorOpMedia:
|
||||
grant.Direction, grant.MaxBytes = hostops.TransferIn, req.Transfer.MaxBytes
|
||||
case contract.ConnectorOpSendFile:
|
||||
if s.Content == nil || req.File != nil {
|
||||
return nil, ErrForbidden
|
||||
}
|
||||
grant.Direction, grant.SHA, grant.Size = hostops.TransferOut, s.Content.SHA, s.Content.Size
|
||||
default:
|
||||
return nil, ErrForbidden
|
||||
}
|
||||
if g.Transfers == nil {
|
||||
return nil, ErrUnavailable
|
||||
}
|
||||
end, err := g.Transfers.Begin(grant)
|
||||
if err != nil {
|
||||
if errors.Is(err, hostops.ErrBadRequest) || errors.Is(err, hostops.ErrConflict) {
|
||||
return nil, ErrForbidden
|
||||
}
|
||||
return nil, ErrUnavailable
|
||||
}
|
||||
return end, nil
|
||||
}
|
||||
|
||||
// finishTransfer ends the handle after the add-on answered. For an inbound
|
||||
// file it writes the size and the hash from the host's own record into the
|
||||
// answer; an add-on that did not finish the file gets no Received.
|
||||
func finishTransfer(end func() hostops.TransferOutcome, out contract.ConnectorResponse) contract.ConnectorResponse {
|
||||
o := end()
|
||||
if out.Received != nil {
|
||||
if o.Direction != hostops.TransferIn || o.Result != hostops.TransferDone || o.SHA == "" {
|
||||
out.Received = nil
|
||||
} else {
|
||||
out.Received.Size, out.Received.SHA256 = o.Bytes, o.SHA
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -69,7 +69,7 @@ var _ = Describe("Forwarded connector answers", func() {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(status) }))
|
||||
defer server.Close()
|
||||
c := &Carrier{config: CarrierConfig{OwnerAddress: server.URL}, client: server.Client()}
|
||||
return c.connectorRequest(context.Background(), "GET", "/addons/connectors/kinds", nil, &[]connectors.Kind{})
|
||||
return c.connectorRequest(context.Background(), "GET", "/addons/connectors/kinds", nil, &[]connectors.Kind{}, connectorTimeout)
|
||||
}
|
||||
|
||||
It("keeps a refusal apart from an outage", func() {
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/addons"
|
||||
"github.com/mudler/LocalAI/core/services/addons/hostops"
|
||||
"github.com/mudler/LocalAI/core/services/studio/connectors"
|
||||
"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// handleSource is an add-on that takes handles: it runs fn during the call.
|
||||
type handleSource struct {
|
||||
transfer bool
|
||||
fn func(contract.ConnectorRequest) contract.ConnectorResponse
|
||||
}
|
||||
|
||||
func (s handleSource) ConnectorKinds(context.Context) ([]addons.ConnectorKind, error) {
|
||||
k := contract.ConnectorKind{ID: "slack", Delivery: contract.DeliveryPull}
|
||||
if s.transfer {
|
||||
k.Files = &contract.ConnectorFiles{Receive: &contract.FileAbility{MaxBytes: 1 << 30}, Transfer: true}
|
||||
}
|
||||
return []addons.ConnectorKind{{AddonID: "slack", ConnectorKind: k}}, nil
|
||||
}
|
||||
|
||||
func (s handleSource) Connector(_ context.Context, _ string, r contract.ConnectorRequest) (contract.ConnectorResponse, error) {
|
||||
return s.fn(r), nil
|
||||
}
|
||||
|
||||
var _ = Describe("Connector calls that carry a file handle", func() {
|
||||
scope := connectors.ConnectorCallScope{AddonID: "slack", Kind: "slack", BindingID: "b1"}
|
||||
var (
|
||||
t *hostops.Transfers
|
||||
host http.Handler
|
||||
)
|
||||
BeforeEach(func() {
|
||||
blobs, err := blob.NewFS(GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
t = hostops.NewTransfers(blobs)
|
||||
b := hostops.New(hostops.NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits())
|
||||
Expect(hostops.RegisterConnectorTransfer(b, t)).To(Succeed())
|
||||
host = b.Handler("slack", "k", []string{"connector."})
|
||||
})
|
||||
put := func(h string, off int64, data []byte, final bool) int {
|
||||
raw, _ := json.Marshal(contract.ConnectorFilePut{Handle: h, Offset: off, Data: data, Final: final})
|
||||
body, _ := json.Marshal(contract.HostRequest{OperationID: "o", Op: contract.ConnectorFilePutOp, Arguments: raw})
|
||||
req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body))
|
||||
req.Header.Set(contract.HostCredentialHeader, "k")
|
||||
rec := httptest.NewRecorder()
|
||||
host.ServeHTTP(rec, req)
|
||||
return rec.Code
|
||||
}
|
||||
media := func(s handleSource, tr *hostops.Transfers) (contract.ConnectorResponse, error) {
|
||||
g := &Gateway{Connectors: s, ConnectorAuthority: allowAllConnectors{}, Transfers: tr}
|
||||
return g.ConnectorCall(context.Background(), scope, contract.ConnectorRequest{Op: contract.ConnectorOpMedia, Kind: "slack", Token: "t",
|
||||
Transfer: &contract.ConnectorTransfer{Handle: "hh", MaxBytes: 100}})
|
||||
}
|
||||
|
||||
It("registers the handle for the call and gives the host's size and hash, not the add-on's", func() {
|
||||
src := handleSource{transfer: true, fn: func(r contract.ConnectorRequest) contract.ConnectorResponse {
|
||||
Expect(t.Live()).To(Equal(1))
|
||||
Expect(put(r.Transfer.Handle, 0, []byte("hello"), true)).To(Equal(http.StatusOK))
|
||||
return contract.ConnectorResponse{Received: &contract.ConnectorReceived{Name: "a.txt", Size: 999, SHA256: "forged"}}
|
||||
}}
|
||||
out, err := media(src, t)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
sum := sha256.Sum256([]byte("hello"))
|
||||
Expect(out.Received.SHA256).To(Equal(hex.EncodeToString(sum[:])))
|
||||
Expect(out.Received.Size).To(BeEquivalentTo(5))
|
||||
Expect(t.Live()).To(BeZero(), "the handle ends with the call")
|
||||
})
|
||||
|
||||
It("drops an answer for a file the add-on did not finish", func() {
|
||||
src := handleSource{transfer: true, fn: func(r contract.ConnectorRequest) contract.ConnectorResponse {
|
||||
Expect(put(r.Transfer.Handle, 0, []byte("half"), false)).To(Equal(http.StatusOK))
|
||||
return contract.ConnectorResponse{Received: &contract.ConnectorReceived{Name: "a.txt", Size: 4, SHA256: "forged"}}
|
||||
}}
|
||||
out, err := media(src, t)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(out.Received).To(BeNil())
|
||||
Expect(t.Live()).To(BeZero())
|
||||
})
|
||||
|
||||
It("refuses a handle for a kind that did not declare transfer, and without a table", func() {
|
||||
src := handleSource{fn: func(contract.ConnectorRequest) contract.ConnectorResponse { return contract.ConnectorResponse{} }}
|
||||
_, err := media(src, t)
|
||||
Expect(err).To(MatchError(ErrForbidden))
|
||||
src.transfer = true
|
||||
_, err = media(src, nil)
|
||||
Expect(err).To(MatchError(ErrUnavailable))
|
||||
})
|
||||
|
||||
It("refuses a handle on an op that moves no file, and an outbound handle with no content", func() {
|
||||
src := handleSource{transfer: true, fn: func(contract.ConnectorRequest) contract.ConnectorResponse { return contract.ConnectorResponse{} }}
|
||||
g := &Gateway{Connectors: src, ConnectorAuthority: allowAllConnectors{}, Transfers: t}
|
||||
_, err := g.ConnectorCall(context.Background(), scope, contract.ConnectorRequest{Op: contract.ConnectorOpSend, Kind: "slack", Transfer: &contract.ConnectorTransfer{Handle: "x"}})
|
||||
Expect(err).To(MatchError(ErrForbidden))
|
||||
_, err = g.ConnectorCall(context.Background(), scope, contract.ConnectorRequest{Op: contract.ConnectorOpSendFile, Kind: "slack", Transfer: &contract.ConnectorTransfer{Handle: "x", Size: 1}})
|
||||
Expect(err).To(MatchError(ErrForbidden), "the frontend names the content in the scope")
|
||||
})
|
||||
})
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/addons"
|
||||
"github.com/mudler/LocalAI/core/services/addons/hostops"
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
)
|
||||
|
||||
@@ -25,6 +26,10 @@ var ErrNotSent = fmt.Errorf("the add-on is not running now and nothing was sent:
|
||||
// Gateway deliberately admits lifecycle commands only until resource admission
|
||||
// and proposal confirmation are installed. Installation admin is not data access.
|
||||
type Gateway struct {
|
||||
// Transfers holds the handles of connector file transfers on the process
|
||||
// that runs the add-ons. Without it a call that carries a handle is
|
||||
// unavailable.
|
||||
Transfers *hostops.Transfers
|
||||
Connectors ConnectorSource
|
||||
ConnectorAuthority ConnectorAuthority
|
||||
// HostSettings is owner-only durable policy; frontends forward, never fall back.
|
||||
|
||||
@@ -54,5 +54,5 @@ func Requests(req *contract.HostOperations, registered, grants []string) Request
|
||||
// must not mistake that startup ordering for unsupported operations. The broker
|
||||
// still requires a registered handler and revalidates every request's authority.
|
||||
func SupportedOperations() []string {
|
||||
return strings.Fields(`finetune.start finetune.get finetune.list finetune.checkpoints finetune.dataset finetune.export finetune.progress finetune.stop quantize.start quantize.get quantize.list quantize.import quantize.progress quantize.stop workspace.materialize workspace.commit workspace.export workspace.import workspace.changes workspace.admin egress.check egress.audit connector.inbound`)
|
||||
return strings.Fields(`finetune.start finetune.get finetune.list finetune.checkpoints finetune.dataset finetune.export finetune.progress finetune.stop quantize.start quantize.get quantize.list quantize.import quantize.progress quantize.stop workspace.materialize workspace.commit workspace.export workspace.import workspace.changes workspace.admin egress.check egress.audit connector.inbound connector.file.put connector.file.get`)
|
||||
}
|
||||
@@ -37,6 +37,7 @@ func registeredBuiltinOperations() []string {
|
||||
})).To(Succeed())
|
||||
Expect(RegisterEgress(b, &fakeEgress{})).To(Succeed())
|
||||
Expect(RegisterConnector(b, admissionInbox{})).To(Succeed())
|
||||
Expect(RegisterConnectorTransfer(b, NewTransfers(nil))).To(Succeed())
|
||||
// Each registration error is checked above, so a duplicate cannot be hidden
|
||||
// by the broker's map. Ops includes dynamic progress/stop and inbound specs.
|
||||
return b.Ops()
|
||||
|
||||
@@ -327,6 +327,10 @@ func (b *Broker) serveInbound(w http.ResponseWriter, r *http.Request, spec Spec,
|
||||
status, code, msg = http.StatusForbidden, contract.HostScope, "not allowed"
|
||||
case errors.Is(err, ErrNotFound):
|
||||
status, code, msg = http.StatusNotFound, contract.HostNotFound, "not found"
|
||||
case errors.Is(err, ErrTransferTooLarge):
|
||||
status, code, msg = http.StatusRequestEntityTooLarge, contract.HostTooLarge, "the file is larger than the limit"
|
||||
case errors.Is(err, ErrConflict):
|
||||
status, code, msg = http.StatusConflict, contract.HostConflict, "conflict"
|
||||
case errors.Is(err, ErrBadRequest):
|
||||
status, code, msg = http.StatusBadRequest, contract.HostBadRequest, err.Error()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,450 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package hostops
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/studio/workspace"
|
||||
"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
)
|
||||
|
||||
// Direction of a transfer, from the host.
|
||||
const (
|
||||
// TransferIn moves a file from an add-on to the host: the add-on puts.
|
||||
TransferIn = "in"
|
||||
// TransferOut moves a file from the host to an add-on: the add-on gets.
|
||||
TransferOut = "out"
|
||||
)
|
||||
|
||||
// Results of a transfer, as the audit shows them.
|
||||
const (
|
||||
TransferDone = "done"
|
||||
TransferTooLarge = "too_large"
|
||||
TransferExpired = "expired"
|
||||
TransferRefused = "refused"
|
||||
TransferIncomplete = "incomplete"
|
||||
)
|
||||
|
||||
// Default times of a handle.
|
||||
const (
|
||||
// DefaultTransferIdle ends a handle that has seen no call for this long.
|
||||
DefaultTransferIdle = 2 * time.Minute
|
||||
// DefaultTransferLife ends a handle this long after it was made, whatever
|
||||
// it does. A connector call is shorter than this.
|
||||
DefaultTransferLife = time.Hour
|
||||
)
|
||||
|
||||
// ErrTransferTooLarge means a put went past the bound of its handle.
|
||||
var ErrTransferTooLarge = errors.New("the file is larger than the limit of the handle")
|
||||
|
||||
// TransferGrant is what the host that issued a handle records about it. The
|
||||
// owner of the add-ons registers it when a connector call that carries a
|
||||
// ConnectorTransfer arrives.
|
||||
type TransferGrant struct {
|
||||
Handle string
|
||||
Direction string
|
||||
// AddonID is the only add-on that may use the handle.
|
||||
AddonID string
|
||||
// BindingID is for the audit.
|
||||
BindingID string
|
||||
// MaxBytes bounds an inbound file.
|
||||
MaxBytes int64
|
||||
// SHA and Size name the outbound content, a blob of the workspace store.
|
||||
SHA string
|
||||
Size int64
|
||||
}
|
||||
|
||||
// TransferOutcome is what a handle ended with. It never holds a name or any
|
||||
// content.
|
||||
type TransferOutcome struct {
|
||||
Direction string
|
||||
BindingID string
|
||||
Bytes int64
|
||||
Result string
|
||||
Duration time.Duration
|
||||
// SHA is the hash of a finished inbound file.
|
||||
SHA string
|
||||
}
|
||||
|
||||
// FileScan is the hook after the last byte of an inbound file. An error
|
||||
// rejects the file. It is a no-op until an owner sets one.
|
||||
type FileScan func(ctx context.Context, sha string, size int64) error
|
||||
|
||||
// Transfers is the table of live handles of one process, and the bytes behind
|
||||
// them. Inbound bytes stream into the blob store of the workspace, which in a
|
||||
// cluster is the shared object store, so no node needs a shared disk. Outbound
|
||||
// bytes stream out of the same store. A file is never held in memory: a put
|
||||
// or a get moves one chunk.
|
||||
type Transfers struct {
|
||||
// Blobs is the blob store of the workspace. UseBlobs sets it when the
|
||||
// files start, which can be after the table is made.
|
||||
Blobs blob.Store
|
||||
Scan FileScan
|
||||
// Audit gets the outcome of every handle when it ends. Nil: none.
|
||||
Audit func(TransferOutcome)
|
||||
// Idle and Life default to DefaultTransferIdle and DefaultTransferLife.
|
||||
Idle, Life time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
m map[string]*xfer
|
||||
}
|
||||
|
||||
type xfer struct {
|
||||
t *Transfers
|
||||
g TransferGrant
|
||||
start time.Time
|
||||
|
||||
mu sync.Mutex // one call at a time per handle
|
||||
received int64
|
||||
pw *io.PipeWriter
|
||||
done chan struct{} // closed when the blob Put of an inbound handle returned
|
||||
sha string
|
||||
size int64
|
||||
putErr error
|
||||
final bool
|
||||
rd io.ReadSeekCloser
|
||||
cancel context.CancelFunc
|
||||
idle *time.Timer
|
||||
life *time.Timer
|
||||
ended bool
|
||||
result string
|
||||
}
|
||||
|
||||
// NewTransfers builds a table over the blob store of the workspace.
|
||||
func NewTransfers(blobs blob.Store) *Transfers {
|
||||
return &Transfers{Blobs: blobs, m: map[string]*xfer{}}
|
||||
}
|
||||
|
||||
// UseBlobs sets the blob store. Handles that begin before it is set are
|
||||
// refused as unavailable.
|
||||
func (t *Transfers) UseBlobs(b blob.Store) {
|
||||
t.mu.Lock()
|
||||
t.Blobs = b
|
||||
t.mu.Unlock()
|
||||
}
|
||||
|
||||
func (t *Transfers) store() blob.Store {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.Blobs
|
||||
}
|
||||
|
||||
func (t *Transfers) idle() time.Duration {
|
||||
if t.Idle > 0 {
|
||||
return t.Idle
|
||||
}
|
||||
return DefaultTransferIdle
|
||||
}
|
||||
|
||||
func (t *Transfers) life() time.Duration {
|
||||
if t.Life > 0 {
|
||||
return t.Life
|
||||
}
|
||||
return DefaultTransferLife
|
||||
}
|
||||
|
||||
// Begin registers a handle. The returned function ends it: it drops a file
|
||||
// that is not finished and reports the outcome. For a finished inbound file the
|
||||
// outcome holds the hash and the size. Calling it again returns the same
|
||||
// outcome. A handle that is already registered is refused, so a handle is
|
||||
// never used for two calls.
|
||||
func (t *Transfers) Begin(g TransferGrant) (end func() TransferOutcome, err error) {
|
||||
if t == nil || t.store() == nil {
|
||||
return nil, errors.New("file transfer is not available here")
|
||||
}
|
||||
if g.Handle == "" || len(g.Handle) > contract.MaxTransferHandle || g.AddonID == "" {
|
||||
return nil, fmt.Errorf("%w: bad handle", ErrBadRequest)
|
||||
}
|
||||
switch g.Direction {
|
||||
case TransferIn:
|
||||
if g.MaxBytes <= 0 {
|
||||
return nil, fmt.Errorf("%w: no bound", ErrBadRequest)
|
||||
}
|
||||
case TransferOut:
|
||||
if g.SHA == "" || g.Size <= 0 {
|
||||
return nil, fmt.Errorf("%w: no content", ErrBadRequest)
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("%w: bad direction", ErrBadRequest)
|
||||
}
|
||||
x := &xfer{t: t, g: g, start: time.Now()}
|
||||
t.mu.Lock()
|
||||
if t.m == nil {
|
||||
t.m = map[string]*xfer{}
|
||||
}
|
||||
if _, dup := t.m[g.Handle]; dup {
|
||||
t.mu.Unlock()
|
||||
return nil, fmt.Errorf("%w: handle in use", ErrConflict)
|
||||
}
|
||||
t.m[g.Handle] = x
|
||||
t.mu.Unlock()
|
||||
x.idle = time.AfterFunc(t.idle(), func() { x.finish(TransferExpired) })
|
||||
x.life = time.AfterFunc(t.life(), func() { x.finish(TransferExpired) })
|
||||
return func() TransferOutcome { return x.finish("") }, nil
|
||||
}
|
||||
|
||||
// Live reports how many handles exist; a test and the sweep use it.
|
||||
func (t *Transfers) Live() int {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return len(t.m)
|
||||
}
|
||||
|
||||
func (t *Transfers) lookup(addonID, handle string) (*xfer, error) {
|
||||
t.mu.Lock()
|
||||
x := t.m[handle]
|
||||
t.mu.Unlock()
|
||||
// The same answer for a handle that does not exist and one of another
|
||||
// add-on: a handle cannot be probed.
|
||||
if x == nil || subtle.ConstantTimeCompare([]byte(x.g.AddonID), []byte(addonID)) != 1 {
|
||||
return nil, ErrNotFound
|
||||
}
|
||||
return x, nil
|
||||
}
|
||||
|
||||
// finish ends the handle once. why is the result when the handle did not end
|
||||
// by its own end function.
|
||||
func (x *xfer) finish(why string) TransferOutcome {
|
||||
x.idle.Stop()
|
||||
x.life.Stop()
|
||||
// A put or a get in progress holds mu; an expiry must not wait for a
|
||||
// stuck one, so it cancels first.
|
||||
if x.cancel != nil {
|
||||
x.cancel()
|
||||
}
|
||||
if x.pw != nil {
|
||||
_ = x.pw.CloseWithError(context.Canceled)
|
||||
}
|
||||
x.mu.Lock()
|
||||
defer x.mu.Unlock()
|
||||
if x.ended {
|
||||
return x.outcome()
|
||||
}
|
||||
x.ended = true
|
||||
x.t.mu.Lock()
|
||||
delete(x.t.m, x.g.Handle)
|
||||
x.t.mu.Unlock()
|
||||
if x.rd != nil {
|
||||
_ = x.rd.Close()
|
||||
}
|
||||
if x.done != nil {
|
||||
<-x.done // the blob Put drops its temp file when its reader fails
|
||||
}
|
||||
switch {
|
||||
case x.result != "":
|
||||
case why != "":
|
||||
x.result = why
|
||||
case x.g.Direction == TransferIn && !x.final:
|
||||
x.result = TransferIncomplete
|
||||
default:
|
||||
x.result = TransferDone
|
||||
}
|
||||
o := x.outcome()
|
||||
if x.t.Audit != nil {
|
||||
x.t.Audit(o)
|
||||
}
|
||||
return o
|
||||
}
|
||||
|
||||
func (x *xfer) outcome() TransferOutcome {
|
||||
o := TransferOutcome{Direction: x.g.Direction, BindingID: x.g.BindingID, Bytes: x.received, Result: x.result, Duration: time.Since(x.start)}
|
||||
if x.g.Direction == TransferOut {
|
||||
o.Bytes = x.size
|
||||
}
|
||||
if x.final && x.result != TransferTooLarge && x.result != TransferRefused {
|
||||
o.SHA = x.sha
|
||||
}
|
||||
return o
|
||||
}
|
||||
|
||||
func (x *xfer) touch() { x.idle.Reset(x.t.idle()) }
|
||||
|
||||
// put writes the next chunk of an inbound file.
|
||||
func (x *xfer) put(ctx context.Context, p contract.ConnectorFilePut) (int64, error) {
|
||||
if x.g.Direction != TransferIn {
|
||||
return 0, ErrNotFound
|
||||
}
|
||||
x.mu.Lock()
|
||||
defer x.mu.Unlock()
|
||||
if x.ended || x.final {
|
||||
return 0, ErrConflict
|
||||
}
|
||||
x.touch()
|
||||
if len(p.Data) > contract.TransferChunkBytes {
|
||||
return 0, fmt.Errorf("%w: chunk too large", ErrBadRequest)
|
||||
}
|
||||
if p.Offset != x.received {
|
||||
// The add-on resumes where the host is.
|
||||
return x.received, nil
|
||||
}
|
||||
if x.received+int64(len(p.Data)) > x.g.MaxBytes {
|
||||
x.result = TransferTooLarge
|
||||
x.abortLocked()
|
||||
return x.received, ErrTransferTooLarge
|
||||
}
|
||||
if x.pw == nil {
|
||||
pr, pw := io.Pipe()
|
||||
bctx, cancel := context.WithCancel(context.Background())
|
||||
x.pw, x.cancel, x.done = pw, cancel, make(chan struct{})
|
||||
go func() {
|
||||
defer close(x.done)
|
||||
sha, size, err := x.t.store().Put(bctx, pr, x.g.MaxBytes)
|
||||
x.sha, x.size, x.putErr = sha, size, err
|
||||
_ = pr.CloseWithError(errors.New("closed"))
|
||||
}()
|
||||
}
|
||||
if len(p.Data) > 0 {
|
||||
if _, err := io.Copy(x.pw, bytes.NewReader(p.Data)); err != nil {
|
||||
return x.received, x.failedLocked(err)
|
||||
}
|
||||
x.received += int64(len(p.Data))
|
||||
}
|
||||
if !p.Final {
|
||||
return x.received, nil
|
||||
}
|
||||
if x.received == 0 {
|
||||
x.result = TransferRefused
|
||||
x.abortLocked()
|
||||
return 0, fmt.Errorf("%w: an empty file", ErrBadRequest)
|
||||
}
|
||||
_ = x.pw.Close()
|
||||
<-x.done
|
||||
switch {
|
||||
case errors.Is(x.putErr, workspace.ErrTooLarge):
|
||||
x.result = TransferTooLarge
|
||||
return x.received, ErrTransferTooLarge
|
||||
case x.putErr != nil:
|
||||
return x.received, x.putErr
|
||||
case x.size != x.received:
|
||||
return x.received, errors.New("the stored size differs from the bytes received")
|
||||
}
|
||||
if x.t.Scan != nil {
|
||||
if err := x.t.Scan(ctx, x.sha, x.size); err != nil {
|
||||
x.result = TransferRefused
|
||||
return x.received, fmt.Errorf("%w: the file was refused", ErrBadRequest)
|
||||
}
|
||||
}
|
||||
x.final = true
|
||||
return x.received, nil
|
||||
}
|
||||
|
||||
// failedLocked turns a failed write into the error of the blob store, which
|
||||
// ended the pipe, so that a bound reached inside the store is too large.
|
||||
func (x *xfer) failedLocked(err error) error {
|
||||
<-x.done
|
||||
if errors.Is(x.putErr, workspace.ErrTooLarge) {
|
||||
x.result = TransferTooLarge
|
||||
return ErrTransferTooLarge
|
||||
}
|
||||
if x.putErr != nil {
|
||||
return x.putErr
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (x *xfer) abortLocked() {
|
||||
if x.pw != nil {
|
||||
_ = x.pw.CloseWithError(context.Canceled)
|
||||
if x.cancel != nil {
|
||||
x.cancel()
|
||||
}
|
||||
<-x.done
|
||||
}
|
||||
}
|
||||
|
||||
// get reads the next chunk of an outbound file.
|
||||
func (x *xfer) get(ctx context.Context, g contract.ConnectorFileGet) (contract.ConnectorFileGetResult, error) {
|
||||
if x.g.Direction != TransferOut {
|
||||
return contract.ConnectorFileGetResult{}, ErrNotFound
|
||||
}
|
||||
x.mu.Lock()
|
||||
defer x.mu.Unlock()
|
||||
if x.ended {
|
||||
return contract.ConnectorFileGetResult{}, ErrConflict
|
||||
}
|
||||
x.touch()
|
||||
if g.Offset < 0 || g.Offset > x.g.Size || g.Length < 1 || g.Length > contract.TransferChunkBytes {
|
||||
return contract.ConnectorFileGetResult{}, fmt.Errorf("%w: offset or length", ErrBadRequest)
|
||||
}
|
||||
if x.rd == nil {
|
||||
r, err := x.t.store().Open(ctx, x.g.SHA)
|
||||
if err != nil {
|
||||
return contract.ConnectorFileGetResult{}, err
|
||||
}
|
||||
n, err := r.Seek(0, io.SeekEnd)
|
||||
if err != nil || n != x.g.Size {
|
||||
_ = r.Close()
|
||||
return contract.ConnectorFileGetResult{}, errors.New("the file changed since the handle was made")
|
||||
}
|
||||
x.rd, x.size = r, n
|
||||
}
|
||||
if _, err := x.rd.Seek(g.Offset, io.SeekStart); err != nil {
|
||||
return contract.ConnectorFileGetResult{}, err
|
||||
}
|
||||
buf := make([]byte, min(g.Length, x.g.Size-g.Offset))
|
||||
n, err := io.ReadFull(x.rd, buf)
|
||||
if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) && !errors.Is(err, io.EOF) {
|
||||
return contract.ConnectorFileGetResult{}, err
|
||||
}
|
||||
buf = buf[:n]
|
||||
return contract.ConnectorFileGetResult{Data: buf, EOF: g.Offset+int64(n) >= x.g.Size}, nil
|
||||
}
|
||||
|
||||
// RegisterConnectorTransfer registers the two operations with which a
|
||||
// connector add-on moves file bytes by handle. Both are inbound class: the
|
||||
// credential of the live incarnation authorizes them, and the handle binds
|
||||
// them to one connector call. The arguments name no user.
|
||||
func RegisterConnectorTransfer(b *Broker, t *Transfers) error {
|
||||
if t == nil {
|
||||
return errors.New("connector file transfer needs a table")
|
||||
}
|
||||
decode := func(raw json.RawMessage, v any) error {
|
||||
d := json.NewDecoder(bytes.NewReader(raw))
|
||||
d.DisallowUnknownFields()
|
||||
if err := d.Decode(v); err != nil {
|
||||
return ErrBadRequest
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if err := b.Register(Spec{Op: contract.ConnectorFilePutOp, Scope: ScopeInbound, Timeout: 2 * time.Minute, Handler: func(ctx context.Context, c Call) (json.RawMessage, error) {
|
||||
var p contract.ConnectorFilePut
|
||||
if err := decode(c.Arguments, &p); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
x, err := t.lookup(c.AddonID, p.Handle)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
got, err := x.put(ctx, p)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return json.Marshal(contract.ConnectorFilePutResult{Received: got})
|
||||
}}); err != nil {
|
||||
return err
|
||||
}
|
||||
return b.Register(Spec{Op: contract.ConnectorFileGetOp, Scope: ScopeInbound, Timeout: 2 * time.Minute, Handler: func(ctx context.Context, c Call) (json.RawMessage, error) {
|
||||
var g contract.ConnectorFileGet
|
||||
if err := decode(c.Arguments, &g); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
x, err := t.lookup(c.AddonID, g.Handle)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r, err := x.get(ctx, g)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return json.Marshal(r)
|
||||
}})
|
||||
}
|
||||
@@ -0,0 +1,284 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package hostops
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"math/rand"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
|
||||
"github.com/mudler/LocalAI/pkg/addons/contract"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type xferFixture struct {
|
||||
t *Transfers
|
||||
h http.Handler
|
||||
outs []TransferOutcome
|
||||
blobs blob.Store
|
||||
dir string
|
||||
shared *storage.FilesystemStore
|
||||
}
|
||||
|
||||
func newXferFixture() *xferFixture {
|
||||
dir := GinkgoT().TempDir()
|
||||
shared, err := storage.NewFilesystemStore(filepath.Join(dir, "objects"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
// The owner node has its own cache disk; the object store is the only
|
||||
// thing the nodes share.
|
||||
blobs := blob.NewObjectStore(shared, filepath.Join(dir, "owner-cache"))
|
||||
f := &xferFixture{blobs: blobs, dir: dir, shared: shared}
|
||||
f.t = NewTransfers(blobs)
|
||||
f.t.Audit = func(o TransferOutcome) { f.outs = append(f.outs, o) }
|
||||
b := New(NewInFlight(), always, contract.DefaultLimits())
|
||||
Expect(RegisterConnectorTransfer(b, f.t)).To(Succeed())
|
||||
f.h = b.Handler("slack", "secret", []string{"connector."})
|
||||
return f
|
||||
}
|
||||
|
||||
// call posts one host operation as the add-on addon.
|
||||
func (f *xferFixture) call(addon, op string, args any) (int, contract.HostResponse) {
|
||||
raw, _ := json.Marshal(args)
|
||||
body, _ := json.Marshal(contract.HostRequest{OperationID: "x", Op: op, Arguments: raw})
|
||||
req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body))
|
||||
req.Header.Set(contract.HostCredentialHeader, "secret")
|
||||
rec := httptest.NewRecorder()
|
||||
h := f.h
|
||||
if addon != "slack" {
|
||||
b := New(NewInFlight(), always, contract.DefaultLimits())
|
||||
Expect(RegisterConnectorTransfer(b, f.t)).To(Succeed())
|
||||
h = b.Handler(addon, "secret", []string{"connector."})
|
||||
}
|
||||
h.ServeHTTP(rec, req)
|
||||
var out contract.HostResponse
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
|
||||
func (f *xferFixture) put(handle string, off int64, data []byte, final bool) (int, int64) {
|
||||
code, r := f.call("slack", contract.ConnectorFilePutOp, contract.ConnectorFilePut{Handle: handle, Offset: off, Data: data, Final: final})
|
||||
var res contract.ConnectorFilePutResult
|
||||
if r.Result != nil {
|
||||
_ = json.Unmarshal(r.Result, &res)
|
||||
}
|
||||
return code, res.Received
|
||||
}
|
||||
|
||||
func (f *xferFixture) get(handle string, off, n int64) (int, contract.ConnectorFileGetResult) {
|
||||
code, r := f.call("slack", contract.ConnectorFileGetOp, contract.ConnectorFileGet{Handle: handle, Offset: off, Length: n})
|
||||
var res contract.ConnectorFileGetResult
|
||||
if r.Result != nil {
|
||||
_ = json.Unmarshal(r.Result, &res)
|
||||
}
|
||||
return code, res
|
||||
}
|
||||
|
||||
// chunks streams n bytes of a seeded sequence, chunk bytes at a time, and
|
||||
// never holds more than one chunk.
|
||||
func chunks(n, chunk int64, each func(off int64, b []byte, last bool)) string {
|
||||
rng := rand.New(rand.NewSource(7))
|
||||
h := sha256.New()
|
||||
buf := make([]byte, chunk)
|
||||
for off := int64(0); off < n; off += chunk {
|
||||
b := buf[:min(chunk, n-off)]
|
||||
_, _ = rng.Read(b)
|
||||
h.Write(b)
|
||||
each(off, b, off+int64(len(b)) >= n)
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
func heap() uint64 {
|
||||
var m runtime.MemStats
|
||||
runtime.GC()
|
||||
runtime.ReadMemStats(&m)
|
||||
return m.HeapAlloc
|
||||
}
|
||||
|
||||
var _ = Describe("connector file transfer by handle", func() {
|
||||
const mb = 1 << 20
|
||||
|
||||
It("moves 20 MB in and 20 MB out with bounded memory", func() {
|
||||
f := newXferFixture()
|
||||
end, err := f.t.Begin(TransferGrant{Handle: "h1", Direction: TransferIn, AddonID: "slack", BindingID: "b1", MaxBytes: 50 * mb})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
base, peak := heap(), uint64(0)
|
||||
var count int
|
||||
want := chunks(20*mb, contract.TransferChunkBytes, func(off int64, b []byte, last bool) {
|
||||
code, got := f.put("h1", off, append([]byte(nil), b...), last)
|
||||
Expect(code).To(Equal(http.StatusOK))
|
||||
Expect(got).To(Equal(off + int64(len(b))))
|
||||
if count++; count%16 == 0 {
|
||||
peak = max(peak, heap())
|
||||
}
|
||||
})
|
||||
out := end()
|
||||
Expect(out.Result).To(Equal(TransferDone))
|
||||
Expect(out.Bytes).To(Equal(int64(20 * mb)))
|
||||
Expect(out.SHA).To(Equal(want))
|
||||
Expect(peak).To(BeNumerically("<", base+16*mb), "the file is never held in memory")
|
||||
Expect(f.t.Live()).To(BeZero())
|
||||
|
||||
// Another node: same object store, its own cache disk.
|
||||
node2 := blob.NewObjectStore(f.shared, filepath.Join(f.dir, "frontend-cache"))
|
||||
r, err := node2.Open(GinkgoT().Context(), out.SHA)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { _ = r.Close() })
|
||||
h := sha256.New()
|
||||
n, err := io.Copy(h, r)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(n).To(Equal(int64(20 * mb)))
|
||||
Expect(hex.EncodeToString(h.Sum(nil))).To(Equal(want))
|
||||
|
||||
// Out: the add-on pulls the same file by a new handle.
|
||||
end2, err := f.t.Begin(TransferGrant{Handle: "h2", Direction: TransferOut, AddonID: "slack", SHA: out.SHA, Size: n})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
hh := sha256.New()
|
||||
var off int64
|
||||
base, peak = heap(), 0
|
||||
for i := 0; ; i++ {
|
||||
code, res := f.get("h2", off, contract.TransferChunkBytes)
|
||||
Expect(code).To(Equal(http.StatusOK))
|
||||
hh.Write(res.Data)
|
||||
off += int64(len(res.Data))
|
||||
if i%16 == 0 {
|
||||
peak = max(peak, heap())
|
||||
}
|
||||
if res.EOF {
|
||||
break
|
||||
}
|
||||
}
|
||||
Expect(off).To(Equal(n))
|
||||
Expect(hex.EncodeToString(hh.Sum(nil))).To(Equal(want))
|
||||
Expect(peak).To(BeNumerically("<", base+16*mb))
|
||||
Expect(end2().Result).To(Equal(TransferDone))
|
||||
Expect(f.outs).To(HaveLen(2))
|
||||
})
|
||||
|
||||
It("resumes at the offset the host has after a gap", func() {
|
||||
f := newXferFixture()
|
||||
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: mb})
|
||||
code, got := f.put("h", 0, []byte("abcd"), false)
|
||||
Expect(code).To(Equal(http.StatusOK))
|
||||
Expect(got).To(BeEquivalentTo(4))
|
||||
_, got = f.put("h", 10, []byte("zz"), false)
|
||||
Expect(got).To(BeEquivalentTo(4), "a gap answers where the host is and writes nothing")
|
||||
_, got = f.put("h", 4, []byte("efgh"), true)
|
||||
Expect(got).To(BeEquivalentTo(8))
|
||||
o := end()
|
||||
sum := sha256.Sum256([]byte("abcdefgh"))
|
||||
Expect(o.SHA).To(Equal(hex.EncodeToString(sum[:])))
|
||||
})
|
||||
|
||||
It("refuses an oversize file and drops the partial one", func() {
|
||||
f := newXferFixture()
|
||||
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 10})
|
||||
code, _ := f.put("h", 0, []byte("12345678"), false)
|
||||
Expect(code).To(Equal(http.StatusOK))
|
||||
code, _ = f.put("h", 8, []byte("9012"), true)
|
||||
Expect(code).To(Equal(http.StatusRequestEntityTooLarge))
|
||||
o := end()
|
||||
Expect(o.Result).To(Equal(TransferTooLarge))
|
||||
Expect(o.SHA).To(BeEmpty())
|
||||
code, _ = f.put("h", 0, []byte("x"), false)
|
||||
Expect(code).To(Equal(http.StatusNotFound), "the handle ended")
|
||||
})
|
||||
|
||||
It("refuses a file whose last chunk says it is short of the bytes sent, an empty file, and a write after final", func() {
|
||||
f := newXferFixture()
|
||||
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
code, _ := f.put("h", 0, nil, true)
|
||||
Expect(code).To(Equal(http.StatusBadRequest), "empty")
|
||||
Expect(end().Result).To(Equal(TransferRefused))
|
||||
|
||||
end, _ = f.t.Begin(TransferGrant{Handle: "g", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
code, _ = f.put("g", 0, []byte("ab"), true)
|
||||
Expect(code).To(Equal(http.StatusOK))
|
||||
code, _ = f.put("g", 2, []byte("cd"), false)
|
||||
Expect(code).To(Equal(http.StatusConflict), "write after final")
|
||||
Expect(end().Result).To(Equal(TransferDone))
|
||||
})
|
||||
|
||||
It("refuses another add-on, an unknown handle and the wrong direction with one answer", func() {
|
||||
f := newXferFixture()
|
||||
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
defer end()
|
||||
other, _ := f.call("telegram", contract.ConnectorFilePutOp, contract.ConnectorFilePut{Handle: "h", Data: []byte("x")})
|
||||
unknown, _ := f.call("slack", contract.ConnectorFilePutOp, contract.ConnectorFilePut{Handle: "nope", Data: []byte("x")})
|
||||
Expect(other).To(Equal(http.StatusNotFound))
|
||||
Expect(unknown).To(Equal(other))
|
||||
code, _ := f.get("h", 0, 10)
|
||||
Expect(code).To(Equal(http.StatusNotFound), "a write handle cannot read")
|
||||
})
|
||||
|
||||
It("does not let a handle be used for two calls", func() {
|
||||
f := newXferFixture()
|
||||
end, err := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
_, err = f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
Expect(err).To(HaveOccurred())
|
||||
end()
|
||||
end() // idempotent
|
||||
Expect(f.outs).To(HaveLen(1))
|
||||
_, err = f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
Expect(err).NotTo(HaveOccurred(), "a handle of an ended call is free; a new call makes its own random one")
|
||||
})
|
||||
|
||||
It("expires a handle that sees no call and removes the partial file", func() {
|
||||
f := newXferFixture()
|
||||
f.t.Idle = 80 * time.Millisecond
|
||||
_, err := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
code, _ := f.put("h", 0, []byte("partial"), false)
|
||||
Expect(code).To(Equal(http.StatusOK))
|
||||
Eventually(f.t.Live, "2s", "20ms").Should(BeZero())
|
||||
code, _ = f.put("h", 7, []byte("x"), false)
|
||||
Expect(code).To(Equal(http.StatusNotFound))
|
||||
Expect(f.outs).To(HaveLen(1))
|
||||
Expect(f.outs[0].Result).To(Equal(TransferExpired))
|
||||
tmp, _ := filepath.Glob(filepath.Join(f.dir, "owner-cache", "tmp", "*"))
|
||||
Expect(tmp).To(BeEmpty(), "no partial file is left")
|
||||
})
|
||||
|
||||
It("refuses a file the scan hook rejects", func() {
|
||||
f := newXferFixture()
|
||||
f.t.Scan = func(_ context.Context, _ string, _ int64) error { return io.ErrClosedPipe }
|
||||
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
|
||||
code, _ := f.put("h", 0, []byte("virus"), true)
|
||||
Expect(code).To(Equal(http.StatusBadRequest))
|
||||
o := end()
|
||||
Expect(o.Result).To(Equal(TransferRefused))
|
||||
Expect(o.SHA).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses an outbound handle whose content changed size and an offset past the end", func() {
|
||||
f := newXferFixture()
|
||||
sha, size, err := f.blobs.Put(GinkgoT().Context(), bytes.NewReader([]byte("hello world")), 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
bad, _ := f.t.Begin(TransferGrant{Handle: "bad", Direction: TransferOut, AddonID: "slack", SHA: sha, Size: size + 1})
|
||||
code, _ := f.get("bad", 0, 5)
|
||||
Expect(code).To(Equal(http.StatusServiceUnavailable))
|
||||
bad()
|
||||
end, _ := f.t.Begin(TransferGrant{Handle: "ok", Direction: TransferOut, AddonID: "slack", SHA: sha, Size: size})
|
||||
defer end()
|
||||
code, res := f.get("ok", 6, contract.TransferChunkBytes+1)
|
||||
Expect(code).To(Equal(http.StatusBadRequest), "a length above the chunk is refused")
|
||||
code, res = f.get("ok", 6, 5)
|
||||
Expect(code).To(Equal(http.StatusOK))
|
||||
Expect(string(res.Data)).To(Equal("world"))
|
||||
Expect(res.EOF).To(BeTrue())
|
||||
code, _ = f.get("ok", size+1, 1)
|
||||
Expect(code).To(Equal(http.StatusBadRequest))
|
||||
})
|
||||
})
|
||||
@@ -266,7 +266,7 @@ func (c *Client) Connector(ctx context.Context, v contract.ConnectorRequest) (co
|
||||
if c.registration.Capabilities.Connector == nil {
|
||||
return r, ErrNotSent
|
||||
}
|
||||
e := c.call(ctx, "/connector", nil, v, &r, contract.ConnectorTimeout)
|
||||
e := c.call(ctx, "/connector", nil, v, &r, v.CallTimeout())
|
||||
return r, e
|
||||
}
|
||||
func (c *Client) Preview(ctx context.Context, v contract.Envelope) (contract.Preview, error) {
|
||||
|
||||
@@ -12,6 +12,16 @@ type ConnectorCallScope struct {
|
||||
AgentID string `json:"agentId,omitempty"`
|
||||
AddonID string `json:"addonId"`
|
||||
Kind string `json:"kind"`
|
||||
// Content names the file an outbound transfer reads: a blob of the
|
||||
// workspace store. The frontend that checked the chat's access sets it;
|
||||
// it never reaches the add-on, which gets only the handle.
|
||||
Content *TransferContent `json:"content,omitempty"`
|
||||
}
|
||||
|
||||
// TransferContent is the stored content behind an outbound handle.
|
||||
type TransferContent struct {
|
||||
SHA string `json:"sha"`
|
||||
Size int64 `json:"size"`
|
||||
}
|
||||
|
||||
const (
|
||||
|
||||
Reference in new issue
Block a user