fix(drive): add streaming download APIs

core/apiresp.go

@@ -15,16 +15,18 @@ package larkcore
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"mime"
 	"net/http"
 	"strconv"
 	"strings"
 )
 
 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("<binary> len %d", len(resp.RawBody))
-	if strings.Contains(contentType, "json") || strings.Contains(contentType, "text") {
+	if resp.Body != nil {
+		body = "<stream>"
+	} 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,

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
+}

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"}}

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
 	}