fix(embedding): extract ONNX runtime atomically via temp file and rename

This commit is contained in:
2026-09-30 01:27:50 +01:00
parent 443f557000
commit da1d96dab4
2 changed files with 77 additions and 5 deletions
+19 -5
View File
@@ -628,21 +628,35 @@ func (s *Service) Close() error {
} }
// writeGunzipped decompresses gz into path with the executable bit the dynamic linker needs. // writeGunzipped decompresses gz into path with the executable bit the dynamic linker needs.
func writeGunzipped(path string, gz []byte) error { // It writes a temp file and renames it, so a crash never leaves a truncated library at path.
func writeGunzipped(path string, gz []byte) (err error) {
zr, err := gzip.NewReader(bytes.NewReader(gz)) zr, err := gzip.NewReader(bytes.NewReader(gz))
if err != nil { if err != nil {
return err return err
} }
defer func() { _ = zr.Close() }() defer func() { _ = zr.Close() }()
// #nosec G302 G304 -- shared library needs 0755; path is under our own cache dir f, err := os.CreateTemp(filepath.Dir(path), filepath.Base(path)+".tmp*")
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
if err != nil { if err != nil {
return err return err
} }
if _, err := io.Copy(f, zr); err != nil { tmp := f.Name()
defer func() {
if err != nil {
_ = os.Remove(tmp)
}
}()
if _, err = io.Copy(f, zr); err != nil {
_ = f.Close() _ = f.Close()
return err return err
} }
return f.Close() if err = f.Close(); err != nil {
return err
}
// #nosec G302 -- shared library needs the executable bit for the dynamic linker
if err = os.Chmod(tmp, 0o755); err != nil {
return err
}
return os.Rename(tmp, path)
} }
+58
View File
@@ -1,9 +1,13 @@
package embedding package embedding
import ( import (
"bytes"
"compress/gzip"
"context" "context"
"errors" "errors"
"math" "math"
"os"
"path/filepath"
"sync" "sync"
"testing" "testing"
"time" "time"
@@ -475,3 +479,57 @@ func cosineSimilarity(a, b []float32) float64 {
return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB)) return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB))
} }
func gzipBytes(t *testing.T, data []byte) []byte {
t.Helper()
var buf bytes.Buffer
zw := gzip.NewWriter(&buf)
_, err := zw.Write(data)
require.NoError(t, err)
require.NoError(t, zw.Close())
return buf.Bytes()
}
func TestWriteGunzipped(t *testing.T) {
payload := bytes.Repeat([]byte("onnx"), 50_000)
good := gzipBytes(t, payload)
tests := []struct {
name string
gz []byte
wantErr bool
wantFile bool
}{
{"valid archive", good, false, true},
{"truncated archive leaves no file", good[:len(good)/2], true, false},
{"not gzip leaves no file", []byte("plain text"), true, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "lib.so")
err := writeGunzipped(path, tt.gz)
if tt.wantErr {
require.Error(t, err)
} else {
require.NoError(t, err)
}
_, statErr := os.Stat(path)
assert.Equal(t, tt.wantFile, statErr == nil)
entries, _ := os.ReadDir(dir)
if !tt.wantFile {
assert.Empty(t, entries, "no partial or temp files may remain")
}
if tt.wantFile {
got, err := os.ReadFile(path)
require.NoError(t, err)
assert.Equal(t, payload, got)
info, err := os.Stat(path)
require.NoError(t, err)
assert.Equal(t, os.FileMode(0o755), info.Mode().Perm())
}
})
}
}