model_test.go 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. package server
  2. import (
  3. "archive/zip"
  4. "bytes"
  5. "errors"
  6. "io"
  7. "os"
  8. "path/filepath"
  9. "slices"
  10. "strings"
  11. "testing"
  12. "github.com/ollama/ollama/api"
  13. )
  14. func createZipFile(t *testing.T, name string) *os.File {
  15. t.Helper()
  16. f, err := os.CreateTemp(t.TempDir(), "")
  17. if err != nil {
  18. t.Fatal(err)
  19. }
  20. zf := zip.NewWriter(f)
  21. defer zf.Close()
  22. zh, err := zf.CreateHeader(&zip.FileHeader{Name: name})
  23. if err != nil {
  24. t.Fatal(err)
  25. }
  26. if _, err := io.Copy(zh, bytes.NewReader([]byte(""))); err != nil {
  27. t.Fatal(err)
  28. }
  29. return f
  30. }
  31. func TestExtractFromZipFile(t *testing.T) {
  32. cases := []struct {
  33. name string
  34. expect []string
  35. err error
  36. }{
  37. {
  38. name: "good",
  39. expect: []string{"good"},
  40. },
  41. {
  42. name: strings.Join([]string{"path", "..", "to", "good"}, string(os.PathSeparator)),
  43. expect: []string{filepath.Join("to", "good")},
  44. },
  45. {
  46. name: strings.Join([]string{"path", "..", "to", "..", "good"}, string(os.PathSeparator)),
  47. expect: []string{"good"},
  48. },
  49. {
  50. name: strings.Join([]string{"path", "to", "..", "..", "good"}, string(os.PathSeparator)),
  51. expect: []string{"good"},
  52. },
  53. {
  54. name: strings.Join([]string{"..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "bad"}, string(os.PathSeparator)),
  55. err: zip.ErrInsecurePath,
  56. },
  57. {
  58. name: strings.Join([]string{"path", "..", "..", "to", "bad"}, string(os.PathSeparator)),
  59. err: zip.ErrInsecurePath,
  60. },
  61. }
  62. for _, tt := range cases {
  63. t.Run(tt.name, func(t *testing.T) {
  64. f := createZipFile(t, tt.name)
  65. defer f.Close()
  66. tempDir := t.TempDir()
  67. if err := extractFromZipFile(tempDir, f, func(api.ProgressResponse) {}); !errors.Is(err, tt.err) {
  68. t.Fatal(err)
  69. }
  70. var matches []string
  71. if err := filepath.Walk(tempDir, func(p string, fi os.FileInfo, err error) error {
  72. if err != nil {
  73. return err
  74. }
  75. if !fi.IsDir() {
  76. matches = append(matches, p)
  77. }
  78. return nil
  79. }); err != nil {
  80. t.Fatal(err)
  81. }
  82. var actual []string
  83. for _, match := range matches {
  84. rel, err := filepath.Rel(tempDir, match)
  85. if err != nil {
  86. t.Error(err)
  87. }
  88. actual = append(actual, rel)
  89. }
  90. if !slices.Equal(actual, tt.expect) {
  91. t.Fatalf("expected %d files, got %d", len(tt.expect), len(matches))
  92. }
  93. })
  94. }
  95. }