mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
feat(distributed): release exact staged keys
Request inputs need transport-neutral cleanup after backend calls. Add authenticated exact-key deletion for HTTP and coordinated cache eviction before shared-object deletion for S3/NATS. Preserve URL metacharacters as filename data, reject unsafe keys, and remove upload sidecars while pruning empty request directories. Assisted-by: Codex:gpt-6
This commit is contained in:
1 parent
2b8bee0b8e
commit
171347ff18
9 files changed
+577
-1
No files matched your search
@@ -439,6 +439,12 @@ func SubjectNodeFilesStage(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.stage"
|
||||
}
|
||||
|
||||
// SubjectNodeFilesRelease tells a serve-backend node to evict one ephemeral cache key.
|
||||
// Reply: {error}
|
||||
func SubjectNodeFilesRelease(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.release"
|
||||
}
|
||||
|
||||
// SubjectNodeFilesTemp tells a serve-backend node to allocate a temp file.
|
||||
// Reply: {local_path, error}
|
||||
func SubjectNodeFilesTemp(nodeID string) string {
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
package nodes
|
||||
|
||||
import "context"
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// FileStager abstracts file transfer between frontend and backend nodes
|
||||
// in distributed mode. Two implementations exist:
|
||||
@@ -29,7 +34,26 @@ type FileStager interface {
|
||||
// StageRemoteToStore uploads a remote file to shared storage.
|
||||
StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error
|
||||
|
||||
// ReleaseRemote removes one ephemeral key from the remote node.
|
||||
ReleaseRemote(ctx context.Context, nodeID, key string) error
|
||||
|
||||
// ListRemoteDir returns relative file paths within a directory on the remote node.
|
||||
// keyPrefix is a storage-style key prefix (e.g. "models/mymodel").
|
||||
ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error)
|
||||
}
|
||||
|
||||
func validateEphemeralReleaseKey(key string) error {
|
||||
if strings.Contains(key, "\\") || path.Clean(key) != key {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
parts := strings.Split(key, "/")
|
||||
if len(parts) != 4 || parts[0] != "ephemeral" {
|
||||
return fmt.Errorf("release key %q must identify one file below ephemeral/", key)
|
||||
}
|
||||
for _, part := range parts[1:] {
|
||||
if part == "" || part == "." || part == ".." {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
@@ -81,6 +82,35 @@ func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token st
|
||||
}
|
||||
}
|
||||
|
||||
// ReleaseRemote removes one exact ephemeral key from a backend node.
|
||||
func (h *HTTPFileStager) ReleaseRemote(ctx context.Context, nodeID, key string) error {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
addr, err := h.httpAddrFor(nodeID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err)
|
||||
}
|
||||
releaseURL := (&url.URL{Scheme: "http", Host: addr, Path: "/v1/files/" + key}).String()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, releaseURL, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating release request for %q: %w", key, err)
|
||||
}
|
||||
if h.token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+h.token)
|
||||
}
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("releasing %q from node %s: %w", key, nodeID, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return fmt.Errorf("releasing %q from node %s: status %d: %s", key, nodeID, resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
|
||||
xlog.Debug("Staging file to remote node via HTTP", "node", nodeID, "localPath", localPath, "key", key)
|
||||
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type releaseTestSubscription struct{}
|
||||
|
||||
func (releaseTestSubscription) Unsubscribe() error { return nil }
|
||||
|
||||
type releaseTestMessaging struct {
|
||||
subject string
|
||||
payload []byte
|
||||
onRequest func()
|
||||
requestCalled bool
|
||||
}
|
||||
|
||||
func (m *releaseTestMessaging) Publish(string, any) error { return nil }
|
||||
func (m *releaseTestMessaging) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) Request(subject string, data []byte, _ time.Duration) ([]byte, error) {
|
||||
m.subject = subject
|
||||
m.payload = append([]byte(nil), data...)
|
||||
m.requestCalled = true
|
||||
if m.onRequest != nil {
|
||||
m.onRequest()
|
||||
}
|
||||
return []byte(`{}`), nil
|
||||
}
|
||||
func (m *releaseTestMessaging) IsConnected() bool { return true }
|
||||
func (m *releaseTestMessaging) Close() {}
|
||||
|
||||
var _ = Describe("File stager exact-key release", func() {
|
||||
startReleaseServer := func(stagingDir, token string) (*HTTPFileStager, func()) {
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
server, err := StartFileTransferServerWithListener(
|
||||
listener,
|
||||
stagingDir,
|
||||
GinkgoT().TempDir(),
|
||||
GinkgoT().TempDir(),
|
||||
token,
|
||||
0,
|
||||
)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return NewHTTPFileStager(func(string) (string, error) {
|
||||
return listener.Addr().String(), nil
|
||||
}, token), func() {
|
||||
Expect(server.Shutdown(context.Background())).To(Succeed())
|
||||
}
|
||||
}
|
||||
|
||||
It("transmits URL metacharacters as the exact key", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(stagingDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
|
||||
key := "ephemeral/request-id/audio/name ?#%2F.wav"
|
||||
exactPath := filepath.Join(categoryDir, "name ?#%2F.wav")
|
||||
wrongPath := filepath.Join(categoryDir, "name ")
|
||||
for _, path := range []string{exactPath, exactPath + hashSidecarSuffix, exactPath + targetSidecarSuffix, wrongPath} {
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
}
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
DeferCleanup(stop)
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", key)).To(Succeed())
|
||||
|
||||
Expect(exactPath).NotTo(BeAnExistingFile())
|
||||
Expect(exactPath + hashSidecarSuffix).NotTo(BeAnExistingFile())
|
||||
Expect(exactPath + targetSidecarSuffix).NotTo(BeAnExistingFile())
|
||||
Expect(wrongPath).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("is idempotent and prunes empty category and request directories", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
path := filepath.Join(stagingDir, "ephemeral", "request-id", "audio", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
DeferCleanup(stop)
|
||||
for range 2 {
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", "ephemeral/request-id/audio/input.wav")).To(Succeed())
|
||||
}
|
||||
|
||||
Expect(filepath.Join(stagingDir, "ephemeral", "request-id", "audio")).NotTo(BeADirectory())
|
||||
Expect(filepath.Join(stagingDir, "ephemeral", "request-id")).NotTo(BeADirectory())
|
||||
Expect(filepath.Join(stagingDir, "ephemeral")).To(BeADirectory())
|
||||
})
|
||||
|
||||
It("rejects non-ephemeral and traversing keys before making a request", func() {
|
||||
resolved := false
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
resolved = true
|
||||
return "127.0.0.1:1", nil
|
||||
}, "token")
|
||||
|
||||
for _, key := range []string{
|
||||
"models/model.gguf",
|
||||
"ephemeral/../models/model.gguf",
|
||||
"ephemeral/request-id/../../model.gguf",
|
||||
"/ephemeral/request-id/audio/input.wav",
|
||||
} {
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", key)).NotTo(Succeed(), key)
|
||||
}
|
||||
Expect(resolved).To(BeFalse())
|
||||
})
|
||||
|
||||
It("rejects symlink escapes", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
outsideDir := GinkgoT().TempDir()
|
||||
outsidePath := filepath.Join(outsideDir, "input.wav")
|
||||
Expect(os.WriteFile(outsidePath, []byte("keep"), 0640)).To(Succeed())
|
||||
requestDir := filepath.Join(stagingDir, "ephemeral", "request-id")
|
||||
Expect(os.MkdirAll(requestDir, 0750)).To(Succeed())
|
||||
Expect(os.Symlink(outsideDir, filepath.Join(requestDir, "audio"))).To(Succeed())
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
DeferCleanup(stop)
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", "ephemeral/request-id/audio/input.wav")).NotTo(Succeed())
|
||||
Expect(outsidePath).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("evicts the worker cache before deleting the shared object", func() {
|
||||
storeRoot := GinkgoT().TempDir()
|
||||
cacheRoot := GinkgoT().TempDir()
|
||||
store, err := storage.NewFilesystemStore(storeRoot)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, cacheRoot)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
key := "ephemeral/request-id/audio/input.wav"
|
||||
Expect(store.Put(context.Background(), key, strings.NewReader("shared"))).To(Succeed())
|
||||
|
||||
client := &releaseTestMessaging{}
|
||||
client.onRequest = func() {
|
||||
exists, existsErr := store.Exists(context.Background(), key)
|
||||
Expect(existsErr).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeTrue())
|
||||
}
|
||||
stager := NewS3NATSFileStager(fm, client)
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node.one", key)).To(Succeed())
|
||||
|
||||
Expect(client.requestCalled).To(BeTrue())
|
||||
Expect(client.subject).To(Equal(messaging.SubjectNodeFilesRelease("node.one")))
|
||||
var payload fileReleaseRequest
|
||||
Expect(json.Unmarshal(client.payload, &payload)).To(Succeed())
|
||||
Expect(payload.Key).To(Equal(key))
|
||||
exists, err := store.Exists(context.Background(), key)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeFalse())
|
||||
})
|
||||
})
|
||||
@@ -48,6 +48,14 @@ type fileStageReply struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type fileReleaseRequest struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
|
||||
type fileReleaseReply struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type fileTempRequest struct{}
|
||||
|
||||
type fileTempReply struct {
|
||||
@@ -181,3 +189,27 @@ func (s *S3NATSFileStager) StageRemoteToStore(ctx context.Context, nodeID, remot
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReleaseRemote evicts one exact ephemeral key from the worker before deleting
|
||||
// the shared object.
|
||||
func (s *S3NATSFileStager) ReleaseRemote(ctx context.Context, nodeID, key string) error {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
reply, err := messaging.RequestJSON[fileReleaseRequest, fileReleaseReply](
|
||||
s.nats,
|
||||
messaging.SubjectNodeFilesRelease(nodeID),
|
||||
fileReleaseRequest{Key: key},
|
||||
30*time.Second,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if reply.Error != "" {
|
||||
return fmt.Errorf("backend release failed: %s", reply.Error)
|
||||
}
|
||||
if err := s.fm.Delete(ctx, key); err != nil {
|
||||
return fmt.Errorf("deleting shared object %q: %w", key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -119,6 +119,8 @@ func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDi
|
||||
handleUpload(w, r, stagingDir, modelsDir, dataDir, key, maxUploadSize)
|
||||
case http.MethodGet:
|
||||
handleDownload(w, r, stagingDir, modelsDir, dataDir, key)
|
||||
case http.MethodDelete:
|
||||
handleRelease(w, r, stagingDir, key)
|
||||
case http.MethodPost:
|
||||
if key == "temp" {
|
||||
handleAllocTemp(w, r, stagingDir)
|
||||
@@ -182,6 +184,57 @@ func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDi
|
||||
return server, nil
|
||||
}
|
||||
|
||||
func handleRelease(w http.ResponseWriter, _ *http.Request, stagingDir, key string) {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
filePath := filepath.Join(stagingDir, filepath.FromSlash(key))
|
||||
if err := validatePathInDir(filePath, stagingDir); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if info, err := os.Lstat(filePath); err == nil && info.IsDir() {
|
||||
http.Error(w, "release key identifies a directory", http.StatusBadRequest)
|
||||
return
|
||||
} else if err != nil && !os.IsNotExist(err) {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
for _, path := range []string{filePath, filePath + hashSidecarSuffix, filePath + targetSidecarSuffix} {
|
||||
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
}
|
||||
categoryDir := filepath.Dir(filePath)
|
||||
requestDir := filepath.Dir(categoryDir)
|
||||
for _, dir := range []string{categoryDir, requestDir} {
|
||||
if err := pruneEmptyDir(dir); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func pruneEmptyDir(dir string) error {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(entries) != 0 {
|
||||
return nil
|
||||
}
|
||||
if err := os.Remove(dir); err != nil && !os.IsNotExist(err) {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func handleHead(w http.ResponseWriter, r *http.Request, stagingDir, modelsDir, dataDir, key string) {
|
||||
if key == "" {
|
||||
http.Error(w, "key is required", http.StatusBadRequest)
|
||||
|
||||
@@ -52,6 +52,8 @@ func (f *fakeFileStager) AllocRemoteTemp(_ context.Context, _ string) (string, e
|
||||
|
||||
func (f *fakeFileStager) StageRemoteToStore(_ context.Context, _, _, _ string) error { return nil }
|
||||
|
||||
func (f *fakeFileStager) ReleaseRemote(_ context.Context, _, _ string) error { return nil }
|
||||
|
||||
func (f *fakeFileStager) ListRemoteDir(_ context.Context, _, _ string) ([]string, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
@@ -57,6 +58,9 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
|
||||
if err != nil {
|
||||
return fmt.Errorf("initializing file manager: %w", err)
|
||||
}
|
||||
if err := subscribeFileRelease(natsClient, nodeID, fm, cacheDir); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Subscribe: files.ensure — download S3 key to local, reply with local path
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesEnsure(nodeID), func(data []byte, reply func([]byte)) {
|
||||
@@ -199,3 +203,119 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
|
||||
xlog.Info("Subscribed to file staging NATS subjects", "nodeID", nodeID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func subscribeFileRelease(natsClient messaging.MessagingClient, nodeID string, fm *storage.FileManager, cacheDir string) error {
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesRelease(nodeID), func(data []byte, reply func([]byte)) {
|
||||
var req struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &req); err != nil {
|
||||
replyJSON(reply, map[string]string{"error": "invalid request"})
|
||||
return
|
||||
}
|
||||
cachePath, err := fm.CachePath(req.Key)
|
||||
if err == nil {
|
||||
err = releaseEphemeralCachePath(cacheDir, req.Key, cachePath)
|
||||
}
|
||||
if err != nil {
|
||||
replyJSON(reply, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
replyJSON(reply, map[string]string{})
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to files.release events: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func releaseEphemeralCacheKey(cacheDir, key string) error {
|
||||
return releaseEphemeralCachePath(cacheDir, key, filepath.Join(cacheDir, filepath.FromSlash(key)))
|
||||
}
|
||||
|
||||
func releaseEphemeralCachePath(cacheDir, key, filePath string) error {
|
||||
if err := validateEphemeralCacheKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateReleasePath(filePath, cacheDir); err != nil {
|
||||
return err
|
||||
}
|
||||
if info, err := os.Lstat(filePath); err == nil && info.IsDir() {
|
||||
return fmt.Errorf("release key identifies a directory")
|
||||
} else if err != nil && !os.IsNotExist(err) {
|
||||
return err
|
||||
}
|
||||
for _, path := range []string{filePath, filePath + ".sha256", filePath + ".sha256.target"} {
|
||||
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, dir := range []string{filepath.Dir(filePath), filepath.Dir(filepath.Dir(filePath))} {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if os.IsNotExist(err) {
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(entries) == 0 {
|
||||
if err := os.Remove(dir); err != nil && !os.IsNotExist(err) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateEphemeralCacheKey(key string) error {
|
||||
if strings.Contains(key, "\\") || path.Clean(key) != key {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
parts := strings.Split(key, "/")
|
||||
if len(parts) != 4 || parts[0] != "ephemeral" {
|
||||
return fmt.Errorf("release key %q must identify one file below ephemeral/", key)
|
||||
}
|
||||
for _, part := range parts[1:] {
|
||||
if part == "" || part == "." || part == ".." {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateReleasePath(targetPath, baseDir string) error {
|
||||
absBase, err := filepath.Abs(baseDir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
realBase, err := filepath.EvalSymlinks(absBase)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
absTarget, err := filepath.Abs(targetPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
realTarget, err := filepath.EvalSymlinks(absTarget)
|
||||
if err != nil {
|
||||
remaining := filepath.Base(absTarget)
|
||||
dir := filepath.Dir(absTarget)
|
||||
for {
|
||||
resolved, resolveErr := filepath.EvalSymlinks(dir)
|
||||
if resolveErr == nil {
|
||||
realTarget = filepath.Join(resolved, remaining)
|
||||
break
|
||||
}
|
||||
remaining = filepath.Join(filepath.Base(dir), remaining)
|
||||
parent := filepath.Dir(dir)
|
||||
if parent == dir {
|
||||
realTarget = filepath.Clean(absTarget)
|
||||
break
|
||||
}
|
||||
dir = parent
|
||||
}
|
||||
}
|
||||
if realTarget != realBase && !strings.HasPrefix(realTarget, realBase+string(filepath.Separator)) {
|
||||
return fmt.Errorf("path %q resolves outside ephemeral cache", targetPath)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type releaseSubscription struct{}
|
||||
|
||||
func (releaseSubscription) Unsubscribe() error { return nil }
|
||||
|
||||
type releaseMessagingClient struct {
|
||||
subject string
|
||||
handler func([]byte, func([]byte))
|
||||
}
|
||||
|
||||
func (m *releaseMessagingClient) Publish(string, any) error { return nil }
|
||||
func (m *releaseMessagingClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) SubscribeReply(subject string, handler func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
m.subject = subject
|
||||
m.handler = handler
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) Request(string, []byte, time.Duration) ([]byte, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) IsConnected() bool { return true }
|
||||
func (m *releaseMessagingClient) Close() {}
|
||||
|
||||
var _ = Describe("Worker exact-key staging release", func() {
|
||||
It("removes only the exact cache file and upload sidecars", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
target := filepath.Join(categoryDir, "input.wav")
|
||||
sibling := filepath.Join(categoryDir, "keep.wav")
|
||||
for _, path := range []string{target, target + ".sha256", target + ".sha256.target", sibling} {
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
}
|
||||
|
||||
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/input.wav")).To(Succeed())
|
||||
Expect(target).NotTo(BeAnExistingFile())
|
||||
Expect(target + ".sha256").NotTo(BeAnExistingFile())
|
||||
Expect(target + ".sha256.target").NotTo(BeAnExistingFile())
|
||||
Expect(sibling).To(BeAnExistingFile())
|
||||
Expect(categoryDir).To(BeADirectory())
|
||||
})
|
||||
|
||||
It("succeeds for a missing file and prunes empty category and request directories", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
|
||||
for range 2 {
|
||||
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/missing.wav")).To(Succeed())
|
||||
}
|
||||
Expect(categoryDir).NotTo(BeADirectory())
|
||||
Expect(filepath.Dir(categoryDir)).NotTo(BeADirectory())
|
||||
Expect(filepath.Join(cacheDir, "ephemeral")).To(BeADirectory())
|
||||
})
|
||||
|
||||
It("rejects traversal and symlink escapes", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
outsideDir := GinkgoT().TempDir()
|
||||
outsidePath := filepath.Join(outsideDir, "input.wav")
|
||||
Expect(os.WriteFile(outsidePath, []byte("keep"), 0640)).To(Succeed())
|
||||
requestDir := filepath.Join(cacheDir, "ephemeral", "request-id")
|
||||
Expect(os.MkdirAll(requestDir, 0750)).To(Succeed())
|
||||
Expect(os.Symlink(outsideDir, filepath.Join(requestDir, "audio"))).To(Succeed())
|
||||
|
||||
for _, key := range []string{
|
||||
"models/model.gguf",
|
||||
"ephemeral/../models/model.gguf",
|
||||
"ephemeral/request-id/audio/../../model.gguf",
|
||||
"ephemeral/request-id/audio/input.wav",
|
||||
} {
|
||||
Expect(releaseEphemeralCacheKey(cacheDir, key)).NotTo(Succeed(), key)
|
||||
}
|
||||
Expect(outsidePath).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("registers an exact release handler", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
path := filepath.Join(cacheDir, "ephemeral", "request-id", "audio", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
fm, err := storage.NewFileManager(nil, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
client := &releaseMessagingClient{}
|
||||
|
||||
Expect(subscribeFileRelease(client, "node.one", fm, cacheDir)).To(Succeed())
|
||||
Expect(client.subject).To(Equal(messaging.SubjectNodeFilesRelease("node.one")))
|
||||
request, err := json.Marshal(map[string]string{"key": "ephemeral/request-id/audio/input.wav"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
var response []byte
|
||||
client.handler(request, func(data []byte) { response = append([]byte(nil), data...) })
|
||||
|
||||
var reply map[string]string
|
||||
Expect(json.Unmarshal(response, &reply)).To(Succeed())
|
||||
Expect(reply["error"]).To(BeEmpty())
|
||||
Expect(path).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("returns validation errors through the release handler", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
fm, err := storage.NewFileManager(nil, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
client := &releaseMessagingClient{}
|
||||
Expect(subscribeFileRelease(client, "node-1", fm, cacheDir)).To(Succeed())
|
||||
|
||||
request, err := json.Marshal(map[string]string{"key": "models/model.gguf"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
var response []byte
|
||||
client.handler(request, func(data []byte) { response = append([]byte(nil), data...) })
|
||||
|
||||
var reply map[string]string
|
||||
Expect(json.Unmarshal(response, &reply)).To(Succeed())
|
||||
Expect(reply["error"]).NotTo(BeEmpty())
|
||||
})
|
||||
})
|
||||
Reference in new issue
Block a user