bluebird/internal/app/download.go

package app

import (
	"database/sql"
	"errors"
	"io"
	"mime"
	"net/http"
	"strconv"
	"time"

	"filippo.io/age"
)

func (a *App) resolve(w http.ResponseWriter, r *http.Request) {
	if !a.allow(r, "pickup", 10) {
		fail(w, 429, "取件尝试过于频繁")
		return
	}
	var in struct {
		Code     string `json:"code"`
		Password string `json:"password"`
	}
	if !decode(w, r, &in) {
		return
	}
	code := normalizeCode(in.Code)
	if len(code) != 10 {
		fail(w, 404, "取件信息无效或已失效")
		return
	}
	p, e := a.policy()
	if e != nil {
		a.serverError(w, e)
		return
	}
	if !p.AnonymousDownload {
		fail(w, 403, "管理员已关闭取件")
		return
	}
	var id, hash string
	var version int
	now := a.now().Unix()
	e = a.db.QueryRow("SELECT id,password_hash,version FROM parcels WHERE code_hash=? AND state='ready' AND EXISTS(SELECT 1 FROM files WHERE parcel=parcels.id AND status='ready' AND remaining>0 AND expires>?)", a.keyed(code), now).Scan(&id, &hash, &version)
	if e != nil {
		if !errors.Is(e, sql.ErrNoRows) {
			a.serverError(w, e)
			return
		}
		fail(w, 404, "取件信息无效或已失效")
		return
	}
	if hash != "" {
		if in.Password == "" {
			respond(w, 401, map[string]any{"error": "请输入取件密码", "password_required": true})
			return
		}
		select {
		case a.passwordSlots <- struct{}{}:
		default:
			fail(w, 429, "密码验证繁忙")
			return
		}
		valid := verifyPassword(in.Password, hash)
		<-a.passwordSlots
		if !valid {
			fail(w, 404, "取件信息无效或已失效")
			return
		}
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	var current int
	e = a.db.QueryRow("SELECT version FROM parcels WHERE id=? AND state='ready' AND code_hash=?", id, a.keyed(code)).Scan(&current)
	if e != nil || current != version {
		fail(w, 404, "取件信息无效或已失效")
		return
	}
	var count int
	if e = a.db.QueryRow("SELECT count(*) FROM grants").Scan(&count); e != nil {
		a.serverError(w, e)
		return
	}
	if count >= a.cfg.MaxRecords {
		fail(w, 429, "取件会话繁忙")
		return
	}
	token := randomToken()
	if _, e = a.db.Exec("INSERT INTO grants VALUES(?,?,?,?)", tokenHash(token), id, version, now+600); e != nil {
		a.serverError(w, e)
		return
	}
	a.cookie(w, r, "bb_pickup", token, "/api/v1/pickup", 600)
	a.pickup(w, r.Clone(r.Context()), id)
}

// Explicit parcel IDs are only used immediately after a validated code exchange.
func (a *App) grant(r *http.Request) (string, error) {
	var id string
	e := a.db.QueryRow("SELECT p.id FROM grants g JOIN parcels p ON p.id=g.parcel WHERE g.token_hash=? AND g.expires>? AND p.state='ready' AND p.version=g.version", tokenHash(cookieValue(r, "bb_pickup")), a.now().Unix()).Scan(&id)
	return id, e
}
func (a *App) fileList(id string) ([]map[string]any, error) {
	rows, e := a.db.Query("SELECT id,name,size,status,remaining,expires FROM files WHERE parcel=? ORDER BY rowid", id)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	out := []map[string]any{}
	for rows.Next() {
		var fid, name, status string
		var size, expires int64
		var remaining int
		if e = rows.Scan(&fid, &name, &size, &status, &remaining, &expires); e != nil {
			return nil, e
		}
		out = append(out, map[string]any{"id": fid, "name": name, "size": size, "status": status, "remaining": remaining, "expires": expires, "available": status == "ready" && remaining > 0 && expires > a.now().Unix()})
	}
	return out, rows.Err()
}
func (a *App) pickup(w http.ResponseWriter, r *http.Request, explicit ...string) {
	id := ""
	if len(explicit) > 0 {
		id = explicit[0]
	} else {
		var e error
		id, e = a.grant(r)
		if e != nil {
			fail(w, 401, "取件授权已失效，请重新输入取件码")
			return
		}
	}
	p, e := a.policy()
	if e != nil {
		a.serverError(w, e)
		return
	}
	if !p.AnonymousDownload {
		fail(w, 403, "管理员已关闭取件")
		return
	}
	files, e := a.fileList(id)
	if e != nil {
		a.serverError(w, e)
		return
	}
	var renewal int
	var mode string
	if e = a.db.QueryRow("SELECT renewal_days,renewal_mode FROM parcels WHERE id=?", id).Scan(&renewal, &mode); e != nil {
		a.serverError(w, e)
		return
	}
	respond(w, 200, map[string]any{"id": id, "files": files, "renewal_days": renewal, "renewal_mode": mode})
}
func (a *App) pickupHandler(w http.ResponseWriter, r *http.Request) { a.pickup(w, r) }

type startWriter struct {
	dst        http.ResponseWriter
	start      func() error
	started    bool
	bytes      int64
	controller *http.ResponseController
	total      time.Time
	idle       time.Duration
}

func (w *startWriter) Write(b []byte) (int, error) {
	if !w.started {
		if e := w.start(); e != nil {
			return 0, e
		}
		w.started = true
	}
	if w.controller != nil {
		end := time.Now().Add(w.idle)
		if end.After(w.total) {
			end = w.total
		}
		_ = w.controller.SetWriteDeadline(end)
	}
	n, e := w.dst.Write(b)
	w.bytes += int64(n)
	return n, e
}
func (a *App) download(w http.ResponseWriter, r *http.Request) {
	ip := a.remoteIP(r)
	if !a.acquire("download", ip) {
		fail(w, 429, "下载并发已满")
		return
	}
	defer a.release("download", ip)
	a.deadlines(w, r)
	a.mu.Lock()
	parcel, e := a.grant(r)
	if e != nil {
		a.mu.Unlock()
		fail(w, 401, "取件授权失效")
		return
	}
	policy, e := a.policy()
	if e != nil {
		a.mu.Unlock()
		a.serverError(w, e)
		return
	}
	if !policy.AnonymousDownload {
		a.mu.Unlock()
		fail(w, 403, "取件已关闭")
		return
	}
	fid := r.PathValue("file")
	var oid, name, status string
	var size, expiry int64
	var remaining int
	var encrypted bool
	e = a.db.QueryRow("SELECT f.object_id,f.name,f.size,f.status,f.remaining,f.expires,o.encrypted FROM files f JOIN objects o ON o.id=f.object_id WHERE f.id=? AND f.parcel=?", fid, parcel).Scan(&oid, &name, &size, &status, &remaining, &expiry, &encrypted)
	if e != nil || status != "ready" || remaining < 1 || expiry <= a.now().Unix() {
		a.mu.Unlock()
		fail(w, 404, "文件已失效")
		return
	}
	file, e := a.root.Open(objectPath(oid))
	if e != nil {
		a.mu.Unlock()
		fail(w, 404, "文件不可用")
		return
	}
	a.activeObjects[oid]++
	defer func() {
		a.mu.Lock()
		_ = file.Close()
		a.activeObjects[oid]--
		if a.activeObjects[oid] == 0 {
			delete(a.activeObjects, oid)
		}
		a.mu.Unlock()
	}()
	var reader io.Reader = file
	if encrypted {
		if a.identity == nil {
			a.mu.Unlock()
			a.serverError(w, errors.New("缺少私钥"))
			return
		}
		reader, e = age.Decrypt(file, a.identity)
		if e != nil {
			a.mu.Unlock()
			a.serverError(w, e)
			return
		}
	}
	if r.Method == "HEAD" {
		a.mu.Unlock()
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Header().Set("Content-Length", strconv.FormatInt(size, 10))
		w.WriteHeader(200)
		return
	}
	attempt := randomID()
	tx, e := a.db.Begin()
	if e != nil {
		a.mu.Unlock()
		a.serverError(w, e)
		return
	}
	result, e := tx.Exec("UPDATE files SET remaining=remaining-1 WHERE id=? AND remaining>0 AND status='ready' AND expires>?", fid, a.now().Unix())
	if e == nil {
		n, _ := result.RowsAffected()
		if n != 1 {
			e = errors.New("次数或期限不足")
		}
	}
	if e == nil {
		_, e = tx.Exec("INSERT INTO downloads(id,file_id,created) VALUES(?,?,?)", attempt, fid, a.now().Unix())
	}
	if e == nil {
		e = tx.Commit()
	} else {
		tx.Rollback()
	}
	if e != nil {
		a.mu.Unlock()
		a.serverError(w, e)
		return
	}
	a.mu.Unlock()
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": name}))
	w.Header().Set("Content-Length", strconv.FormatInt(size, 10))
	w.Header().Set("Accept-Ranges", "none")
	writer := &startWriter{dst: w, controller: http.NewResponseController(w), total: time.Now().Add(time.Duration(a.cfg.MaxTransferSeconds) * time.Second), idle: time.Duration(a.cfg.TransferIdleSeconds) * time.Second, start: func() error { _, e := a.db.Exec("UPDATE downloads SET started=1 WHERE id=?", attempt); return e }}
	// A zero-byte download is still an accepted download and must consume a count.
	if size == 0 {
		var b [1]byte
		n, re := reader.Read(b[:])
		if n != 0 || re != io.EOF {
			e = errors.New("空文件内容异常")
		}
		if e == nil {
			e = writer.start()
		}
		if e == nil {
			writer.started = true
			w.WriteHeader(200)
		}
	} else {
		_, e = io.CopyBuffer(writer, io.LimitReader(reader, size+1), make([]byte, 128<<10))
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	if !writer.started {
		a.refund(attempt, fid)
		w.Header().Del("Content-Length")
		w.Header().Del("Content-Disposition")
		a.serverError(w, errors.New("下载开始前失败"))
		return
	}
	if e == nil && writer.bytes == size {
		if e = a.completeDownload(attempt, fid, parcel); e != nil {
			a.log.Error("download completion failed", "error", e)
		}
	} else {
		a.log.Warn("download interrupted", "file", fid, "bytes", writer.bytes)
	}
}
func (a *App) refund(attempt, fid string) {
	tx, e := a.db.Begin()
	if e != nil {
		return
	}
	defer tx.Rollback()
	res, e := tx.Exec("UPDATE downloads SET refunded=1 WHERE id=? AND started=0 AND refunded=0", attempt)
	if e != nil {
		return
	}
	n, _ := res.RowsAffected()
	if n == 1 {
		_, e = tx.Exec("UPDATE files SET remaining=remaining+1 WHERE id=?", fid)
	}
	if e == nil {
		_ = tx.Commit()
	}
}
func (a *App) completeDownload(attempt, fid, parcel string) error {
	tx, e := a.db.Begin()
	if e != nil {
		return e
	}
	defer tx.Rollback()
	res, e := tx.Exec("UPDATE downloads SET completed=1 WHERE id=? AND completed=0", attempt)
	if e != nil {
		return e
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return tx.Commit()
	}
	var days, remaining int
	var mode, state, status string
	var expiry, hard int64
	e = tx.QueryRow("SELECT p.renewal_days,p.renewal_mode,p.state,f.status,f.expires,f.hard_expires,f.remaining FROM files f JOIN parcels p ON p.id=f.parcel WHERE f.id=?", fid).Scan(&days, &mode, &state, &status, &expiry, &hard, &remaining)
	if e != nil {
		return e
	}
	if state == "ready" && status == "ready" && remaining > 0 && days > 0 {
		newExpiry := max(expiry, a.now().Unix()+int64(days)*86400)
		if mode == "add" {
			newExpiry = max(expiry, a.now().Unix()) + int64(days)*86400
		}
		newExpiry = min(newExpiry, hard)
		if newExpiry > expiry {
			if _, e = tx.Exec("UPDATE files SET expires=? WHERE id=?", newExpiry, fid); e != nil {
				return e
			}
			if _, e = tx.Exec("UPDATE parcels SET expires=(SELECT max(expires) FROM files WHERE parcel=?) WHERE id=?", parcel, parcel); e != nil {
				return e
			}
		}
	}
	return tx.Commit()
}