Skip to content

Files upload refactored #70

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

Files upload refactored #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
252 changes: 130 additions & 122 deletions bot.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,23 @@ type BotAPI struct {
mu sync.RWMutex
}

type requestPayload struct {
body io.Reader
closer io.Closer
contentType string
}

func (p requestPayload) close() {
if p.closer != nil {
_ = p.closer.Close()
}
}

type requestDebug struct {
params Params
fileCount int
}

// NewBotAPI creates a new BotAPI instance.
//
// It requires a token, provided by @BotFather on Telegram.
Expand Down Expand Up @@ -93,25 +110,106 @@ func buildParams(in Params) url.Values {
return out
}

// MakeRequest makes a request to a specific endpoint with our token.
func (bot *BotAPI) MakeRequest(endpoint string, params Params) (*APIResponse, error) {
return bot.MakeRequestWithContext(context.Background(), endpoint, params)
func buildFormPayload(params Params) requestPayload {
values := buildParams(params)
reader := strings.NewReader(values.Encode())
return requestPayload{
body: reader,
contentType: "application/x-www-form-urlencoded",
}
}

func (bot *BotAPI) MakeRequestWithContext(ctx context.Context, endpoint string, params Params) (*APIResponse, error) {
func buildMultipartPayload(params Params, files []RequestFile) (requestPayload, error) {
reader, writer := io.Pipe()
multipartWriter := multipart.NewWriter(writer)

go func() {
defer writer.Close()
defer multipartWriter.Close()

for field, value := range params {
if err := multipartWriter.WriteField(field, value); err != nil {
writer.CloseWithError(err)
return
}
}

for _, file := range files {
source, err := resolveRequestFileData(file.Data)
if err != nil {
writer.CloseWithError(err)
return
}

if source.kindIsUpload() {
desc, err := source.openUpload()
if err != nil {
writer.CloseWithError(err)
return
}

part, err := multipartWriter.CreateFormFile(file.Name, desc.name)
if err != nil {
desc.reader.Close()
writer.CloseWithError(err)
return
}

if _, err = io.Copy(part, desc.reader); err != nil {
desc.reader.Close()
writer.CloseWithError(err)
return
}

if err = desc.reader.Close(); err != nil {
writer.CloseWithError(err)
return
}

continue
}

value, err := source.referenceValue()
if err != nil {
writer.CloseWithError(err)
return
}

if err = multipartWriter.WriteField(file.Name, value); err != nil {
writer.CloseWithError(err)
return
}
}
}()

return requestPayload{
body: reader,
closer: reader,
contentType: multipartWriter.FormDataContentType(),
}, nil
}

func (bot *BotAPI) executeRequest(ctx context.Context, endpoint string, payload requestPayload, debugInfo requestDebug) (*APIResponse, error) {
defer payload.close()

if bot.Debug {
log.Printf("Endpoint: %s, params: %v\n", endpoint, params)
if debugInfo.fileCount > 0 {
log.Printf("Endpoint: %s, params: %v, with %d files\n", endpoint, debugInfo.params, debugInfo.fileCount)
} else {
log.Printf("Endpoint: %s, params: %v\n", endpoint, debugInfo.params)
}
}

method := fmt.Sprintf(bot.apiEndpoint, bot.Token, endpoint)

values := buildParams(params)

req, err := http.NewRequestWithContext(ctx, "POST", method, strings.NewReader(values.Encode()))
req, err := http.NewRequestWithContext(ctx, "POST", method, payload.body)
if err != nil {
return &APIResponse{}, err
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

if payload.contentType != "" {
req.Header.Set("Content-Type", payload.contentType)
}

resp, err := bot.Client.Do(req)
if err != nil {
Expand Down Expand Up @@ -146,6 +244,16 @@ func (bot *BotAPI) MakeRequestWithContext(ctx context.Context, endpoint string,
return &apiResp, nil
}

// MakeRequest makes a request to a specific endpoint with our token.
func (bot *BotAPI) MakeRequest(endpoint string, params Params) (*APIResponse, error) {
return bot.MakeRequestWithContext(context.Background(), endpoint, params)
}

func (bot *BotAPI) MakeRequestWithContext(ctx context.Context, endpoint string, params Params) (*APIResponse, error) {
payload := buildFormPayload(params)
return bot.executeRequest(ctx, endpoint, payload, requestDebug{params: params})
}

// decodeAPIResponse decode response and return slice of bytes if debug enabled.
// If debug disabled, just decode http.Response.Body stream to APIResponse struct
// for efficient memory usage
Expand Down Expand Up @@ -176,101 +284,15 @@ func (bot *BotAPI) UploadFiles(endpoint string, params Params, files []RequestFi
}

func (bot *BotAPI) UploadFilesWithContext(ctx context.Context, endpoint string, params Params, files []RequestFile) (*APIResponse, error) {
r, w := io.Pipe()
m := multipart.NewWriter(w)

// This code modified from the very helpful @HirbodBehnam
// https://github.com/go-telegram-bot-api/telegram-bot-api/issues/354#issuecomment-663856473
go func() {
defer w.Close()
defer m.Close()

for field, value := range params {
if err := m.WriteField(field, value); err != nil {
w.CloseWithError(err)
return
}
}

for _, file := range files {
if file.Data.NeedsUpload() {
name, reader, err := file.Data.UploadData()
if err != nil {
w.CloseWithError(err)
return
}

part, err := m.CreateFormFile(file.Name, name)
if err != nil {
w.CloseWithError(err)
return
}

if _, err := io.Copy(part, reader); err != nil {
w.CloseWithError(err)
return
}

if closer, ok := reader.(io.ReadCloser); ok {
if err = closer.Close(); err != nil {
w.CloseWithError(err)
return
}
}
} else {
value := file.Data.SendData()

if err := m.WriteField(file.Name, value); err != nil {
w.CloseWithError(err)
return
}
}
}
}()

if bot.Debug {
log.Printf("Endpoint: %s, params: %v, with %d files\n", endpoint, params, len(files))
}

method := fmt.Sprintf(bot.apiEndpoint, bot.Token, endpoint)

req, err := http.NewRequestWithContext(ctx, "POST", method, r)
payload, err := buildMultipartPayload(params, files)
if err != nil {
return nil, err
}

req.Header.Set("Content-Type", m.FormDataContentType())

resp, err := bot.Client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()

var apiResp APIResponse
bytes, err := bot.decodeAPIResponse(resp.Body, &apiResp)
if err != nil {
return &apiResp, err
}

if bot.Debug {
log.Printf("Endpoint: %s, response: %s\n", endpoint, string(bytes))
}

if !apiResp.Ok {
var parameters ResponseParameters

if apiResp.Parameters != nil {
parameters = *apiResp.Parameters
}

return &apiResp, &Error{
Message: apiResp.Description,
ResponseParameters: parameters,
}
}

return &apiResp, nil
return bot.executeRequest(ctx, endpoint, payload, requestDebug{
params: params,
fileCount: len(files),
})
}

// GetFileDirectURL returns direct URL to file
Expand Down Expand Up @@ -313,16 +335,6 @@ func (bot *BotAPI) IsMessageToMe(message Message) bool {
return strings.Contains(message.Text, "@"+bot.Self.UserName)
}

func hasFilesNeedingUpload(files []RequestFile) bool {
for _, file := range files {
if file.Data.NeedsUpload() {
return true
}
}

return false
}

// Request sends a Chattable to Telegram, and returns the APIResponse.
func (bot *BotAPI) Request(c Chattable) (*APIResponse, error) {
return bot.RequestWithContext(context.Background(), c)
Expand All @@ -335,18 +347,11 @@ func (bot *BotAPI) RequestWithContext(ctx context.Context, c Chattable) (*APIRes
}

if t, ok := c.(Fileable); ok {
files := t.files()
payload := payloadFromFileable(t)
params = payload.applyInline(params)

// If we have files that need to be uploaded, we should delegate the
// request to UploadFile.
if hasFilesNeedingUpload(files) {
return bot.UploadFiles(t.method(), params, files)
}

// However, if there are no files to be uploaded, there's likely things
// that need to be turned into params instead.
for _, file := range files {
params[file.Name] = file.Data.SendData()
if payload.needsUpload() {
return bot.UploadFilesWithContext(ctx, t.method(), params, payload.filesSlice())
}
}

Expand Down Expand Up @@ -576,9 +581,12 @@ func WriteToHTTPResponse(w http.ResponseWriter, c Chattable) error {
}

if t, ok := c.(Fileable); ok {
if hasFilesNeedingUpload(t.files()) {
payload := payloadFromFileable(t)
if payload.needsUpload() {
return errors.New("unable to use http response to upload files")
}

params = payload.applyInline(params)
}

values := buildParams(params)
Expand Down
Loading