mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
fix(embedding): extract ONNX runtime atomically via temp file and rename
This commit is contained in:
@@ -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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user