feat: minor enhancements to /models/apply (#297)

.github/workflows/image.yml

@@ -9,6 +9,10 @@ on:
     tags:
       - '*'
 
+concurrency:
+  group: ci-${{ github.head_ref || github.ref }}-${{ github.repository }}
+  cancel-in-progress: true
+
 jobs:
   docker:
     runs-on: ubuntu-latest

.github/workflows/test.yml

@@ -9,6 +9,10 @@ on:
     tags:
       - '*'
 
+concurrency:
+  group: ci-tests-${{ github.head_ref || github.ref }}-${{ github.repository }}
+  cancel-in-progress: true
+
 jobs:
   ubuntu-latest:
     runs-on: ubuntu-latest

api/gallery.go

@@ -86,6 +86,8 @@ func (g *galleryApplier) start(c context.Context, cm *ConfigMerger) {
 					continue
 				}
 
+				config.Files = append(config.Files, op.req.AdditionalFiles...)
+
 				if err := gallery.Apply(g.modelPath, op.req.Name, &config); err != nil {
 					updateError(err)
 					continue
@@ -108,6 +110,7 @@ func (g *galleryApplier) start(c context.Context, cm *ConfigMerger) {
 type ApplyGalleryModelRequest struct {
 	URL             string         `json:"url"`
 	Name            string         `json:"name"`
+	AdditionalFiles []gallery.File `json:"files"`
 }
 
 func getOpStatus(g *galleryApplier) func(c *fiber.Ctx) error {

pkg/gallery/models.go

@@ -50,9 +50,9 @@ type Config struct {
 }
 
 type File struct {
-	Filename string `yaml:"filename"`
-	SHA256   string `yaml:"sha256"`
-	URI      string `yaml:"uri"`
+	Filename string `yaml:"filename" json:"filename"`
+	SHA256   string `yaml:"sha256" json:"sha256"`
+	URI      string `yaml:"uri" json:"uri"`
 }
 
 type PromptTemplate struct {
@@ -77,6 +77,21 @@ func ReadConfigFile(filePath string) (*Config, error) {
 	return &config, nil
 }
 
+func inTrustedRoot(path string, trustedRoot string) error {
+	for path != "/" {
+		path = filepath.Dir(path)
+		if path == trustedRoot {
+			return nil
+		}
+	}
+	return fmt.Errorf("path is outside of trusted root")
+}
+
+func verifyPath(path, basePath string) error {
+	c := filepath.Clean(filepath.Join(basePath, path))
+	return inTrustedRoot(c, basePath)
+}
+
 func Apply(basePath, nameOverride string, config *Config) error {
 	// Create base path if it doesn't exist
 	err := os.MkdirAll(basePath, 0755)
@@ -88,6 +103,9 @@ func Apply(basePath, nameOverride string, config *Config) error {
 	for _, file := range config.Files {
 		log.Debug().Msgf("Checking %q exists and matches SHA", file.Filename)
 
+		if err := verifyPath(file.Filename, basePath); err != nil {
+			return err
+		}
 		// Create file path
 		filePath := filepath.Join(basePath, file.Filename)
 
@@ -173,6 +191,9 @@ func Apply(basePath, nameOverride string, config *Config) error {
 
 	// Write prompt template contents to separate files
 	for _, template := range config.PromptTemplates {
+		if err := verifyPath(template.Name+".tmpl", basePath); err != nil {
+			return err
+		}
 		// Create file path
 		filePath := filepath.Join(basePath, template.Name+".tmpl")
 
@@ -195,6 +216,10 @@ func Apply(basePath, nameOverride string, config *Config) error {
 		name = nameOverride
 	}
 
+	if err := verifyPath(name+".yaml", basePath); err != nil {
+		return err
+	}
+
 	configFilePath := filepath.Join(basePath, name+".yaml")
 
 	// Read and update config file as map[string]interface{}

pkg/gallery/models_test.go

@@ -26,5 +26,30 @@ var _ = Describe("Model test", func() {
 				Expect(err).ToNot(HaveOccurred())
 			}
 		})
+		It("renames model correctly", func() {
+			tempdir, err := os.MkdirTemp("", "test")
+			Expect(err).ToNot(HaveOccurred())
+			defer os.RemoveAll(tempdir)
+			c, err := ReadConfigFile(filepath.Join(os.Getenv("FIXTURES"), "gallery_simple.yaml"))
+			Expect(err).ToNot(HaveOccurred())
+
+			err = Apply(tempdir, "foo", c)
+			Expect(err).ToNot(HaveOccurred())
+
+			for _, f := range []string{"cerebras", "cerebras-completion.tmpl", "cerebras-chat.tmpl", "foo.yaml"} {
+				_, err = os.Stat(filepath.Join(tempdir, f))
+				Expect(err).ToNot(HaveOccurred())
+			}
+		})
+		It("catches path traversals", func() {
+			tempdir, err := os.MkdirTemp("", "test")
+			Expect(err).ToNot(HaveOccurred())
+			defer os.RemoveAll(tempdir)
+			c, err := ReadConfigFile(filepath.Join(os.Getenv("FIXTURES"), "gallery_simple.yaml"))
+			Expect(err).ToNot(HaveOccurred())
+
+			err = Apply(tempdir, "../../../foo", c)
+			Expect(err).To(HaveOccurred())
+		})
 	})
 })