Initial commit: docker compose config

server/tgbot/upload/manager.go

@@ -0,0 +1,288 @@
+package upload
+
+import (
+	"errors"
+	"fmt"
+	"math"
+	"path/filepath"
+	"strconv"
+	"sync"
+	"time"
+
+	tele "gopkg.in/telebot.v4"
+
+	"server/log"
+	"server/torr"
+	"server/torr/state"
+)
+
+// TrFunc is set by tgbot for localization (avoids circular import)
+var TrFunc func(int64, string) string
+
+// EscapeFunc is set by tgbot for HTML escaping (avoids circular import)
+var EscapeFunc func(string) string
+
+func tr(uid int64, key string) string {
+	if TrFunc != nil {
+		return TrFunc(uid, key)
+	}
+	return key
+}
+
+func escapeHtml(s string) string {
+	if EscapeFunc != nil {
+		return EscapeFunc(s)
+	}
+	return s
+}
+
+type Worker struct {
+	id          int
+	c           tele.Context
+	msg         *tele.Message
+	torrentHash string
+	isCancelled bool
+	from        int
+	to          int
+	ti          *state.TorrentStatus
+}
+
+type Manager struct {
+	queue     []*Worker
+	working   map[int]*Worker
+	ids       int
+	wrkSync   sync.Mutex
+	queueLock sync.Mutex
+}
+
+func (m *Manager) Start() {
+	m.working = make(map[int]*Worker)
+	go m.work()
+}
+
+func (m *Manager) AddRange(c tele.Context, hash string, from, to int) {
+	m.queueLock.Lock()
+	defer m.queueLock.Unlock()
+
+	if len(m.queue) > 50 {
+		c.Bot().Send(c.Recipient(), fmt.Sprintf(tr(c.Sender().ID, "upload_queue_full"), len(m.queue)))
+		return
+	}
+
+	m.ids++
+	if m.ids > math.MaxInt {
+		m.ids = 0
+	}
+
+	var msg *tele.Message
+	var err error
+
+	for i := 0; i < 20; i++ {
+		msg, err = c.Bot().Send(c.Recipient(), fmt.Sprintf(tr(c.Sender().ID, "upload_connecting"), hash))
+		if err == nil {
+			break
+		}
+		log.TLogln("tg upload retry", i+1, "/", 20)
+		if i < 19 {
+			backoff := time.Duration(1<<uint(i)) * 100 * time.Millisecond
+			if backoff > 5*time.Second {
+				backoff = 5 * time.Second
+			}
+			time.Sleep(backoff)
+		}
+	}
+
+	if err != nil {
+		log.TLogln("tg upload send err", err)
+		return
+	}
+
+	t := torr.GetTorrent(hash)
+	if t == nil {
+		c.Bot().Edit(msg, tr(c.Sender().ID, "torrent_not_found")+":\n<code>"+hash+"</code>")
+		return
+	}
+	t.WaitInfo()
+	for t.Status().Stat != state.TorrentWorking {
+		time.Sleep(time.Second)
+		t = torr.GetTorrent(hash)
+		if t == nil {
+			return
+		}
+	}
+	ti := t.Status()
+
+	if from == 1 && to == -1 {
+		to = len(ti.FileStats)
+	}
+	if from < 1 {
+		from = 1
+	}
+	if to > len(ti.FileStats) {
+		to = len(ti.FileStats)
+	}
+	if from > to {
+		from, to = to, from
+	}
+	if to > len(ti.FileStats) {
+		to = len(ti.FileStats)
+	}
+
+	w := &Worker{
+		id:          m.ids,
+		c:           c,
+		torrentHash: hash,
+		msg:         msg,
+		ti:          ti,
+		from:        from,
+		to:          to,
+	}
+
+	m.queue = append(m.queue, w)
+}
+
+func (m *Manager) Cancel(id int) {
+	m.queueLock.Lock()
+	defer m.queueLock.Unlock()
+	for i, w := range m.queue {
+		if w.id == id {
+			w.isCancelled = true
+			w.c.Bot().Delete(w.msg)
+			m.queue = append(m.queue[:i], m.queue[i+1:]...)
+			return
+		}
+	}
+	if wrk, ok := m.working[id]; ok {
+		wrk.isCancelled = true
+		return
+	}
+}
+
+func (m *Manager) work() {
+	for {
+		m.queueLock.Lock()
+		if len(m.working) > 0 {
+			m.queueLock.Unlock()
+			m.sendQueueStatus()
+			time.Sleep(time.Second)
+			continue
+		}
+		if len(m.queue) == 0 {
+			m.queueLock.Unlock()
+			time.Sleep(time.Second)
+			continue
+		}
+		wrk := m.queue[0]
+		m.queue = m.queue[1:]
+		m.working[wrk.id] = wrk
+		m.queueLock.Unlock()
+
+		m.sendQueueStatus()
+
+		loading(wrk)
+
+		m.queueLock.Lock()
+		delete(m.working, wrk.id)
+		m.queueLock.Unlock()
+	}
+}
+
+func (m *Manager) sendQueueStatus() {
+	m.queueLock.Lock()
+	defer m.queueLock.Unlock()
+	for i, wrk := range m.queue {
+		if wrk.msg == nil || wrk.c.Sender() == nil {
+			continue
+		}
+		torrKbd := &tele.ReplyMarkup{}
+		torrKbd.Inline([]tele.Row{torrKbd.Row(torrKbd.Data(tr(wrk.c.Sender().ID, "upload_cancel"), "cancel", strconv.Itoa(wrk.id)))}...)
+
+		msg := fmt.Sprintf(tr(wrk.c.Sender().ID, "upload_queue_pos"), i+1)
+
+		wrk.c.Bot().Edit(wrk.msg, msg, torrKbd)
+	}
+}
+
+func loading(wrk *Worker) {
+	iserr := false
+
+	t := torr.GetTorrent(wrk.torrentHash)
+	if t == nil {
+		wrk.c.Bot().Edit(wrk.msg, tr(wrk.c.Sender().ID, "torrent_not_found")+":\n<code>"+wrk.torrentHash+"</code>")
+		return
+	}
+	t.WaitInfo()
+	for t.Status().Stat != state.TorrentWorking {
+		time.Sleep(time.Second)
+		t = torr.GetTorrent(wrk.torrentHash)
+		if t == nil {
+			return
+		}
+	}
+	wrk.ti = t.Status()
+
+	for i := wrk.from - 1; i <= wrk.to-1; i++ {
+		file := wrk.ti.FileStats[i]
+		if wrk.isCancelled {
+			return
+		}
+
+		err := uploadFile(wrk, file, i+1, len(wrk.ti.FileStats))
+		if err != nil {
+			errstr := fmt.Sprintf(tr(wrk.c.Sender().ID, "upload_error"), err)
+			wrk.c.Bot().Edit(wrk.msg, errstr)
+			iserr = true
+			break
+		}
+	}
+	if !iserr {
+		wrk.c.Bot().Delete(wrk.msg)
+	}
+}
+
+func uploadFile(wrk *Worker, file *state.TorrentFileStat, fi, fc int) error {
+	caption := filepath.Base(file.Path)
+	torrFile, err := NewTorrFile(wrk, file)
+	if err != nil {
+		return err
+	}
+
+	var wa sync.WaitGroup
+	wa.Add(1)
+	complete := false
+	go func() {
+		for !complete {
+			updateLoadStatus(wrk, torrFile, fi, fc)
+			time.Sleep(1 * time.Second)
+		}
+		wa.Done()
+	}()
+
+	d := &tele.Document{}
+	d.FileName = file.Path
+	d.Caption = caption
+	d.File.FileReader = torrFile
+
+	for i := 0; i < 20; i++ {
+		err = wrk.c.Send(d)
+		if err == nil || errors.Is(err, ERR_STOPPED) {
+			break
+		}
+		log.TLogln("tg upload retry", i+1, "/", 20)
+		if i < 19 {
+			backoff := time.Duration(1<<uint(i)) * 100 * time.Millisecond
+			if backoff > 5*time.Second {
+				backoff = 5 * time.Second
+			}
+			time.Sleep(backoff)
+		}
+	}
+
+	complete = true
+	wa.Wait()
+	torrFile.Close()
+	if errors.Is(err, ERR_STOPPED) {
+		err = nil
+	}
+	return err
+}

server/tgbot/upload/queue.go

@@ -0,0 +1,130 @@
+package upload
+
+import (
+	"fmt"
+	"strconv"
+	"time"
+
+	"github.com/dustin/go-humanize"
+	tele "gopkg.in/telebot.v4"
+	"server/torr"
+)
+
+type DLQueue struct {
+	id        int
+	c         tele.Context
+	hash      string
+	fileID    string
+	fileName  string
+	updateMsg *tele.Message
+}
+
+var manager = &Manager{}
+
+func Start() {
+	manager.Start()
+}
+
+func ShowQueue(c tele.Context) error {
+	msg := ""
+	manager.queueLock.Lock()
+	defer manager.queueLock.Unlock()
+	if len(manager.queue) == 0 && len(manager.working) == 0 {
+		return c.Send(tr(c.Sender().ID, "queue_empty"))
+	}
+	if len(manager.working) > 0 {
+		msg += tr(c.Sender().ID, "upload_working") + ":\n"
+		i := 0
+		for _, dlQueue := range manager.working {
+			s := "#" + strconv.Itoa(i+1) + ": <code>" + dlQueue.torrentHash + "</code>\n"
+			if len(msg+s) > 1024 {
+				c.Send(msg)
+				msg = ""
+			}
+			msg += s
+			i++
+		}
+		if len(msg) > 0 {
+			c.Send(msg)
+			msg = ""
+		}
+	}
+	if len(manager.queue) > 0 {
+		msg = tr(c.Sender().ID, "upload_in_queue") + ":\n"
+		for i, dlQueue := range manager.queue {
+			s := "#" + strconv.Itoa(i+1) + ": <code>" + dlQueue.torrentHash + "</code>\n"
+			if len(msg+s) > 1024 {
+				c.Send(msg)
+				msg = ""
+			}
+			msg += s
+		}
+		if len(msg) > 0 {
+			c.Send(msg)
+			msg = ""
+		}
+	}
+	return nil
+}
+
+func AddRange(c tele.Context, hash string, from, to int) {
+	manager.AddRange(c, hash, from, to)
+}
+
+func Cancel(id int) {
+	manager.Cancel(id)
+}
+
+func updateLoadStatus(wrk *Worker, file *TorrFile, fi, fc int) {
+	if wrk.msg == nil {
+		return
+	}
+	t := torr.GetTorrent(wrk.torrentHash)
+	if t == nil {
+		return
+	}
+	ti := t.Status()
+	if wrk.isCancelled {
+		wrk.c.Bot().Edit(wrk.msg, tr(wrk.c.Sender().ID, "upload_stopping"))
+	} else {
+		wrk.c.Send(tele.UploadingVideo)
+		if ti.DownloadSpeed == 0 {
+			ti.DownloadSpeed = 1.0
+		}
+		wait := time.Duration(float64(file.Remaining())/ti.DownloadSpeed) * time.Second
+		speed := humanize.IBytes(uint64(ti.DownloadSpeed)) + "/sec"
+		peers := fmt.Sprintf("%v · %v/%v", ti.ConnectedSeeders, ti.ActivePeers, ti.TotalPeers)
+		prc := fmt.Sprintf("%.2f%% %v / %v", float64(file.offset)*100.0/float64(file.size), humanize.IBytes(uint64(file.offset)), humanize.IBytes(uint64(file.size)))
+
+		name := file.name
+		if name == ti.Title {
+			name = ""
+		}
+
+		uid := wrk.c.Sender().ID
+		msg := tr(uid, "upload_title") + ":\n" +
+			"<b>" + escapeHtml(ti.Title) + "</b>\n"
+		if name != "" {
+			msg += "<i>" + escapeHtml(name) + "</i>\n"
+		}
+		msg += "<b>" + tr(uid, "upload_hash") + ":</b> <code>" + file.hash + "</code>\n"
+		if file.offset < file.size {
+			msg += "<b>" + tr(uid, "upload_speed") + ": </b>" + speed + "\n" +
+				"<b>" + tr(uid, "upload_remaining") + ": </b>" + wait.String() + "\n" +
+				"<b>" + tr(uid, "upload_peers") + ": </b>" + peers + "\n" +
+				"<b>" + tr(uid, "upload_progress") + ": </b>" + prc
+		}
+		if fc > 1 {
+			msg += "\n<b>" + tr(uid, "upload_files") + ": </b>" + strconv.Itoa(fi) + "/" + strconv.Itoa(fc)
+		}
+		if file.offset >= file.size {
+			msg += "\n<b>" + tr(uid, "upload_finishing") + "</b>"
+			wrk.c.Bot().Edit(wrk.msg, msg)
+			return
+		}
+
+		torrKbd := &tele.ReplyMarkup{}
+		torrKbd.Inline([]tele.Row{torrKbd.Row(torrKbd.Data(tr(wrk.c.Sender().ID, "upload_cancel"), "cancel", strconv.Itoa(wrk.id)))}...)
+		wrk.c.Bot().Edit(wrk.msg, msg, torrKbd)
+	}
+}

server/tgbot/upload/torrfile.go

@@ -0,0 +1,98 @@
+package upload
+
+import (
+	"errors"
+	"fmt"
+	"path/filepath"
+
+	"github.com/anacrolix/torrent"
+
+	sets "server/settings"
+	"server/log"
+	"server/tgbot/config"
+	"server/torr"
+	"server/torr/state"
+	"server/torr/storage/torrstor"
+)
+
+var ERR_STOPPED = errors.New("stopped")
+
+type TorrFile struct {
+	hash   string
+	name   string
+	wrk    *Worker
+	offset int64
+	size   int64
+	id     int
+
+	reader *torrstor.Reader
+}
+
+func NewTorrFile(wrk *Worker, stFile *state.TorrentFileStat) (*TorrFile, error) {
+	uid := int64(0)
+	if wrk.c != nil && wrk.c.Sender() != nil {
+		uid = wrk.c.Sender().ID
+	}
+	if config.Cfg != nil && config.Cfg.HostTG != "" && stFile.Length > 2*1024*1024*1024 {
+		return nil, errors.New(tr(uid, "upload_file_too_large_2gb"))
+	}
+	if (config.Cfg == nil || config.Cfg.HostTG == "") && stFile.Length > 50*1024*1024 {
+		return nil, errors.New(tr(uid, "upload_file_too_large_50mb"))
+	}
+
+	tf := new(TorrFile)
+	tf.hash = wrk.torrentHash
+	tf.name = filepath.Base(stFile.Path)
+	tf.wrk = wrk
+	tf.size = stFile.Length
+
+	t := torr.GetTorrent(wrk.torrentHash)
+	t.WaitInfo()
+
+	files := t.Files()
+	var file *torrent.File
+	for _, tfile := range files {
+		if tfile.Path() == stFile.Path {
+			file = tfile
+			break
+		}
+	}
+	if file == nil {
+		return nil, fmt.Errorf("file with id %v not found", stFile.Id)
+	}
+	if int64(sets.MaxSize) > 0 && file.Length() > int64(sets.MaxSize) {
+		log.TLogln("tg upload err size", file.DisplayPath(), "max", sets.MaxSize)
+		return nil, fmt.Errorf("file size exceeded max allowed %d bytes", sets.MaxSize)
+	}
+
+	reader := t.NewReader(file)
+	if reader == nil {
+		return nil, errors.New("cannot create torrent reader")
+	}
+	if sets.BTsets != nil && sets.BTsets.ResponsiveMode {
+		reader.SetResponsive()
+	}
+	tf.reader = reader
+
+	return tf, nil
+}
+
+func (t *TorrFile) Read(p []byte) (n int, err error) {
+	if t.wrk.isCancelled {
+		return 0, ERR_STOPPED
+	}
+	n, err = t.reader.Read(p)
+	t.offset += int64(n)
+	return
+}
+
+func (t *TorrFile) Remaining() int64 {
+	return t.size - t.offset
+}
+
+func (t *TorrFile) Close() {
+	if t.reader != nil {
+		t.reader.Close()
+		t.reader = nil
+	}
+}