diff --git a/core/apiresp.go b/core/apiresp.go index 59d21267..7e1921a5 100644 --- a/core/apiresp.go +++ b/core/apiresp.go @@ -15,6 +15,7 @@ package larkcore import ( "encoding/json" "fmt" + "io" "mime" "net/http" "strconv" @@ -22,9 +23,10 @@ import ( ) type ApiResp struct { - StatusCode int `json:"-"` - Header http.Header `json:"-"` - RawBody []byte `json:"-"` + StatusCode int `json:"-"` + Header http.Header `json:"-"` + RawBody []byte `json:"-"` + Body io.ReadCloser `json:"-"` } func (resp ApiResp) Write(writer http.ResponseWriter) { @@ -34,6 +36,13 @@ func (resp ApiResp) Write(writer http.ResponseWriter) { writer.Header().Add(k, v) } } + if resp.Body != nil { + defer resp.Body.Close() + if _, err := io.Copy(writer, resp.Body); err != nil { + panic(err) + } + return + } if _, err := writer.Write(resp.RawBody); err != nil { panic(err) } @@ -61,7 +70,9 @@ func (resp ApiResp) LogId() string { func (resp ApiResp) String() string { contentType := resp.Header.Get(contentTypeHeader) body := fmt.Sprintf(" len %d", len(resp.RawBody)) - if strings.Contains(contentType, "json") || strings.Contains(contentType, "text") { + if resp.Body != nil { + body = "" + } else if strings.Contains(contentType, "json") || strings.Contains(contentType, "text") { body = string(resp.RawBody) } return fmt.Sprintf("StatusCode: %d, Header:%v, Content-Type: %s, Body: %v", resp.StatusCode, diff --git a/core/httptransport.go b/core/httptransport.go index 615bbbc8..c4ef9da0 100644 --- a/core/httptransport.go +++ b/core/httptransport.go @@ -229,6 +229,44 @@ func doSend(ctx context.Context, rawRequest *http.Request, httpClient HttpClient }, nil } +func doSendStream(ctx context.Context, rawRequest *http.Request, httpClient HttpClient, logger Logger) (*ApiResp, error) { + if httpClient == nil { + httpClient = http.DefaultClient + } + resp, err := httpClient.Do(rawRequest) + if err != nil { + if er, ok := err.(*url.Error); ok { + if er.Timeout() { + return nil, &ClientTimeoutError{msg: er.Error()} + } + + if e, ok := er.Err.(*net.OpError); ok && e.Op == "dial" { + return nil, &DialFailedError{msg: er.Error()} + } + } + return nil, err + } + + if resp.StatusCode == http.StatusGatewayTimeout { + logID := resp.Header.Get(HttpHeaderKeyLogId) + if logID == "" { + logID = resp.Header.Get(HttpHeaderKeyRequestId) + } + logger.Info(ctx, fmt.Sprintf("req path:%s, server time out,requestId:%s", + rawRequest.URL.RequestURI(), logID)) + if resp.Body != nil { + _ = resp.Body.Close() + } + return nil, &ServerTimeoutError{msg: "server time out error"} + } + + return &ApiResp{ + StatusCode: resp.StatusCode, + Header: resp.Header, + Body: resp.Body, + }, nil +} + func Request(ctx context.Context, req *ApiReq, config *Config, options ...RequestOptionFunc) (*ApiResp, error) { option := &RequestOption{} for _, optionFunc := range options { @@ -257,7 +295,43 @@ func Request(ctx context.Context, req *ApiReq, config *Config, options ...Reques } +func RequestStream(ctx context.Context, req *ApiReq, config *Config, options ...RequestOptionFunc) (*ApiResp, error) { + option := &RequestOption{} + for _, optionFunc := range options { + optionFunc(option) + } + + if len(req.SupportedAccessTokenTypes) == 0 { + req.SupportedAccessTokenTypes = append(req.SupportedAccessTokenTypes, AccessTokenTypeNone) + } + + err := validateTokenType(req.SupportedAccessTokenTypes, option) + if err != nil { + return nil, err + } + accessTokenType, err := determineTokenType(req.SupportedAccessTokenTypes, option, config) + if err != nil { + return nil, err + } + err = validate(config, option, accessTokenType) + if err != nil { + return nil, err + } + + return doRequestStream(ctx, req, accessTokenType, config, option) +} + +type sendApiRespFunc func(context.Context, *http.Request, HttpClient, Logger) (*ApiResp, error) + func doRequest(ctx context.Context, httpReq *ApiReq, accessTokenType AccessTokenType, config *Config, option *RequestOption) (*ApiResp, error) { + return doRequestWithSender(ctx, httpReq, accessTokenType, config, option, doSend, false) +} + +func doRequestStream(ctx context.Context, httpReq *ApiReq, accessTokenType AccessTokenType, config *Config, option *RequestOption) (*ApiResp, error) { + return doRequestWithSender(ctx, httpReq, accessTokenType, config, option, doSendStream, true) +} + +func doRequestWithSender(ctx context.Context, httpReq *ApiReq, accessTokenType AccessTokenType, config *Config, option *RequestOption, send sendApiRespFunc, streamBody bool) (*ApiResp, error) { var rawResp *ApiResp var errResult error for i := 0; i < 2; i++ { @@ -276,7 +350,7 @@ func doRequest(ctx context.Context, httpReq *ApiReq, accessTokenType AccessToken } else { config.Logger.Debug(ctx, fmt.Sprintf("req:%s,%s", httpReq.HttpMethod, httpReq.ApiPath)) } - rawResp, err = doSend(ctx, req, config.HttpClient, config.Logger) + rawResp, err = send(ctx, req, config.HttpClient, config.Logger) if config.LogReqAtDebug { if shouldSkipCodeErrorPreDecode(httpReq.ApiPath) { statusCode := 0 @@ -297,8 +371,16 @@ func doRequest(ctx context.Context, httpReq *ApiReq, accessTokenType AccessToken continue } - fileDownloadSuccess := option.FileDownload && rawResp.StatusCode == http.StatusOK - if fileDownloadSuccess || !strings.Contains(rawResp.Header.Get(contentTypeHeader), contentTypeJson) { + fileDownloadSuccess := option.FileDownload && isFileDownloadStatus(rawResp.StatusCode) + if streamBody && !fileDownloadSuccess { + if err := rawResp.readAndCloseBody(); err != nil { + return nil, err + } + } + if fileDownloadSuccess { + break + } + if !strings.Contains(rawResp.Header.Get(contentTypeHeader), contentTypeJson) { break } if shouldSkipCodeErrorPreDecode(httpReq.ApiPath) { @@ -335,3 +417,20 @@ func doRequest(ctx context.Context, httpReq *ApiReq, accessTokenType AccessToken } return rawResp, nil } + +func (resp *ApiResp) readAndCloseBody() error { + if resp == nil || resp.Body == nil { + return nil + } + body, err := readResponseBody(resp.Body) + if err != nil { + return err + } + resp.RawBody = body + resp.Body = nil + return nil +} + +func isFileDownloadStatus(statusCode int) bool { + return statusCode == http.StatusOK || statusCode == http.StatusPartialContent +} diff --git a/core/httptransport_test.go b/core/httptransport_test.go index 6d1306eb..ac94f927 100644 --- a/core/httptransport_test.go +++ b/core/httptransport_test.go @@ -55,6 +55,22 @@ func (b *closeTrackingBody) Close() error { return nil } +type readTrackingBody struct { + reader *strings.Reader + reads int32 + closed int32 +} + +func (b *readTrackingBody) Read(p []byte) (int, error) { + atomic.AddInt32(&b.reads, 1) + return b.reader.Read(p) +} + +func (b *readTrackingBody) Close() error { + atomic.StoreInt32(&b.closed, 1) + return nil +} + type httpClientStub struct { resp *http.Response err error @@ -123,6 +139,88 @@ func TestDoSend_CloseBodyOnGatewayTimeout(t *testing.T) { } } +func TestRequestStreamDoesNotReadSuccessfulDownloadBody(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, "http://example.com/", nil) + if err != nil { + t.Fatalf("new request failed: %v", err) + } + body := &readTrackingBody{reader: strings.NewReader("stream-body")} + client := httpClientStub{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{contentTypeHeader: []string{"application/octet-stream"}}, + Body: body, + Request: req, + }} + config := mockConfig() + config.HttpClient = client + + resp, err := RequestStream(context.Background(), &ApiReq{ + HttpMethod: http.MethodGet, + ApiPath: "/download", + SupportedAccessTokenTypes: []AccessTokenType{AccessTokenTypeUser}, + }, config, WithUserAccessToken("user-token"), WithFileDownload()) + if err != nil { + t.Fatalf("request stream failed: %v", err) + } + if atomic.LoadInt32(&body.reads) != 0 { + t.Fatalf("expected stream body unread before caller consumes it") + } + + data, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read stream failed: %v", err) + } + if string(data) != "stream-body" { + t.Fatalf("unexpected stream body: %s", data) + } + if atomic.LoadInt32(&body.reads) == 0 { + t.Fatalf("expected caller read to consume stream body") + } + if err := resp.Body.Close(); err != nil { + t.Fatalf("close stream failed: %v", err) + } + if atomic.LoadInt32(&body.closed) != 1 { + t.Fatalf("expected stream body closed") + } +} + +func TestRequestStreamReadsAndClosesErrorBody(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, "http://example.com/", nil) + if err != nil { + t.Fatalf("new request failed: %v", err) + } + body := &readTrackingBody{reader: strings.NewReader(`{"code":999,"msg":"failed"}`)} + client := httpClientStub{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{contentTypeHeader: []string{contentTypeJson}}, + Body: body, + Request: req, + }} + config := mockConfig() + config.HttpClient = client + + resp, err := RequestStream(context.Background(), &ApiReq{ + HttpMethod: http.MethodGet, + ApiPath: "/download", + SupportedAccessTokenTypes: []AccessTokenType{AccessTokenTypeUser}, + }, config, WithUserAccessToken("user-token"), WithFileDownload()) + if err != nil { + t.Fatalf("request stream failed: %v", err) + } + if resp.Body != nil { + t.Fatalf("expected error body to be closed and cleared") + } + if got := string(resp.RawBody); got != `{"code":999,"msg":"failed"}` { + t.Fatalf("unexpected raw body: %s", got) + } + if atomic.LoadInt32(&body.reads) == 0 { + t.Fatalf("expected error body to be read") + } + if atomic.LoadInt32(&body.closed) != 1 { + t.Fatalf("expected error body closed") + } +} + func TestDetermineTokenTypeRejectsAppOnlyInClientAssertionMode(t *testing.T) { config := mockConfig() config.ClientAssertionProvider = &mockClientAssertionProvider{token: &Token{Value: "assertion"}} diff --git a/core/utils.go b/core/utils.go index fcb693f6..72277760 100644 --- a/core/utils.go +++ b/core/utils.go @@ -505,8 +505,12 @@ func UserAgent(source string) string { } func readResponse(resp *http.Response) ([]byte, error) { - defer resp.Body.Close() - respBody, err := ioutil.ReadAll(resp.Body) + return readResponseBody(resp.Body) +} + +func readResponseBody(body io.ReadCloser) ([]byte, error) { + defer body.Close() + respBody, err := ioutil.ReadAll(body) if err != nil { return nil, err } diff --git a/service/drive/v1/resource.go b/service/drive/v1/resource.go index aa3d38f5..24894914 100644 --- a/service/drive/v1/resource.go +++ b/service/drive/v1/resource.go @@ -49,6 +49,10 @@ func New(config *larkcore.Config) *V1 { } } +func isDownloadSuccessStatus(statusCode int) bool { + return statusCode == http.StatusOK || statusCode == http.StatusPartialContent +} + type exportTask struct { config *larkcore.Config } @@ -144,7 +148,7 @@ func (e *exportTask) Download(ctx context.Context, req *DownloadExportTaskReq, o // 反序列响应结果 resp := &DownloadExportTaskResp{ApiResp: apiResp} // 如果是下载，则设置响应结果 - if apiResp.StatusCode == http.StatusOK { + if isDownloadSuccessStatus(apiResp.StatusCode) { resp.File = bytes.NewBuffer(apiResp.RawBody) resp.FileName = larkcore.FileNameByHeader(apiResp.Header) return resp, err @@ -156,6 +160,30 @@ func (e *exportTask) Download(ctx context.Context, req *DownloadExportTaskReq, o return resp, err } +// DownloadStream 下载导出文件，返回响应 body 流，调用方需要关闭 resp.ApiResp.Body。 +func (e *exportTask) DownloadStream(ctx context.Context, req *DownloadExportTaskReq, options ...larkcore.RequestOptionFunc) (*DownloadExportTaskResp, error) { + apiReq := req.apiReq + apiReq.ApiPath = "/open-apis/drive/v1/export_tasks/file/:file_token/download" + apiReq.HttpMethod = http.MethodGet + apiReq.SupportedAccessTokenTypes = []larkcore.AccessTokenType{larkcore.AccessTokenTypeTenant, larkcore.AccessTokenTypeUser} + options = append(options, larkcore.WithFileDownload()) + apiResp, err := larkcore.RequestStream(ctx, apiReq, e.config, options...) + if err != nil { + return nil, err + } + resp := &DownloadExportTaskResp{ApiResp: apiResp} + if isDownloadSuccessStatus(apiResp.StatusCode) { + resp.File = apiResp.Body + resp.FileName = larkcore.FileNameByHeader(apiResp.Header) + return resp, err + } + err = apiResp.JSONUnmarshalBody(resp, e.config) + if err != nil { + return nil, err + } + return resp, err +} + // Get 查询导出任务结果 // // - 根据[创建导出任务](/ssl::ttdoc//uAjLw4CM/ukTMukTMukTM/reference/drive-v1/export_task/create)的ticket查询导出任务的结果，前提条件需要先调用创建导出任务接口。;;通过该接口获取到下载文件的 token 后调用[下载导出文件接口](/ssl::ttdoc//uAjLw4CM/ukTMukTMukTM/reference/drive-v1/export_task/download)将文件进行下载 @@ -342,7 +370,7 @@ func (f *file) Download(ctx context.Context, req *DownloadFileReq, options ...la // 反序列响应结果 resp := &DownloadFileResp{ApiResp: apiResp} // 如果是下载，则设置响应结果 - if apiResp.StatusCode == http.StatusOK { + if isDownloadSuccessStatus(apiResp.StatusCode) { resp.File = bytes.NewBuffer(apiResp.RawBody) resp.FileName = larkcore.FileNameByHeader(apiResp.Header) return resp, err @@ -354,6 +382,30 @@ func (f *file) Download(ctx context.Context, req *DownloadFileReq, options ...la return resp, err } +// DownloadStream 下载文件，返回响应 body 流，调用方需要关闭 resp.ApiResp.Body。 +func (f *file) DownloadStream(ctx context.Context, req *DownloadFileReq, options ...larkcore.RequestOptionFunc) (*DownloadFileResp, error) { + apiReq := req.apiReq + apiReq.ApiPath = "/open-apis/drive/v1/files/:file_token/download" + apiReq.HttpMethod = http.MethodGet + apiReq.SupportedAccessTokenTypes = []larkcore.AccessTokenType{larkcore.AccessTokenTypeUser, larkcore.AccessTokenTypeTenant} + options = append(options, larkcore.WithFileDownload()) + apiResp, err := larkcore.RequestStream(ctx, apiReq, f.config, options...) + if err != nil { + return nil, err + } + resp := &DownloadFileResp{ApiResp: apiResp} + if isDownloadSuccessStatus(apiResp.StatusCode) { + resp.File = apiResp.Body + resp.FileName = larkcore.FileNameByHeader(apiResp.Header) + return resp, err + } + err = apiResp.JSONUnmarshalBody(resp, f.config) + if err != nil { + return nil, err + } + return resp, err +} + // GetSubscribe 查询云文档事件订阅状态 // // - 该接口**仅支持文档拥有者**查询自己文档的订阅状态，可订阅的文档类型为**旧版文档**、**新版文档**、**电子表格**和**多维表格**。在调用该接口之前请确保正确[配置事件回调网址和订阅事件类型](https://open.feishu.cn/document/ukTMukTMukTM/uUTNz4SN1MjL1UzM#2eb3504a)，事件类型参考[事件列表](https://open.feishu.cn/document/ukTMukTMukTM/uYDNxYjL2QTM24iN0EjN/event-list)。 @@ -1204,7 +1256,7 @@ func (m *media) Download(ctx context.Context, req *DownloadMediaReq, options ... // 反序列响应结果 resp := &DownloadMediaResp{ApiResp: apiResp} // 如果是下载，则设置响应结果 - if apiResp.StatusCode == http.StatusOK { + if isDownloadSuccessStatus(apiResp.StatusCode) { resp.File = bytes.NewBuffer(apiResp.RawBody) resp.FileName = larkcore.FileNameByHeader(apiResp.Header) return resp, err @@ -1216,6 +1268,30 @@ func (m *media) Download(ctx context.Context, req *DownloadMediaReq, options ... return resp, err } +// DownloadStream 下载素材，返回响应 body 流，调用方需要关闭 resp.ApiResp.Body。 +func (m *media) DownloadStream(ctx context.Context, req *DownloadMediaReq, options ...larkcore.RequestOptionFunc) (*DownloadMediaResp, error) { + apiReq := req.apiReq + apiReq.ApiPath = "/open-apis/drive/v1/medias/:file_token/download" + apiReq.HttpMethod = http.MethodGet + apiReq.SupportedAccessTokenTypes = []larkcore.AccessTokenType{larkcore.AccessTokenTypeUser, larkcore.AccessTokenTypeTenant} + options = append(options, larkcore.WithFileDownload()) + apiResp, err := larkcore.RequestStream(ctx, apiReq, m.config, options...) + if err != nil { + return nil, err + } + resp := &DownloadMediaResp{ApiResp: apiResp} + if isDownloadSuccessStatus(apiResp.StatusCode) { + resp.File = apiResp.Body + resp.FileName = larkcore.FileNameByHeader(apiResp.Header) + return resp, err + } + err = apiResp.JSONUnmarshalBody(resp, m.config) + if err != nil { + return nil, err + } + return resp, err +} + // UploadAll 上传素材 // // - 将文件、图片、视频等素材文件上传到指定云文档中。素材文件在云空间中不会显示，只会显示在对应云文档中。 diff --git a/service/drive/v1/resource_test.go b/service/drive/v1/resource_test.go new file mode 100644 index 00000000..145f0f94 --- /dev/null +++ b/service/drive/v1/resource_test.go @@ -0,0 +1,167 @@ +package larkdrive + +import ( + "context" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + + larkcore "github.com/larksuite/oapi-sdk-go/v3/core" +) + +type streamDownloadBody struct { + reader *strings.Reader + reads int32 + closed int32 +} + +func (b *streamDownloadBody) Read(p []byte) (int, error) { + atomic.AddInt32(&b.reads, 1) + return b.reader.Read(p) +} + +func (b *streamDownloadBody) Close() error { + atomic.StoreInt32(&b.closed, 1) + return nil +} + +type streamDownloadClient struct { + t *testing.T + body *streamDownloadBody + expectedPath string + statusCode int +} + +func (c *streamDownloadClient) Do(req *http.Request) (*http.Response, error) { + c.t.Helper() + if req.URL.Path != c.expectedPath { + c.t.Fatalf("unexpected request path: %s", req.URL.Path) + } + if got := req.Header.Get("Authorization"); got != "Bearer user-token" { + c.t.Fatalf("unexpected authorization header: %s", got) + } + return &http.Response{ + StatusCode: c.statusCode, + Header: http.Header{ + "Content-Type": []string{"application/octet-stream"}, + "Content-Disposition": []string{`attachment; filename="large.bin"`}, + }, + Body: c.body, + Request: req, + }, nil +} + +func TestDownloadStreamDoesNotBufferResponse(t *testing.T) { + tests := []struct { + name string + expectedPath string + statusCode int + download func(context.Context, *V1) (io.Reader, *larkcore.ApiResp, string, error) + }{ + { + name: "export task", + expectedPath: "/open-apis/drive/v1/export_tasks/file/file-token/download", + statusCode: http.StatusOK, + download: func(ctx context.Context, client *V1) (io.Reader, *larkcore.ApiResp, string, error) { + resp, err := client.ExportTask.DownloadStream( + ctx, + NewDownloadExportTaskReqBuilder().FileToken("file-token").Build(), + larkcore.WithUserAccessToken("user-token"), + ) + if err != nil { + return nil, nil, "", err + } + return resp.File, resp.ApiResp, resp.FileName, nil + }, + }, + { + name: "file", + expectedPath: "/open-apis/drive/v1/files/file-token/download", + statusCode: http.StatusOK, + download: func(ctx context.Context, client *V1) (io.Reader, *larkcore.ApiResp, string, error) { + resp, err := client.File.DownloadStream( + ctx, + NewDownloadFileReqBuilder().FileToken("file-token").Build(), + larkcore.WithUserAccessToken("user-token"), + ) + if err != nil { + return nil, nil, "", err + } + return resp.File, resp.ApiResp, resp.FileName, nil + }, + }, + { + name: "media", + expectedPath: "/open-apis/drive/v1/medias/file-token/download", + statusCode: http.StatusOK, + download: func(ctx context.Context, client *V1) (io.Reader, *larkcore.ApiResp, string, error) { + resp, err := client.Media.DownloadStream( + ctx, + NewDownloadMediaReqBuilder().FileToken("file-token").Build(), + larkcore.WithUserAccessToken("user-token"), + ) + if err != nil { + return nil, nil, "", err + } + return resp.File, resp.ApiResp, resp.FileName, nil + }, + }, + { + name: "media range", + expectedPath: "/open-apis/drive/v1/medias/file-token/download", + statusCode: http.StatusPartialContent, + download: func(ctx context.Context, client *V1) (io.Reader, *larkcore.ApiResp, string, error) { + resp, err := client.Media.DownloadStream( + ctx, + NewDownloadMediaReqBuilder().FileToken("file-token").Build(), + larkcore.WithUserAccessToken("user-token"), + ) + if err != nil { + return nil, nil, "", err + } + return resp.File, resp.ApiResp, resp.FileName, nil + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := &streamDownloadBody{reader: strings.NewReader("large-file")} + config := &larkcore.Config{ + BaseUrl: "https://open.feishu.cn", + AppId: "app-id", + AppSecret: "app-secret", + HttpClient: &streamDownloadClient{t: t, body: body, expectedPath: tt.expectedPath, statusCode: tt.statusCode}, + Logger: larkcore.NewDefaultLogger(larkcore.LogLevelError), + Serializable: &larkcore.DefaultSerialization{}, + } + + file, apiResp, fileName, err := tt.download(context.Background(), New(config)) + if err != nil { + t.Fatalf("download stream failed: %v", err) + } + if fileName != "large.bin" { + t.Fatalf("unexpected file name: %s", fileName) + } + if atomic.LoadInt32(&body.reads) != 0 { + t.Fatalf("expected download body unread before caller consumes it") + } + + data, err := io.ReadAll(file) + if err != nil { + t.Fatalf("read streamed file failed: %v", err) + } + if string(data) != "large-file" { + t.Fatalf("unexpected streamed file body: %s", data) + } + if err := apiResp.Body.Close(); err != nil { + t.Fatalf("close streamed response failed: %v", err) + } + if atomic.LoadInt32(&body.closed) != 1 { + t.Fatalf("expected streamed response body closed") + } + }) + } +}