From 8afa3d0ddaecb104abde91543e8a59ac33f140c8 Mon Sep 17 00:00:00 2001 From: csznet Date: Tue, 5 Mar 2024 17:39:25 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3=E4=B8=8B=E8=BD=BD=E5=A4=B1?= =?UTF-8?q?=E8=B4=A5=E3=80=81=E5=A2=9E=E5=8A=A0=E4=B8=8B=E8=BD=BD=E8=BF=9B?= =?UTF-8?q?=E5=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- assets/templates/footer.tmpl | 1 + control/control.go | 42 +++++++++++++++++++++++++----------- main.go | 1 + utils/utils.go | 21 +++++++++++++----- 4 files changed, 47 insertions(+), 18 deletions(-) diff --git a/assets/templates/footer.tmpl b/assets/templates/footer.tmpl index fed768f..ed91446 100644 --- a/assets/templates/footer.tmpl +++ b/assets/templates/footer.tmpl @@ -40,6 +40,7 @@ } var temp = "tgstate-blob"; temp = temp + '\n' + file.name; + temp = temp + '\nsize' + file.size; uploadNextChunk() .then(() => { console.log(temp); // 所有块上传完成后打印 diff --git a/control/control.go b/control/control.go index abdf522..0e34d3d 100644 --- a/control/control.go +++ b/control/control.go @@ -2,13 +2,14 @@ package control import ( "encoding/json" + "fmt" "html/template" "io" "log" "net/http" "path/filepath" + "strconv" "strings" - "time" "csz.net/tgstate/assets" "csz.net/tgstate/conf" @@ -89,24 +90,28 @@ func D(w http.ResponseWriter, r *http.Request) { } // 发起HTTP GET请求来获取Telegram图片 - resp, err := http.Get(utils.GetDownloadUrl(id)) + fileUrl, _ := utils.GetDownloadUrl(id) + resp, err := http.Get(fileUrl) if err != nil { http.Error(w, "Failed to fetch content", http.StatusInternalServerError) return } - defer resp.Body.Close() w.Header().Set("Content-Disposition", "inline") // 设置为 "inline" 以支持在线播放 // 检查Content-Type是否为图片类型 if !strings.HasPrefix(resp.Header.Get("Content-Type"), "application/octet-stream") { - // 设置响应的状态码为 404 w.WriteHeader(http.StatusNotFound) - // 写入响应内容 w.Write([]byte("404 Not Found")) return } - // 读取前512个字节以用于文件类型检测 - buffer := make([]byte, 10*1024*1024) + contentLength, err := strconv.Atoi(resp.Header.Get("Content-Length")) + if err != nil { + log.Println("获取Content-Length出错:", err) + return + } + buffer := make([]byte, contentLength) n, err := resp.Body.Read(buffer) + fmt.Println(string(buffer)) + defer resp.Body.Close() if err != nil && err != io.ErrUnexpectedEOF { log.Println("读取响应主体数据时发生错误:", err) return @@ -116,24 +121,34 @@ func D(w http.ResponseWriter, r *http.Request) { content := string(buffer) lines := strings.Split(content, "\n") log.Println("分块文件:" + lines[1]) + var fileSize string + var startLine = 2 + if strings.HasPrefix(lines[2], "size") { + fileSize = lines[2][len("size"):] + startLine = startLine + 1 + } w.Header().Set("Content-Type", "application/octet-stream") w.Header().Set("Content-Disposition", "attachment; filename=\""+lines[1]+"\"") - for i := 2; i < len(lines); i++ { - blobResp, err := http.Get(utils.GetDownloadUrl(strings.ReplaceAll(lines[i], " ", ""))) + w.Header().Set("Content-Length", fileSize) + for i := startLine; i < len(lines); i++ { + fileStatus := false + var fileUrl string + for !fileStatus { + http.DefaultTransport.(*http.Transport).CloseIdleConnections() + fileUrl, fileStatus = utils.GetDownloadUrl(strings.ReplaceAll(lines[i], " ", "")) + } + blobResp, err := http.Get(fileUrl) if err != nil { http.Error(w, "Failed to fetch content", http.StatusInternalServerError) return } - // 将文件名设置到Content-Disposition标头 - blobResp.Header.Set("Content-Disposition", "attachment; filename=\""+lines[1]+"\"") - defer blobResp.Body.Close() _, err = io.Copy(w, blobResp.Body) + blobResp.Body.Close() if err != nil { log.Println("写入响应主体数据时发生错误:", err) return } } - time.Sleep(10 * time.Second) } else { // 使用DetectContentType函数检测文件类型 w.Header().Set("Content-Type", http.DetectContentType(buffer)) @@ -144,6 +159,7 @@ func D(w http.ResponseWriter, r *http.Request) { return } _, err = io.Copy(w, resp.Body) + resp.Body.Close() if err != nil { log.Println(http.StatusInternalServerError) return diff --git a/main.go b/main.go index 34a3c06..d0c687c 100644 --- a/main.go +++ b/main.go @@ -21,6 +21,7 @@ func main() { fmt.Println("请先设置Bot Token和对象") return } + utils.GetBot() go utils.BotDo() web() } diff --git a/utils/utils.go b/utils/utils.go index cdf6561..92b3a76 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -54,19 +54,30 @@ func UpDocument(fileData tgbotapi.FileReader) string { return resp } -func GetDownloadUrl(fileID string) string { - bot, err := tgbotapi.NewBotAPI(conf.BotToken) +var tgbot *tgbotapi.BotAPI + +func GetBot() { + var err error + tgbot, err = tgbotapi.NewBotAPI(conf.BotToken) if err != nil { + log.Println("初始化bot失败") log.Panic(err) } +} + +func GetDownloadUrl(fileID string) (string, bool) { // 使用 getFile 方法获取文件信息 - file, err := bot.GetFile(tgbotapi.FileConfig{FileID: fileID}) + file, err := tgbot.GetFile(tgbotapi.FileConfig{FileID: fileID}) if err != nil { - log.Panic(err) + log.Println("获取文件失败【" + fileID + "】") + log.Println(err) + return "", false } + log.Println("获取文件成功【" + fileID + "】") // 获取文件下载链接 fileURL := file.Link(conf.BotToken) - return fileURL + + return fileURL, true } func BotDo() { bot, err := tgbotapi.NewBotAPI(conf.BotToken)