drive/backend/handlers/webdav.go

package handlers

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"git.elijahkuntz.com/Elijah/drive/config"
	"git.elijahkuntz.com/Elijah/drive/database"
	"git.elijahkuntz.com/Elijah/drive/middleware"
	"github.com/gofiber/fiber/v2"
	"github.com/valyala/fasthttp/fasthttpadaptor"
	"golang.org/x/net/webdav"
)

type statusCapture struct {
	http.ResponseWriter
	statusCode int
}

func (sc *statusCapture) WriteHeader(code int) {
	sc.statusCode = code
	sc.ResponseWriter.WriteHeader(code)
}

type WebDAVHandler struct {
	Handler *webdav.Handler
	DB      *database.DB
	Config  *config.Config
	Guard   *middleware.BruteForceGuard
	
	authCache map[string]time.Time
	cacheMu   sync.RWMutex
}

func NewWebDAVHandler(db *database.DB, cfg *config.Config, guard *middleware.BruteForceGuard) *WebDAVHandler {
	fs := webdav.Dir(cfg.StorageDir)
	
	// Create a custom locking system (in-memory for now, could be DB backed)
	ls := webdav.NewMemLS()

	h := &webdav.Handler{
		Prefix:     "/webdav",
		FileSystem: fs,
		LockSystem: ls,
		Logger: func(r *http.Request, err error) {
			if err != nil {
				fmt.Printf("WebDAV [%s] %s: %v\n", r.Method, r.URL.Path, err)
			}
		},
	}

	return &WebDAVHandler{
		Handler:   h,
		DB:        db,
		Config:    cfg,
		Guard:     guard,
		authCache: make(map[string]time.Time),
	}
}

func (h *WebDAVHandler) Handle(c *fiber.Ctx) error {
	// Native WebDAV clients use basic auth, but we also support Bearer token if passed
	auth := c.Get("Authorization")
	var user, pass string
	var hasBasicAuth bool

	if strings.HasPrefix(auth, "Basic ") {
		if payload, err := base64.StdEncoding.DecodeString(auth[6:]); err == nil {
			parts := strings.SplitN(string(payload), ":", 2)
			if len(parts) == 2 {
				user = parts[0]
				pass = parts[1]
				hasBasicAuth = true
			}
		}
	}

	if hasBasicAuth {
		hasher := sha256.Sum256([]byte(user + ":" + pass))
		cacheKey := hex.EncodeToString(hasher[:])
		
		h.cacheMu.RLock()
		expireTime, ok := h.authCache[cacheKey]
		h.cacheMu.RUnlock()

		if !ok || time.Now().After(expireTime) {
			// Verify basic auth against DB
			var storedHash string
			err := h.DB.QueryRow("SELECT password_hash FROM users WHERE username = ?", user).Scan(&storedHash)
			if err != nil || !verifyPassword(pass, storedHash) {
				h.Guard.RecordFailure(c.IP())
				c.Set("WWW-Authenticate", `Basic realm="Drive WebDAV"`)
				return c.SendStatus(fiber.StatusUnauthorized)
			}
			h.Guard.RecordSuccess(c.IP())
			
			// Cache successful auth for 5 minutes
			h.cacheMu.Lock()
			// Bounded cache to prevent DoS
			if len(h.authCache) >= 1000 {
				now := time.Now()
				for k, v := range h.authCache {
					if now.After(v) {
						delete(h.authCache, k)
					}
				}
			}
			if len(h.authCache) < 1000 {
				h.authCache[cacheKey] = time.Now().Add(5 * time.Minute)
			}
			h.cacheMu.Unlock()
		}
	} else if c.Locals("userID") == nil {
		c.Set("WWW-Authenticate", `Basic realm="Drive WebDAV"`)
		return c.SendStatus(fiber.StatusUnauthorized)
	}

	// Adapt Fiber's FastHTTP to standard net/http for the WebDAV handler
	fasthttpadaptor.NewFastHTTPHandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		relPath := strings.TrimPrefix(r.URL.Path, "/webdav")
		
		capture := &statusCapture{ResponseWriter: w, statusCode: 200}
		h.Handler.ServeHTTP(capture, r)
		
		if r.Method == "PUT" && capture.statusCode >= 200 && capture.statusCode < 300 {
			h.syncToDB(relPath)
		}
	})(c.Context())

	return nil
}

// syncToDB manually synchronizes a file written via WebDAV to the SQLite database
func (h *WebDAVHandler) syncToDB(relPath string) {
	fullPath := filepath.Join(h.Config.StorageDir, filepath.FromSlash(relPath))
	info, err := os.Stat(fullPath)
	if err != nil {
		return
	}

	checksum, _ := generateChecksum(fullPath) // Reusing the helper from files.go
	cleanPath := filepath.ToSlash(relPath)
	
	mimeType := "application/octet-stream"
	if ext := filepath.Ext(cleanPath); ext != "" {
		if m := mime.TypeByExtension(ext); m != "" {
			mimeType = m
		}
	}

	h.DB.Exec(`
		INSERT INTO files (path, name, size, mime_type, checksum, is_dir)
		VALUES (?, ?, ?, ?, ?, ?)
		ON CONFLICT(path) DO UPDATE SET
			size = excluded.size,
			checksum = excluded.checksum,
			updated_at = CURRENT_TIMESTAMP
	`, cleanPath, filepath.Base(cleanPath), info.Size(), mimeType, checksum, info.IsDir())
}