writer.go

package main

import (
	"context"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/google/osv.dev/go/logger"
	"github.com/google/osv.dev/go/osv/clients"
)

// writeMsg holds the data for a file to be written.
type writeMsg struct {
	path     string
	mimeType string
	data     []byte
}

// writer is a worker that receives writeMsgs and writes them to either a GCS
// bucket or a local directory.
func writer(ctx context.Context, cancel context.CancelFunc, inCh <-chan writeMsg, client clients.CloudStorage, pathPrefix string, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		select {
		case msg, ok := <-inCh:
			if !ok {
				// Channel closed.
				return
			}
			path := filepath.Join(pathPrefix, msg.path)
			if client != nil {
				// Write to the bucket.
				err := client.WriteObject(ctx, path, msg.data, &clients.WriteOptions{
					ContentType: msg.mimeType,
				})
				if err != nil {
					logger.Error("failed to write file", slog.String("path", path), slog.Any("err", err))
					cancel()

					break
				}
			} else {
				cleanPath := filepath.Clean(msg.path)
				if filepath.IsAbs(cleanPath) || cleanPath == ".." || strings.HasPrefix(cleanPath, ".."+string(filepath.Separator)) {
					logger.Error("invalid file path", slog.String("path", msg.path))
					cancel()

					break
				}
				basePath, err := filepath.Abs(pathPrefix)
				if err != nil {
					logger.Error("failed to get absolute path", slog.String("path", pathPrefix), slog.Any("err", err))
					cancel()

					break
				}
				localPath, err := filepath.Abs(filepath.Join(basePath, cleanPath))
				if err != nil {
					logger.Error("failed to get absolute path", slog.String("path", cleanPath), slog.Any("err", err))
					cancel()

					break
				}
				relPath, err := filepath.Rel(basePath, localPath)
				if err != nil || relPath == ".." || strings.HasPrefix(relPath, ".."+string(filepath.Separator)) {
					logger.Error("invalid file path", slog.String("path", msg.path), slog.Any("err", err))
					cancel()

					break
				}
				// Write locally.
				dir := filepath.Dir(localPath)
				if err := os.MkdirAll(dir, 0755); err != nil {
					logger.Error("failed to create directories", slog.String("dir", dir), slog.Any("err", err))
					cancel()

					break
				}
				if err := os.WriteFile(localPath, msg.data, 0600); err != nil {
					logger.Error("failed to write file", slog.String("path", localPath), slog.Any("err", err))
					cancel()

					break
				}
			}
		case <-ctx.Done():
			return
		}
	}
}