cletus-plugins/sdk/cmd/cletus-plugin/workspace.go

package main

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path"
	"path/filepath"
	"strings"
)

const testWorkspaceName = "Test folder"

// workspace answers host.workspace.* from one directory, as one shared
// folder the user owns.
func (h *testHost) workspace(method, workspace, key, prefix, src, data string) (any, error) {
	if h.workspaceDir == "" {
		return nil, errors.New("no shared folders in the test host; pass -workspace <dir> to simulate one")
	}
	keyPath := func(key string) (string, error) {
		rel, ok := strings.CutPrefix(key, testWorkspaceID+"/files/")
		if !ok || rel == "" {
			return "", fmt.Errorf("%q is not a file in a shared folder (keys look like %s/files/<path>)", key, testWorkspaceID)
		}
		return h.workspaceFile(rel)
	}
	folder := func() error {
		if workspace != testWorkspaceID && workspace != testWorkspaceName {
			return fmt.Errorf("you are not a member of a shared folder %q", workspace)
		}
		return nil
	}
	switch method {
	case "host.workspace.folders":
		return []map[string]string{{"id": testWorkspaceID, "name": testWorkspaceName, "role": "owner"}}, nil
	case "host.workspace.list":
		if err := folder(); err != nil {
			return nil, err
		}
		var out []map[string]any
		err := filepath.WalkDir(h.workspaceDir, func(p string, d fs.DirEntry, err error) error {
			if err != nil || d.IsDir() || len(out) >= 1000 {
				return err
			}
			rel := filepath.ToSlash(mustRel(h.workspaceDir, p))
			if !strings.HasPrefix(rel, prefix) {
				return nil
			}
			info, err := d.Info()
			if err != nil {
				return err
			}
			etag, err := fileETag(p)
			if err != nil {
				return err
			}
			out = append(out, map[string]any{"key": testWorkspaceID + "/files/" + rel, "size": info.Size(), "etag": etag, "modified": info.ModTime().UTC()})
			return nil
		})
		return out, err
	case "host.workspace.fetch":
		file, err := keyPath(key)
		if err != nil {
			return nil, err
		}
		etag, err := fileETag(file)
		if err != nil {
			return nil, fmt.Errorf("%s does not exist", key)
		}
		// Like Cletus: a cached copy in the working folder's scratch area,
		// which host.files.resolve accepts.
		dst := filepath.Join(h.filesDir, ".cletus", "tmp", "workspace", filepath.FromSlash(strings.TrimPrefix(key, testWorkspaceID+"/files/")))
		if err := copyFile(file, dst); err != nil {
			return nil, err
		}
		return map[string]string{"path": dst, "etag": etag}, nil
	case "host.workspace.read":
		file, err := keyPath(key)
		if err != nil {
			return nil, err
		}
		info, err := os.Stat(file)
		if err != nil {
			return nil, fmt.Errorf("%s does not exist", key)
		}
		if info.Size() > 2<<20 {
			return nil, fmt.Errorf("%s is larger than 2 MiB; use host.workspace.fetch", key)
		}
		body, err := os.ReadFile(file)
		if err != nil {
			return nil, err
		}
		etag, _ := fileETag(file)
		return map[string]any{"data": base64.StdEncoding.EncodeToString(body), "etag": etag, "size": len(body)}, nil
	case "host.workspace.put":
		if err := folder(); err != nil {
			return nil, err
		}
		if key == "" || strings.Contains(key, "..") || strings.HasPrefix(key, "/") || strings.Contains(key, "//") {
			return nil, fmt.Errorf("key %q must be a relative path with no empty or .. segments", key)
		}
		dst, err := h.workspaceFile(key)
		if err != nil {
			return nil, err
		}
		var body []byte
		if src != "" {
			resolved, err := h.resolve(src)
			if err != nil {
				return nil, err
			}
			if body, err = os.ReadFile(resolved); err != nil {
				return nil, err
			}
		} else if body, err = base64.StdEncoding.DecodeString(data); err != nil {
			return nil, errors.New("data must be base64")
		}
		if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
			return nil, err
		}
		if err := os.WriteFile(dst, body, 0o644); err != nil {
			return nil, err
		}
		etag, _ := fileETag(dst)
		return map[string]any{"key": testWorkspaceID + "/files/" + key, "size": len(body), "etag": etag}, nil
	case "host.workspace.delete":
		file, err := keyPath(key)
		if err != nil {
			return nil, err
		}
		if err := os.Remove(file); err != nil && !os.IsNotExist(err) {
			return nil, err
		}
		return map[string]string{"key": key}, nil
	case "host.workspace.create":
		return nil, errors.New("the test host has one shared folder; it can't create another")
	}
	return nil, fmt.Errorf("unknown plugin host operation %q", method)
}

func (h *testHost) workspaceFile(rel string) (string, error) {
	file := filepath.Join(h.workspaceDir, filepath.FromSlash(path.Clean("/"+rel)))
	if !inside(h.workspaceDir, file) {
		return "", errors.New("the key is outside the shared folder")
	}
	return file, nil
}

func mustRel(root, p string) string {
	rel, _ := filepath.Rel(root, p)
	return rel
}

func fileETag(p string) (string, error) {
	f, err := os.Open(p)
	if err != nil {
		return "", err
	}
	defer f.Close()
	sum := sha256.New()
	if _, err := io.Copy(sum, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(sum.Sum(nil))[:32], nil
}

func copyFile(src, dst string) error {
	if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
		return err
	}
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		return err
	}
	return out.Close()
}