AFI/afi_webdav.go

package main

import (
	"encoding/xml"
	"fmt"
	"io"
	"mime"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// ---- Minimal WebDAV 实现 ----

// 内存锁系统
type webdavLock struct {
	token   string
	path    string
	expires time.Time
	owner   string
}

type memLS struct {
	mu    sync.Mutex
	locks map[string]*webdavLock
}

func newMemLS() *memLS {
	return &memLS{locks: make(map[string]*webdavLock)}
}

func (ls *memLS) create(path, owner string, timeout time.Duration) string {
	ls.mu.Lock()
	defer ls.mu.Unlock()
	token := fmt.Sprintf("opaquelocktoken:%x", time.Now().UnixNano())
	ls.locks[token] = &webdavLock{
		token: token,
		path:  path,
		owner: owner,
	}
	if timeout > 0 {
		ls.locks[token].expires = time.Now().Add(timeout)
	}
	return token
}

func (ls *memLS) refresh(token string) {
	ls.mu.Lock()
	defer ls.mu.Unlock()
	if l, ok := ls.locks[token]; ok {
		l.expires = time.Now().Add(5 * time.Minute)
	}
}

func (ls *memLS) unlock(token string) bool {
	ls.mu.Lock()
	defer ls.mu.Unlock()
	_, ok := ls.locks[token]
	delete(ls.locks, token)
	return ok
}

func (ls *memLS) isLocked(path string) bool {
	ls.mu.Lock()
	defer ls.mu.Unlock()
	// 清理过期锁
	now := time.Now()
	for _, l := range ls.locks {
		if !l.expires.IsZero() && now.After(l.expires) {
			delete(ls.locks, l.token)
		}
	}
	// 检查路径本身及其所有父路径
	for _, l := range ls.locks {
		if l.path == path || strings.HasPrefix(path, l.path+"/") {
			return true
		}
	}
	return false
}

var lockSystem = newMemLS()

// checkLocked 如果路径被锁定则返回 423 Locked 响应，true 表示已锁定
func checkLocked(w http.ResponseWriter, p string) bool {
	if lockSystem.isLocked(p) {
		http.Error(w, "Resource is locked", http.StatusLocked)
		return true
	}
	return false
}

// ---- XML 结构 ----

type propfindXML struct {
	XMLName xml.Name `xml:"DAV: propfind"`
	Prop    propXML  `xml:"prop"`
}

type propXML struct {
	XMLName             xml.Name        `xml:"DAV: prop"`
	CreationDate        string          `xml:"DAV: creationdate,omitempty"`
	DisplayName         string          `xml:"DAV: displayname,omitempty"`
	GetContentLanguage  string          `xml:"DAV: getcontentlanguage,omitempty"`
	GetContentLength    string          `xml:"DAV: getcontentlength,omitempty"`
	GetContentType      string          `xml:"DAV: getcontenttype,omitempty"`
	GetETag             string          `xml:"DAV: getetag,omitempty"`
	GetLastModified     string          `xml:"DAV: getlastmodified,omitempty"`
	ResourceType        *resourceTypeXML `xml:"DAV: resourcetype,omitempty"`
	SupportedLock       string          `xml:"DAV: supportedlock,omitempty"`
	IsCollection        string          `xml:"DAV: iscollection,omitempty"`
	AFIComment          string          `xml:"AFI: comment,omitempty"`
}

type resourceTypeXML struct {
	Collection *struct{} `xml:"DAV: collection"`
}

type multistatusXML struct {
	XMLName   xml.Name      `xml:"DAV: multistatus"`
	Xmlns     string        `xml:"xmlns,attr"`
	XmlnsAFI  string        `xml:"xmlns:AFI,attr"`
	Responses []responseXML `xml:"response"`
}

type responseXML struct {
	XMLName   xml.Name      `xml:"DAV: response"`
	Href      string        `xml:"href"`
	Propstats []propstatXML `xml:"propstat"`
}

type propstatXML struct {
	XMLName xml.Name `xml:"DAV: propstat"`
	Status  string   `xml:"status"`
	Prop    propXML  `xml:"prop"`
}

type lockDiscoveryXML struct {
	XMLName    xml.Name       `xml:"DAV: lockdiscovery"`
	ActiveLock *activeLockXML `xml:"activelock,omitempty"`
}

type activeLockXML struct {
	LockType  lockTypeXML   `xml:"locktype"`
	LockScope lockScopeXML  `xml:"lockscope"`
	Depth     string        `xml:"depth"`
	Owner     *ownerXML     `xml:"owner,omitempty"`
	Timeout   string        `xml:"timeout"`
	LockToken *lockTokenXML `xml:"locktoken,omitempty"`
	LockRoot  string        `xml:"lockroot"`
}

type lockTypeXML struct {
	Write string `xml:"write"`
}

type lockScopeXML struct {
	Exclusive string `xml:"exclusive"`
}

type ownerXML struct {
	Href string `xml:"href"`
}

type lockTokenXML struct {
	Href string `xml:"href"`
}

// ---- WebDAV Handler ----

func webdavHandler(w http.ResponseWriter, r *http.Request) {
	// OPTIONS 提前放行（免认证）
	if r.Method == "OPTIONS" {
		w.Header().Set("DAV", "1")
		w.Header().Set("Allow", "GET, HEAD, PUT, DELETE, PROPFIND, MKCOL, MOVE, COPY, LOCK, UNLOCK, OPTIONS")
		w.WriteHeader(http.StatusOK)
		return
	}

	// PROPFIND / GET / HEAD 免认证，其余方法需要认证
	readMethod := r.Method == "PROPFIND" || r.Method == "GET" || r.Method == "HEAD"
	if !readMethod && !checkAuthRequest(w, r) {
		w.Header().Set("WWW-Authenticate", `Basic realm="AFI WebDAV"`)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	// 去掉前缀得到内部路径
	davPrefix := *extraPath + "dav"
	p := strings.TrimPrefix(r.URL.Path, davPrefix)
	if !strings.HasPrefix(p, "/") {
		p = "/" + p
	}
	fullPath := enforcePath(p)

	switch r.Method {
	case "PROPFIND":
		handlePropfind(w, r, p, fullPath)
	case "PROPPATCH":
		handleProppatch(w, r, p, fullPath)
	case "MKCOL":
		handleMkcol(w, r, p, fullPath)
	case "GET", "HEAD":
		handleGet(w, r, fullPath)
	case "PUT":
		handlePut(w, r, p, fullPath)
	case "DELETE":
		handleDelete(w, r, p, fullPath)
	case "MOVE":
		handleMove(w, r, p, fullPath, davPrefix)
	case "COPY":
		handleCopy(w, r, p, fullPath, davPrefix)
	case "LOCK":
		handleLock(w, r, p)
	case "UNLOCK":
		handleUnlock(w, r, p)
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

func handlePropfind(w http.ResponseWriter, r *http.Request, p, fullPath string) {
	stat, err := os.Stat(fullPath)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	depth := r.Header.Get("Depth")
	if depth == "" {
		depth = "1"
	}

	var entries []responseXML

	if stat.IsDir() && depth != "0" {
		if depth == "infinity" {
			// 递归遍历所有子目录（Depth: infinity）
			filepath.Walk(fullPath, func(walkPath string, walkInfo os.FileInfo, err error) error {
				if err != nil {
					return err
				}
				if walkPath == fullPath {
					return nil
				}
				if *skipHidden && strings.HasPrefix(filepath.Base(walkPath), ".") {
					if walkInfo.IsDir() {
						return filepath.SkipDir
					}
					return nil
				}
				if !*symlinks && walkInfo.Mode()&os.ModeSymlink != 0 {
					return nil
				}
				relPath, _ := filepath.Rel(fullPath, walkPath)
				href := path.Clean(p + "/" + filepath.ToSlash(relPath))
				if walkInfo.IsDir() && !strings.HasSuffix(href, "/") {
					href += "/"
				}
				// 从所在目录的 DESCRIPT.ION 读取注释
				dirDesc := parseDescriptIon(filepath.Dir(walkPath))
				comment := dirDesc[walkInfo.Name()]
				entries = append(entries, makePropResponse(href, walkInfo, walkInfo.IsDir(), comment))
				return nil
			})
		} else {
			// depth "1" - 只列出直接子项
			files, err := os.ReadDir(fullPath)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			dirDesc := parseDescriptIon(fullPath)
			for _, entry := range files {
				name := entry.Name()
				if *skipHidden && strings.HasPrefix(name, ".") {
					continue
				}
				childPath := path.Clean(p + "/" + name)
				childFull := filepath.Join(fullPath, name)
				childStat, err := os.Stat(childFull)
				if err != nil {
					continue
				}
				isDir := entry.IsDir()
				entries = append(entries, makePropResponse(childPath, childStat, isDir, dirDesc[name]))
			}
		}
	} else {
		isDir := stat.IsDir()
		href := path.Clean(p)
		if isDir && !strings.HasSuffix(href, "/") {
			href += "/"
		}
		parentDesc := parseDescriptIon(filepath.Dir(fullPath))
		entries = append(entries, makePropResponse(href, stat, isDir, parentDesc[stat.Name()]))
	}

	writeMultistatus(w, entries)
}

func makePropResponse(href string, fi os.FileInfo, isDir bool, comment string) responseXML {
	var rt *resourceTypeXML
	if isDir {
		rt = &resourceTypeXML{Collection: &struct{}{}}
	}

	size := ""
	if !isDir {
		size = fmt.Sprintf("%d", fi.Size())
	}

	mimeType := ""
	if !isDir {
		mimeType = "application/octet-stream"
		if ext := filepath.Ext(href); ext != "" {
			if t := mime.TypeByExtension(ext); t != "" {
				mimeType = t
			}
		}
	}

	etag := fmt.Sprintf(`"%x-%x"`, fi.Size(), fi.ModTime().UnixNano())
	lastMod := fi.ModTime().UTC().Format(http.TimeFormat)
	created := fi.ModTime().UTC().Format(time.RFC3339)

	return responseXML{
		Href: href,
		Propstats: []propstatXML{{
			Status: "HTTP/1.1 200 OK",
			Prop: propXML{
				CreationDate:     created,
				DisplayName:      path.Base(href),
				GetContentLength: size,
				GetContentType:   mimeType,
				GetETag:          etag,
				GetLastModified:  lastMod,
				ResourceType:     rt,
				IsCollection:     fmt.Sprintf("%t", isDir),
				AFIComment:       comment,
			},
		}},
	}
}

func writeMultistatus(w http.ResponseWriter, entries []responseXML) {
	w.Header().Set("Content-Type", "application/xml; charset=utf-8")
	w.WriteHeader(http.StatusMultiStatus)
	resp := multistatusXML{
		Xmlns:     "DAV:",
		XmlnsAFI:  "AFI:",
		Responses: entries,
	}
	enc := xml.NewEncoder(w)
	enc.Indent("", "  ")
	w.Write([]byte(xml.Header))
	enc.Encode(resp)
}

func handleMkcol(w http.ResponseWriter, _ *http.Request, p, fullPath string) {
	if checkLocked(w, p) {
		return
	}
	err := os.MkdirAll(fullPath, os.ModePerm)
	if err != nil {
		http.Error(w, err.Error(), http.StatusConflict)
		return
	}
	w.WriteHeader(http.StatusCreated)
}

func handlePut(w http.ResponseWriter, r *http.Request, p, fullPath string) {
	if checkLocked(w, p) {
		return
	}
	// Create parent dirs if needed
	os.MkdirAll(filepath.Dir(fullPath), os.ModePerm)

	f, err := os.Create(fullPath)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer f.Close()

	_, err = io.Copy(f, r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusCreated)
}

func handleDelete(w http.ResponseWriter, _ *http.Request, p, fullPath string) {
	if checkLocked(w, p) {
		return
	}
	err := os.RemoveAll(fullPath)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func handleMove(w http.ResponseWriter, r *http.Request, p, fullPath, davPrefix string) {
	if checkLocked(w, p) {
		return
	}
	dest := r.Header.Get("Destination")
	if dest == "" {
		http.Error(w, "Destination header required", http.StatusBadRequest)
		return
	}

	// 从 Destination URL 提取路径
	u, err := r.URL.Parse(dest)
	if err != nil {
		http.Error(w, "Invalid Destination header: "+err.Error(), http.StatusBadRequest)
		return
	}
	destPath := strings.TrimPrefix(u.Path, davPrefix)
	destFull := enforcePath(destPath)

	if checkLocked(w, destPath) {
		return
	}

	err = os.Rename(fullPath, destFull)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func handleCopy(w http.ResponseWriter, r *http.Request, p, fullPath, davPrefix string) {
	if checkLocked(w, p) {
		return
	}
	dest := r.Header.Get("Destination")
	if dest == "" {
		http.Error(w, "Destination header required", http.StatusBadRequest)
		return
	}

	u, err := r.URL.Parse(dest)
	if err != nil {
		http.Error(w, "Invalid Destination header: "+err.Error(), http.StatusBadRequest)
		return
	}
	destPath := strings.TrimPrefix(u.Path, davPrefix)
	destFull := enforcePath(destPath)

	if checkLocked(w, destPath) {
		return
	}

	err = copyFileOrDir(fullPath, destFull)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func handleGet(w http.ResponseWriter, r *http.Request, fullPath string) {
	stat, err := os.Stat(fullPath)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	if stat.IsDir() {
		http.Error(w, "Not a file", http.StatusMethodNotAllowed)
		return
	}

	file, err := os.Open(fullPath)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer file.Close()

	// ETag 基于文件大小和修改时间
	etag := fmt.Sprintf(`"%x-%x"`, stat.Size(), stat.ModTime().UnixNano())
	w.Header().Set("ETag", etag)

	// Content-Type 优先用 mime.TypeByExtension
	ext := filepath.Ext(stat.Name())
	if mimeType := mime.TypeByExtension(ext); mimeType != "" {
		w.Header().Set("Content-Type", mimeType)
	}

	// ServeContent 自动处理 Range/Partial Content (206), Content-Length, Last-Modified
	http.ServeContent(w, r, stat.Name(), stat.ModTime(), file)
}

// ---- DESCRIPT.ION 读写 ----

func writeDescriptIon(dirPath, name, description string) error {
	descPath := filepath.Join(dirPath, "DESCRIPT.ION")

	entries := make(map[string]string)
	if data, err := os.ReadFile(descPath); err == nil {
		for _, line := range strings.Split(string(data), "\n") {
			line = strings.TrimSpace(line)
			if line == "" {
				continue
			}
			var targetName, desc string
			if strings.HasPrefix(line, `"`) {
				end := strings.Index(line[1:], `"`)
				if end != -1 {
					targetName = line[1 : end+1]
					desc = strings.TrimSpace(line[end+2:])
				}
			} else {
				parts := strings.SplitN(line, " ", 2)
				targetName = parts[0]
				if len(parts) > 1 {
					desc = strings.TrimSpace(parts[1])
				}
			}
			if targetName != "" {
				cleanName := strings.Trim(targetName, `"`)
				entries[cleanName] = strings.TrimSpace(desc)
			}
		}
	}

	if description == "" {
		delete(entries, name)
	} else {
		entries[name] = description
	}

	// 排序后写入，保持确定性输出
	var names []string
	for n := range entries {
		names = append(names, n)
	}
	sort.Strings(names)

	var buf strings.Builder
	for _, target := range names {
		desc := entries[target]
		if strings.Contains(target, " ") {
			buf.WriteString(`"` + target + `" `)
		} else {
			buf.WriteString(target + " ")
		}
		buf.WriteString(desc + "\n")
	}

	return os.WriteFile(descPath, []byte(buf.String()), 0644)
}

// ---- PROPPATCH ----

func handleProppatch(w http.ResponseWriter, r *http.Request, p, fullPath string) {
	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	type afiProp struct {
		Comment string `xml:"AFI: comment"`
	}
	type setOrRemove struct {
		Prop afiProp `xml:"prop"`
	}
	type propertyupdateXML struct {
		XMLName xml.Name     `xml:"DAV: propertyupdate"`
		Set     *setOrRemove `xml:"set"`
		Remove  *setOrRemove `xml:"remove"`
	}

	var req propertyupdateXML
	if err := xml.Unmarshal(body, &req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	dirPath := filepath.Dir(fullPath)
	name := filepath.Base(fullPath)

	if req.Remove != nil {
		if err := writeDescriptIon(dirPath, name, ""); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	} else if req.Set != nil {
		if err := writeDescriptIon(dirPath, name, req.Set.Prop.Comment); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}

	// 构建响应 prop
	respProp := propXML{}
	if req.Set != nil {
		respProp.AFIComment = req.Set.Prop.Comment
	}

	// 返回 207 Multi-Status
	w.Header().Set("Content-Type", "application/xml; charset=utf-8")
	w.WriteHeader(http.StatusMultiStatus)
	resp := multistatusXML{
		Xmlns:    "DAV:",
		XmlnsAFI: "AFI:",
		Responses: []responseXML{{
			Href: p,
			Propstats: []propstatXML{{
				Status: "HTTP/1.1 200 OK",
				Prop:   respProp,
			}},
		}},
	}
	enc := xml.NewEncoder(w)
	enc.Indent("", "  ")
	w.Write([]byte(xml.Header))
	enc.Encode(resp)
}

func copyFileOrDir(src, dst string) error {
	si, err := os.Stat(src)
	if err != nil {
		return err
	}

	if si.IsDir() {
		return copyDir(src, dst)
	}

	return copyFile(src, dst)
}

func copyFile(src, dst string) error {
	os.MkdirAll(filepath.Dir(dst), os.ModePerm)

	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	return err
}

func copyDir(src, dst string) error {
	os.MkdirAll(dst, os.ModePerm)
	entries, err := os.ReadDir(src)
	if err != nil {
		return err
	}
	for _, e := range entries {
		if err := copyFileOrDir(filepath.Join(src, e.Name()), filepath.Join(dst, e.Name())); err != nil {
			return err
		}
	}
	return nil
}

func parseTimeout(s string) time.Duration {
	// RFC 4918: "Second-1234" or "Infinite"
	if strings.HasPrefix(s, "Second-") {
		n, err := strconv.Atoi(strings.TrimPrefix(s, "Second-"))
		if err == nil && n > 0 {
			return time.Duration(n) * time.Second
		}
	}
	return 5 * time.Minute // 默认 5 分钟
}

func handleLock(w http.ResponseWriter, r *http.Request, p string) {
	// 解析 body 获取 owner
	owner := ""
	if r.Body != nil {
		body, _ := io.ReadAll(r.Body)
		var lockReq struct {
			Owner struct {
				Href string `xml:"href"`
			} `xml:"owner"`
		}
		xml.Unmarshal(body, &lockReq)
		owner = lockReq.Owner.Href
	}

	timeout := parseTimeout(r.Header.Get("Timeout"))
	token := lockSystem.create(p, owner, timeout)

	timeoutStr := "Infinite"
	if timeout > 0 {
		timeoutStr = fmt.Sprintf("Second-%d", int(timeout.Seconds()))
	}

	w.Header().Set("Lock-Token", "<"+token+">")
	w.Header().Set("Content-Type", "application/xml; charset=utf-8")
	w.WriteHeader(http.StatusOK)

	resp := struct {
		XMLName      xml.Name       `xml:"DAV: prop"`
		LockDiscovery lockDiscoveryXML `xml:"lockdiscovery"`
	}{
		LockDiscovery: lockDiscoveryXML{
			ActiveLock: &activeLockXML{
				LockType:  lockTypeXML{Write: ""},
				LockScope: lockScopeXML{Exclusive: ""},
				Depth:     "Infinity",
				Owner:     &ownerXML{Href: owner},
				Timeout:   timeoutStr,
				LockToken: &lockTokenXML{Href: token},
				LockRoot:  p,
			},
		},
	}
	w.Write([]byte(xml.Header))
	enc := xml.NewEncoder(w)
	enc.Indent("", "  ")
	enc.Encode(resp)
}

func handleUnlock(w http.ResponseWriter, r *http.Request, p string) {
	token := r.Header.Get("Lock-Token")
	token = strings.Trim(token, "<>")
	if token == "" {
		http.Error(w, "Lock-Token header required", http.StatusBadRequest)
		return
	}
	if !lockSystem.unlock(token) {
		http.Error(w, "No matching lock found", http.StatusConflict)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}